Update app.py

app.py

@@ -2,42 +2,35 @@ from fastapi import FastAPI, UploadFile, File
 from fastapi.responses import JSONResponse
 from PIL import Image
 import io
-
-from transformers import AutoProcessor, AutoModelForImageClassification
 import torch
+import os
+
+from transformers import AutoFeatureExtractor, AutoModelForImageClassification
+
+# Fix permissions by setting cache location
+os.environ["HF_HOME"] = "/app/hf_cache"
 
-# Load model and processor (only once at startup)
 MODEL_NAME = "google/vit-base-patch16-224"
-processor = AutoProcessor.from_pretrained(MODEL_NAME)
+extractor = AutoFeatureExtractor.from_pretrained(MODEL_NAME)
 model = AutoModelForImageClassification.from_pretrained(MODEL_NAME)
 
 app = FastAPI()
 
-@app.get("/")
-def root():
-    return {"message": "Send an image to POST /analyze-image"}
-
 @app.post("/analyze-image")
 async def analyze_image(file: UploadFile = File(...)):
     try:
-        # Read and convert the uploaded image
         contents = await file.read()
         image = Image.open(io.BytesIO(contents)).convert("RGB")
 
-        # Preprocess
-        inputs = processor(images=image, return_tensors="pt")
+        inputs = extractor(images=image, return_tensors="pt")
         with torch.no_grad():
             outputs = model(**inputs)
 
-        # Get top prediction
         logits = outputs.logits
         predicted_class_id = logits.argmax(-1).item()
         label = model.config.id2label[predicted_class_id]
 
-        return JSONResponse(content={
-            "filename": file.filename,
-            "predicted_label": label
-        })
+        return JSONResponse(content={"filename": file.filename, "predicted_label": label})
 
     except Exception as e:
         return JSONResponse(content={"error": str(e)}, status_code=400)