| from ultralytics import YOLO |
| from flask import Flask, request, jsonify |
| from PIL import Image, ImageDraw |
| import io |
|
|
| app = Flask(__name__) |
| model = YOLO('best.pt') |
| 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)) |
|
|
| |
| |
| 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) |
|
|
