Select Git revision
findsalamanderinfo.py
-
Herman Andersen Dyrkorn authoredHerman Andersen Dyrkorn authored
Code owners
Assign users and groups as approvers for specific file changes. Learn more.
findsalamanderinfo.py 2.75 KiB
from flask import request, jsonify
from flask_restful import Resource
from flask_jwt_extended import get_jwt_identity, jwt_required
from algorithm.straighten_with_dlc import straighten
from cv2 import cv2
import os
from api import limiter
from image_encoder.image_encoder import *
from path_constants import image_type
import glob
ALLOWED_EXTENSIONS = ['jpg', 'png', 'jpeg']
class FindSalamanderInfo(Resource):
"""
Endpoint for processing the salamanders abdominal pattern.
POST: Receives an image of a salamander, processes it, and returns the processed image back to the user.
"""
decorators = [limiter.limit("5/minute")]
@staticmethod
@jwt_required
def post():
"""
param image: an image file of a salamander
:return: JSON object containing the processed image or error messages
"""
if "image" not in request.files:
return jsonify({"message": "no file given", "status": 400})
image = request.files['image']
user_id = get_jwt_identity()
if image.filename != '':
if image and allowed_image(image.filename):
temp_img_path = os.path.abspath("./temp_images/")
list_of_temp_images = glob.glob(os.path.join(temp_img_path, str(user_id) + '*.*'))
for temp_image in list_of_temp_images:
os.remove(temp_image)
image.filename = str(user_id) + "." + extension(image.filename)
image.save(os.path.join(temp_img_path, image.filename))
img = cv2.imread(os.path.join("./temp_images", str(user_id) + "." + extension(image.filename)))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# run AI on image
str_image, _, _, _, _ = straighten(img)
if str_image is not None:
str_image = cv2.cvtColor(str_image, cv2.COLOR_BGR2RGB)
cv2.imwrite(os.path.join(temp_img_path, str(user_id) + "_str." + image_type), str_image)
encoded = encode(os.path.join(temp_img_path, str(user_id) + "_str." + image_type))[2:-1]
gender = "AI_sex"
return jsonify({"sex": gender, "species": "AI_species", "status": 200, "image": encoded})
else:
return jsonify({"message": "image could not be processed", "status": 400})
else:
return jsonify({"message": "file extension not allowed", "status": 400})
else:
return jsonify({"message": "something wrong with image", "status": 400})
def allowed_image(filename):
return '.' in filename and \
filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
def extension(filename):
return filename.rsplit('.', 1)[1].lower()