Update main.py

main.py

@@ -1,55 +1,75 @@
-import os
 import io
-import logging
-from typing import Tuple
-
+import torch
+import timm
+from PIL import Image
 from fastapi import FastAPI, File, UploadFile, HTTPException
-from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
-from PIL import Image
+from torchvision import transforms
 
-# Roboflow inference
-from inference import get_model
+# =========================
+# App Init
+# =========================
+app = FastAPI(title="Vehicle Type Classifier API")
 
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger("vehicle-predictor")
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print("Using device:", device)
 
-# FastAPI setup
-app = FastAPI(title="Vehicle Type Predictor")
+MODEL_PATH = "vehicle_classifier_best.pth"
+model = None
+class_names = None
 
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],  # you can tighten this later if needed
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
 
-# Load Roboflow model at startup
-ROBOFLOW_API_KEY = os.environ.get("ROBOFLOW_API_KEY")
-MODEL_ID = "vehicle-classification-eapcd/19"
-
-if ROBOFLOW_API_KEY is None:
-    logger.error("❌ ROBOFLOW_API_KEY not found in environment variables")
-    model = None
-else:
-    try:
-        logger.info("🚀 Loading Roboflow model...")
-        model = get_model(model_id=MODEL_ID, api_key=ROBOFLOW_API_KEY)
-        logger.info("✅ Roboflow model loaded successfully")
-    except Exception as e:
-        logger.exception("❌ Failed to load Roboflow model")
-        model = None
-
-
-# Response model
+# =========================
+# Response Schema
+# =========================
 class PredictionResponse(BaseModel):
     label: str
     confidence: float
 
 
+# =========================
+# Image Transform
+# =========================
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(
+        mean=[0.485, 0.456, 0.406],
+        std=[0.229, 0.224, 0.225]
+    )
+])
+
+
+# =========================
+# Load Model on Startup
+# =========================
+@app.on_event("startup")
+def load_model():
+    global model, class_names
+
+    checkpoint = torch.load(MODEL_PATH, map_location=device)
+    class_names = checkpoint["class_names"]
+    num_classes = len(class_names)
+
+    model = timm.create_model(
+        "convnext_tiny",
+        pretrained=False,
+        num_classes=num_classes
+    )
+
+    model.load_state_dict(checkpoint["model_state_dict"])
+    model.to(device)
+    model.eval()
+
+    print("✅ Model loaded successfully")
+
+
+# =========================
+# Prediction Endpoint
+# =========================
 @app.post("/predict", response_model=PredictionResponse)
 async def predict(file: UploadFile = File(...)):
+
     if model is None:
         raise HTTPException(status_code=503, detail="Model not loaded")
 
@@ -57,30 +77,31 @@ async def predict(file: UploadFile = File(...)):
         raise HTTPException(status_code=400, detail="File must be an image")
 
     try:
-        contents = await file.read()
-
-        # Roboflow accepts PIL Image directly
-        img = Image.open(io.BytesIO(contents)).convert("RGB")
+        image_bytes = await file.read()
+        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
 
-        # Run inference
-        result = model.infer(img)
+        img_tensor = transform(image).unsqueeze(0).to(device)
 
-        if not result.get("predictions"):
-            raise HTTPException(status_code=500, detail="No predictions returned")
+        with torch.no_grad():
+            outputs = model(img_tensor)
+            probs = torch.softmax(outputs, dim=1)
+            pred_idx = torch.argmax(probs, dim=1).item()
 
-        # Take top prediction
-        pred = result["predictions"][0]
-        label = pred.get("class", "Unknown")
-        confidence = float(pred.get("confidence", 0.0))
+        label = class_names[pred_idx]
+        confidence = probs[0][pred_idx].item()
 
-        logger.info(f"Predicted {label} ({confidence:.4f}) for {file.filename}")
-        return PredictionResponse(label=label, confidence=confidence)
+        return PredictionResponse(
+            label=label,
+            confidence=round(confidence, 4)
+        )
 
     except Exception as e:
-        logger.exception("Prediction failed")
-        raise HTTPException(status_code=500, detail="Prediction failed")
+        raise HTTPException(status_code=500, detail=str(e))
 
 
+# =========================
+# Health Check
+# =========================
 @app.get("/health")
-def health():
-    return {"status": "ok", "model_loaded": model is not None}
+def root():
+    return {"status": "Vehicle Classifier API is running 🚗🏍️"}