Update app.py

app.py

@@ -3,24 +3,41 @@ from fastapi.responses import JSONResponse
 from PIL import Image
 import io
 
+from transformers import AutoProcessor, AutoModelForImageClassification
+import torch
+
+# Load model and processor (only once at startup)
+MODEL_NAME = "google/vit-base-patch16-224"
+processor = AutoProcessor.from_pretrained(MODEL_NAME)
+model = AutoModelForImageClassification.from_pretrained(MODEL_NAME)
+
 app = FastAPI()
 
 @app.get("/")
-def read_root():
-    return {"message": "Space is running. Use POST /process-image"}
+def root():
+    return {"message": "Send an image to POST /analyze-image"}
 
-@app.post("/process-image")
-async def process_image(file: UploadFile = File(...)):
+@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))
-        width, height = image.size
+        image = Image.open(io.BytesIO(contents)).convert("RGB")
+
+        # Preprocess
+        inputs = processor(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,
-            "width": width,
-            "height": height,
-            "format": image.format
+            "predicted_label": label
         })
+
     except Exception as e:
         return JSONResponse(content={"error": str(e)}, status_code=400)