File size: 1,551 Bytes
97523ab |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
from ultralytics import YOLO
from flask import Flask, request, jsonify
from PIL import Image, ImageDraw
import io
app = Flask(__name__)
model = YOLO('last.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)
|