Update index.py

index.py

@@ -1,51 +1,73 @@
+from fastapi import FastAPI, File, UploadFile, HTTPException
 from ultralytics import YOLO
-from flask import Flask, request, jsonify
-from PIL import Image, ImageDraw
+from PIL import Image
 import io
+from typing import List
+from pydantic import BaseModel
+import uvicorn
 
-app = Flask(__name__)
-model = YOLO('last.pt')  # Load your YOLO model
+# Define the response model
+class Prediction(BaseModel):
+    x1: float
+    y1: float
+    x2: float
+    y2: float
+    confidence: float
+    class_name: str
+
+class PredictionResponse(BaseModel):
+    predictions: List[Prediction]
+
+# Initialize FastAPI app and model
+app = FastAPI(title="Skin Condition Detection API")
+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
+@app.post("/classify", response_model=PredictionResponse)
+async def classify_image(file: UploadFile = File(...)):
+    """
+    Endpoint to classify skin conditions in an uploaded image
+    """
+    # Validate file
+    if not file.content_type.startswith("image/"):
+        raise HTTPException(status_code=400, detail="File must be an image")
+    
+    try:
+        # Read and process image
+        contents = await file.read()
+        image = Image.open(io.BytesIO(contents))
+        resized_image = image.copy()
+        resized_image.thumbnail((640, 640))
+        
+        # Get predictions
+        results = model(resized_image)[0]
+        predictions = []
 
-    file = request.files['image']
-    if file.filename == '':
-        return jsonify({"error": "Empty image file"}), 400
+        if results.boxes is not None:
+            boxes = results.boxes.xyxy
+            confidences = results.boxes.conf
+            classes = results.boxes.cls
 
-    image = Image.open(io.BytesIO(file.read()))
-    resized_image = image.copy()
-    resized_image.thumbnail((640, 640))
+            for i in range(len(boxes)):
+                box = boxes[i]
+                confidence = confidences[i].item()
+                class_id = int(classes[i].item())
+                
+                prediction = Prediction(
+                    x1=float(box[0].item()),
+                    y1=float(box[1].item()),
+                    x2=float(box[2].item()),
+                    y2=float(box[3].item()),
+                    confidence=confidence,
+                    class_name=class_names[class_id]
+                )
+                predictions.append(prediction)
 
-    # 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)
+        return PredictionResponse(predictions=predictions)
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+# For local development
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)