Upload app.py

app.py

@@ -3,17 +3,16 @@ import base64
 import torch
 from PIL import Image
 from fastapi import FastAPI, File, UploadFile, HTTPException
-from fastapi.responses import JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from ultralytics import YOLO
 from transformers import AutoImageProcessor, AutoModelForImageClassification
 
-# --- Configuration ---
+# --- App Config ---
 app = FastAPI(
     title="Food & Vegetable AI API",
-    description="Combined API for Food Classification (ViT) and Fruit/Veg Detection (YOLO)",
-    version="2.0.0"
+    description="Separate APIs for ViT Classification and YOLO Detection",
+    version="2.1.0"
 )
 
 app.add_middleware(
@@ -29,83 +28,127 @@ class Base64ImageRequest(BaseModel):
 
 # --- Model Loading ---
 print("Loading models...")
+
 try:
-    # 1. ViT Food Classifier
-    vit_model = AutoModelForImageClassification.from_pretrained("eslamxm/vit-base-food101")
-    vit_processor = AutoImageProcessor.from_pretrained("eslamxm/vit-base-food101")
-    
-    # 2. YOLO Fruit/Veg Detector
-    # Ensure this file is in your root directory
-    yolo_model = YOLO('yolo_fruits_and_vegetables_v3.pt')
-    
-    print("✓ All models loaded successfully!")
+    vit_model = AutoModelForImageClassification.from_pretrained(
+        "eslamxm/vit-base-food101"
+    )
+    vit_processor = AutoImageProcessor.from_pretrained(
+        "eslamxm/vit-base-food101"
+    )
+
+    yolo_model = YOLO("yolo_fruits_and_vegetables_v3.pt")
+
+    print("✓ Models loaded successfully")
 except Exception as e:
-    print(f"✗ Error loading models: {e}")
+    print(f"✗ Model loading failed: {e}")
     vit_model = None
     yolo_model = None
 
-# --- Utility Functions ---
-def process_pil_image(image: Image.Image):
-    """Common logic for YOLO detection and ViT classification."""
-    results = {"detection": None, "classification": None}
-    
-    # YOLO Inference
-    if yolo_model:
-        y_results = yolo_model(image)
-        detections = []
-        summary = {}
-        for r in y_results:
-            for i in range(len(r.boxes)):
-                label = yolo_model.names[int(r.boxes.cls[i])]
-                detections.append({
-                    "label": label,
-                    "confidence": float(r.boxes.conf[i]),
-                    "bbox": r.boxes.xyxy[i].tolist()
-                })
-                summary[label] = summary.get(label, 0) + 1
-        results["detection"] = {"detections": detections, "summary": summary}
-
-    # ViT Inference
-    if vit_model:
-        inputs = vit_processor(images=image, return_tensors="pt")
-        with torch.no_grad():
-            outputs = vit_model(**inputs)
-        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
-        pred_id = probs.argmax().item()
-        results["classification"] = {
-            "label": vit_model.config.id2label[pred_id],
-            "confidence": round(probs[0][pred_id].item(), 4)
-        }
-        
-    return results
-
-# --- Endpoints ---
 
+# --- Utility ---
+def load_image_from_bytes(image_bytes: bytes) -> Image.Image:
+    return Image.open(io.BytesIO(image_bytes)).convert("RGB")
+
+
+# --- YOLO Detection ---
+def run_yolo(image: Image.Image):
+    if not yolo_model:
+        raise HTTPException(status_code=500, detail="YOLO model not loaded")
+
+    results = yolo_model(image)
+
+    detections = []
+    summary = {}
+
+    for r in results:
+        for i in range(len(r.boxes)):
+            label = yolo_model.names[int(r.boxes.cls[i])]
+            detections.append({
+                "label": label,
+                "confidence": float(r.boxes.conf[i]),
+                "bbox": r.boxes.xyxy[i].tolist()
+            })
+            summary[label] = summary.get(label, 0) + 1
+
+    return {
+        "detections": detections,
+        "summary": summary
+    }
+
+
+# --- ViT Classification ---
+def run_vit(image: Image.Image):
+    if not vit_model:
+        raise HTTPException(status_code=500, detail="ViT model not loaded")
+
+    inputs = vit_processor(images=image, return_tensors="pt")
+
+    with torch.no_grad():
+        outputs = vit_model(**inputs)
+
+    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
+    pred_id = probs.argmax().item()
+
+    return {
+        "label": vit_model.config.id2label[pred_id],
+        "confidence": round(probs[0][pred_id].item(), 4)
+    }
+
+
+# --- Routes ---
 @app.get("/")
 async def root():
-    return {"message": "AI API is online. Use /predict-upload or /predict-base64."}
+    return {
+        "message": "API running",
+        "endpoints": ["/predict-vit", "/predict-yolo"]
+    }
 
-@app.post("/predict-upload")
-async def predict_file(file: UploadFile = File(...)):
-    """Upload a raw image file for full analysis."""
+
+# ---------- YOLO Endpoint ----------
+@app.post("/predict-yolo")
+async def predict_yolo(file: UploadFile = File(...)):
     try:
-        image = Image.open(file.file).convert("RGB")
-        return process_pil_image(image)
+        image = load_image_from_bytes(await file.read())
+        return {"detection": run_yolo(image)}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-@app.post("/predict-base64")
-async def predict_base64(request: Base64ImageRequest):
-    """Send a base64 string (useful for mobile apps)."""
+
+# ---------- ViT Endpoint ----------
+@app.post("/predict-vit")
+async def predict_vit(file: UploadFile = File(...)):
     try:
-        header, encoded = request.image.split(",", 1) if "," in request.image else (None, request.image)
+        image = load_image_from_bytes(await file.read())
+        return {"classification": run_vit(image)}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+# ---------- Base64 Support (optional) ----------
+@app.post("/predict-vit-base64")
+async def predict_vit_base64(request: Base64ImageRequest):
+    try:
+        _, encoded = request.image.split(",", 1) if "," in request.image else (None, request.image)
         image_bytes = base64.b64decode(encoded)
-        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-        return process_pil_image(image)
+        image = load_image_from_bytes(image_bytes)
+        return {"classification": run_vit(image)}
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Base64 processing failed: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"ViT base64 error: {str(e)}")
+
+
+@app.post("/predict-yolo-base64")
+async def predict_yolo_base64(request: Base64ImageRequest):
+    try:
+        _, encoded = request.image.split(",", 1) if "," in request.image else (None, request.image)
+        image_bytes = base64.b64decode(encoded)
+        image = load_image_from_bytes(image_bytes)
+        return {"detection": run_yolo(image)}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"YOLO base64 error: {str(e)}")
+
 
+# --- Run ---
 if __name__ == "__main__":
     import uvicorn
-    # 7860 is the required port for Hugging Face Spaces
     uvicorn.run(app, host="0.0.0.0", port=7860)