|
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) |
|
|