Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,51 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from ultralytics import YOLO
|
2 |
+
from flask import Flask, request, jsonify
|
3 |
+
from PIL import Image, ImageDraw
|
4 |
+
import io
|
5 |
+
|
6 |
+
app = Flask(__name__)
|
7 |
+
model = YOLO('last.pt') # Load your YOLO model
|
8 |
+
class_names = ['Acne', 'Dark circles', 'blackheads', 'eczema', 'rosacea', 'whiteheads', 'wrinkles']
|
9 |
+
|
10 |
+
@app.route('/classify', methods=['POST'])
|
11 |
+
def classify_image():
|
12 |
+
if 'image' not in request.files:
|
13 |
+
return jsonify({"error": "No image provided"}), 400
|
14 |
+
|
15 |
+
file = request.files['image']
|
16 |
+
if file.filename == '':
|
17 |
+
return jsonify({"error": "Empty image file"}), 400
|
18 |
+
|
19 |
+
image = Image.open(io.BytesIO(file.read()))
|
20 |
+
resized_image = image.copy()
|
21 |
+
resized_image.thumbnail((640, 640))
|
22 |
+
|
23 |
+
# Get results from the model
|
24 |
+
|
25 |
+
results = model(resized_image)[0]
|
26 |
+
|
27 |
+
predictions = []
|
28 |
+
|
29 |
+
if results.boxes is not None:
|
30 |
+
boxes = results.boxes.xyxy
|
31 |
+
confidences = results.boxes.conf
|
32 |
+
classes = results.boxes.cls
|
33 |
+
|
34 |
+
for i in range(len(boxes)):
|
35 |
+
box = boxes[i]
|
36 |
+
confidence = confidences[i].item()
|
37 |
+
class_id = int(classes[i].item())
|
38 |
+
prediction = {
|
39 |
+
"x1": box[0].item(),
|
40 |
+
"y1": box[1].item(),
|
41 |
+
"x2": box[2].item(),
|
42 |
+
"y2": box[3].item(),
|
43 |
+
"confidence": confidence,
|
44 |
+
"class": class_names[class_id],
|
45 |
+
}
|
46 |
+
predictions.append(prediction)
|
47 |
+
|
48 |
+
return jsonify({"predictions": predictions})
|
49 |
+
|
50 |
+
if __name__ == '__main__':
|
51 |
+
app.run(host='127.0.0.1', port=5000, debug=True)
|