Update app.py

app.py

@@ -1,73 +1,99 @@
-from fastapi import FastAPI, File, UploadFile, HTTPException
+import streamlit as st
 from ultralytics import YOLO
-from PIL import Image
+from PIL import Image, ImageDraw
 import io
-from typing import List
-from pydantic import BaseModel
-import uvicorn
+import json
 
-# Define the response model
-class Prediction(BaseModel):
-    x1: float
-    y1: float
-    x2: float
-    y2: float
-    confidence: float
-    class_name: str
+# Set page config
+st.set_page_config(page_title="Skin Condition Classifier", layout="wide")
 
-class PredictionResponse(BaseModel):
-    predictions: List[Prediction]
+# Initialize the model
+@st.cache_resource
+def load_model():
+    return YOLO('best.pt')
 
-# Initialize FastAPI app and model
-app = FastAPI(title="Skin Condition Detection API")
-model = YOLO('best.pt')
+model = load_model()
 class_names = ['Acne', 'Dark circles', 'blackheads', 'eczema', 'rosacea', 'whiteheads', 'wrinkles']
 
-@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")
+def process_image(image):
+    # Resize image
+    resized_image = image.copy()
+    resized_image.thumbnail((640, 640))
     
-    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 = []
+    
+    if results.boxes is not None:
+        boxes = results.boxes.xyxy
+        confidences = results.boxes.conf
+        classes = results.boxes.cls
         
-        # Get predictions
-        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 = 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)
+        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 predictions
 
-        return PredictionResponse(predictions=predictions)
+def draw_boxes(image, predictions):
+    draw = ImageDraw.Draw(image)
+    for pred in predictions:
+        x1, y1, x2, y2 = pred["x1"], pred["y1"], pred["x2"], pred["y2"]
+        draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
+        label = f"{pred['class']} ({pred['confidence']:.2f})"
+        draw.text((x1, y1-10), label, fill="red")
+    return image
 
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+def main():
+    st.title("Skin Condition Classifier")
+    
+    # File uploader
+    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+    
+    if uploaded_file is not None:
+        # Create columns for layout
+        col1, col2 = st.columns(2)
+        
+        # Display original image
+        image = Image.open(uploaded_file)
+        col1.subheader("Original Image")
+        col1.image(image, use_column_width=True)
+        
+        # Process image and get predictions
+        predictions = process_image(image)
+        
+        # Draw boxes on image
+        image_with_boxes = image.copy()
+        image_with_boxes = draw_boxes(image_with_boxes, predictions)
+        
+        # Display annotated image
+        col2.subheader("Detected Conditions")
+        col2.image(image_with_boxes, use_column_width=True)
+        
+        # Display predictions in a nice format
+        st.subheader("Detailed Results")
+        for pred in predictions:
+            st.write(f"- Detected {pred['class']} with {pred['confidence']:.2f} confidence")
+        
+        # Add download button for JSON results
+        if predictions:
+            json_results = json.dumps({"predictions": predictions}, indent=2)
+            st.download_button(
+                label="Download Results as JSON",
+                data=json_results,
+                file_name="predictions.json",
+                mime="application/json"
+            )
 
-# For local development
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    main()