figh8back's picture
Update backend/app.py
2c9ad3b verified
import os
from typing import List
import numpy as np
import pandas as pd
from PIL import Image
import tensorflow as tf
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image
from tensorflow.keras.preprocessing.image import load_img
from tensorflow.keras.preprocessing.image import img_to_array
from models.skin_tone.skin_tone_knn import identify_skin_tone
from flask import Flask, request
from flask_restful import Api, Resource, reqparse, abort
import werkzeug
from models.recommender.rec import recs_essentials, makeup_recommendation
import base64
from io import BytesIO
import logging
app = Flask(__name__)
api = Api(app)
# Set up logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)
class_names1 = ['Dry_skin', 'Normal_skin', 'Oil_skin']
class_names2 = ['Low', 'Moderate', 'Severe']
skin_tone_dataset = 'models/skin_tone/skin_tone_dataset.csv'
def get_model():
global model1, model2
model1 = load_model('./models/skin_model')
print('Model 1 loaded')
model2 = load_model('./models/acne_model')
print("Model 2 loaded!")
def load_image(img_path):
img = image.load_img(img_path, target_size=(224, 224))
img_tensor = image.img_to_array(img)
img_tensor = np.expand_dims(img_tensor, axis=0)
img_tensor /= 255.
return img_tensor
def prediction_skin(img_path):
new_image = load_image(img_path)
pred1 = model1.predict(new_image)
if len(pred1[0]) > 1:
pred_class1 = class_names1[tf.argmax(pred1[0])]
else:
pred_class1 = class_names1[int(tf.round(pred1[0]))]
return pred_class1
def prediction_acne(img_path):
new_image = load_image(img_path)
pred2 = model2.predict(new_image)
if len(pred2[0]) > 1:
pred_class2 = class_names2[tf.argmax(pred2[0])]
else:
pred_class2 = class_names2[int(tf.round(pred2[0]))]
return pred_class2
get_model()
# Parsing arguments for image and recommendations
img_put_args = reqparse.RequestParser()
img_put_args.add_argument("file", help="Please provide a valid image file", required=True)
rec_args = reqparse.RequestParser()
rec_args.add_argument("tone", type=int, help="Argument required", required=True)
rec_args.add_argument("type", type=str, help="Argument required", required=True)
rec_args.add_argument("features", type=dict, help="Argument required", required=True)
# Recommendation Class
class Recommendation(Resource):
logger.info(f"Received recommendation request before --------")
def put(self):
args = rec_args.parse_args()
# Log the incoming recommendation request
logger.info(f"Received recommendation request with data: {args}")
features = args['features']
tone = args['tone']
skin_type = args['type'].lower()
skin_tone = 'light to medium'
# Adjust skin tone based on the tone input
if tone <= 2:
skin_tone = 'fair to light'
elif tone >= 4:
skin_tone = 'medium to dark'
# Log the skin tone and type
logger.info(f"Skin tone: {skin_tone}, Skin type: {skin_type}")
fv = []
for key, value in features.items():
fv.append(int(value))
# Log the features sent for recommendation
logger.info(f"Features: {fv}")
try:
general = recs_essentials(fv, None)
makeup = makeup_recommendation(skin_tone, skin_type)
# Log the recommendation data being returned
logger.info(f"Generated recommendations: General: {general}, Makeup: {makeup}")
return {'general': general, 'makeup': makeup}
except Exception as e:
logger.error(f"Error during recommendation generation: {str(e)}")
return {'error': 'Error processing recommendations'}, 400
# Skin Metrics Class
class SkinMetrics(Resource):
def put(self):
args = img_put_args.parse_args()
# Log the incoming image request
logger.info(f"Received image for skin metrics analysis: {args}")
file = args['file']
starter = file.find(',')
image_data = file[starter+1:]
image_data = bytes(image_data, encoding="ascii")
im = Image.open(BytesIO(base64.b64decode(image_data+b'==')))
filename = 'image.png'
file_path = os.path.join('./static', filename)
im.save(file_path)
skin_type = prediction_skin(file_path).split('_')[0]
acne_type = prediction_acne(file_path)
tone = identify_skin_tone(file_path, dataset=skin_tone_dataset)
# Log the predictions for skin type, acne type, and skin tone
logger.info(f"Predicted skin type: {skin_type}, acne type: {acne_type}, tone: {tone}")
return {'type': skin_type, 'tone': str(tone), 'acne': acne_type}, 200
api.add_resource(SkinMetrics, "/upload")
api.add_resource(Recommendation, "/recommend")
# if __name__ == "__main__":
# app.run(debug=False)
if __name__ == "__main__":
app.run(host='0.0.0.0', port=5000)