Select Git revision
-
Andreas Blakli authoredAndreas Blakli 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()