Create app.py

app.py

@@ -0,0 +1,60 @@
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+import torch
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+from PIL import Image
+import aiohttp
+import io
+import base64
+
+app = FastAPI()
+
+# Load model (skshmjn/Pokemon-classifier-gen9-1025)
+MODEL_NAME = "skshmjn/Pokemon-classifier-gen9-1025"
+processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
+model = AutoModelForImageClassification.from_pretrained(MODEL_NAME)
+
+class ClassifyRequest(BaseModel):
+    image_url: str = None
+    image_data: str = None  # base64
+
+@app.get("/health")
+async def health():
+    return {"status": "ok", "model": MODEL_NAME}
+
+@app.post("/classify")
+async def classify(request: ClassifyRequest):
+    try:
+        # Load image
+        if request.image_url:
+            async with aiohttp.ClientSession() as session:
+                async with session.get(request.image_url) as resp:
+                    image_bytes = await resp.read()
+                    image = Image.open(io.BytesIO(image_bytes))
+        elif request.image_data:
+            image_bytes = base64.b64decode(request.image_data)
+            image = Image.open(io.BytesIO(image_bytes))
+        else:
+            raise HTTPException(400, "No image provided")
+        
+        # Preprocess
+        inputs = processor(images=image, return_tensors="pt")
+        
+        # Predict
+        with torch.no_grad():
+            outputs = model(**inputs)
+            logits = outputs.logits
+            probs = torch.softmax(logits, dim=1)
+            top_prob, top_idx = torch.max(probs, dim=1)
+        
+        # Get Pokemon name
+        pokemon_name = model.config.id2label[top_idx.item()]
+        confidence = top_prob.item()
+        
+        return {
+            "name": pokemon_name.lower(),
+            "confidence": confidence
+        }
+    
+    except Exception as e:
+        raise HTTPException(500, str(e))