findsalamanderinfo.py

from flask import request, jsonify
from flask_restful import Resource
from flask_jwt_extended import get_jwt_identity, jwt_required
from algorithm.train_src.new_straightening import straighten
from algorithm.train_src.estimate_gender import estimate_image
import cv2
import os
from api import limiter
from image_encoder.image_encoder import *
from path_constants import image_type, abs_path_imageai_model, abs_path_imageai_config
ALLOWED_EXTENSIONS = ['jpg', 'png', 'jpeg']


# findSalamanderInfo endpoint
class FindSalamanderInfo(Resource):
    decorators = [limiter.limit("5/minute")]

    @staticmethod
    @jwt_required
    def post():
        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/")
                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)), 1)

                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"
                    # gender, _, _, _ = estimate_image(abs_path_imageai_model,abs_path_imageai_config, img)
                    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()