from ultralytics import YOLO
from flask import Flask, request, jsonify
from PIL import Image, ImageDraw
import io

app = Flask(__name__)
model = YOLO('best.pt')  # Load your YOLO model
class_names = ['Acne', 'Dark circles', 'blackheads', 'eczema', 'rosacea', 'whiteheads', 'wrinkles']

@app.route('/classify', methods=['POST'])
def classify_image():
    if 'image' not in request.files:
        return jsonify({"error": "No image provided"}), 400

    file = request.files['image']
    if file.filename == '':
        return jsonify({"error": "Empty image file"}), 400

    image = Image.open(io.BytesIO(file.read()))
    resized_image = image.copy()
    resized_image.thumbnail((640, 640))

    # Get results from the model
    
    results = model(resized_image)[0]
    
    predictions = []

    if results.boxes is not None:
        boxes = results.boxes.xyxy
        confidences = results.boxes.conf
        classes = results.boxes.cls

        for i in range(len(boxes)):
            box = boxes[i]
            confidence = confidences[i].item()
            class_id = int(classes[i].item())
            prediction = {
                "x1": box[0].item(),
                "y1": box[1].item(),
                "x2": box[2].item(),
                "y2": box[3].item(),
                "confidence": confidence,
                "class": class_names[class_id],
            }
            predictions.append(prediction)

    return jsonify({"predictions": predictions})

if __name__ == '__main__':
    app.run(host='127.0.0.1', port=5000, debug=True)