Update app.py

app.py

@@ -7,7 +7,7 @@ import torch
 import uvicorn
 from fastapi import FastAPI, UploadFile, File
 from fastapi.responses import JSONResponse
-from transformers import AutoModelForImageClassification, AutoImageProcessor
+from transformers import AutoConfig, AutoModelForImageClassification, AutoImageProcessor
 from PIL import Image
 
 # ================== Config ==================
@@ -25,24 +25,38 @@ else:
 
 # ================== Load Model ==================
 print("📦 Loading model...")
+
+# Load config first to avoid HF hub errors
+config = AutoConfig.from_pretrained(str(MODEL_DIR), local_files_only=True)
+
+# Load model with local config
 model = AutoModelForImageClassification.from_pretrained(
     str(MODEL_DIR),
+    config=config,
     local_files_only=True
 )
+
+# Load image processor
 processor = AutoImageProcessor.from_pretrained(
     str(MODEL_DIR),
     local_files_only=True
 )
+
 model.eval()
 print("✅ Model loaded successfully")
 
 # ================== FastAPI App ==================
-app = FastAPI()
+app = FastAPI(title="Cattle/Buffalo Classifier API")
 
 @app.get("/")
 def home():
     return {"message": "🚀 Cattle/Buffalo Model API is running!"}
 
+@app.get("/labels")
+def get_labels():
+    """Return all possible class labels."""
+    return {"labels": list(model.config.id2label.values())}
+
 @app.post("/predict")
 async def predict(file: UploadFile = File(...)):
     try:
@@ -52,12 +66,13 @@ async def predict(file: UploadFile = File(...)):
         # Preprocess
         inputs = processor(images=image, return_tensors="pt")
 
+        # Forward pass
         with torch.no_grad():
             outputs = model(**inputs)
             logits = outputs.logits
             predicted_class_idx = logits.argmax(-1).item()
 
-        # Map to labels
+        # Map to label
         label = model.config.id2label[predicted_class_idx]
 
         return JSONResponse(content={"prediction": label})