Upload app.py

app.py

@@ -0,0 +1,75 @@
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import JSONResponse
+from PIL import Image
+from io import BytesIO
+import torch
+import torch.nn.functional as F
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+
+app = FastAPI()
+
+@app.get("/")
+async def root():
+    return {"message": "API is running"}
+
+# Load model and processor
+MODEL_NAME = "ivandrian11/fruit-classifier"
+processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
+model = AutoModelForImageClassification.from_pretrained(MODEL_NAME)
+model.eval()  # set to evaluation mode
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(DEVICE)
+
+VALID_CLASSES = ['apple', 'banana', 'orange', 'tomato', 'bitter gourd', 'capsicum']
+
+CLASS_MAPPING = {
+    'apple': 'apple',
+    'banana': 'banana',
+    'orange': 'orange',
+    'tomato': 'tomato',
+    'bitter gourd': 'bitter gourd',
+    'bitter melon': 'bitter gourd',
+    'bell pepper': 'capsicum',
+    'pepper': 'capsicum',
+    'capsicum': 'capsicum',
+    'green pepper': 'capsicum',
+    'red pepper': 'capsicum',
+    'yellow pepper': 'capsicum',
+    'granny smith': 'apple',
+    'fuji apple': 'apple',
+    'gala apple': 'apple',
+    'navel orange': 'orange',
+    'valencia orange': 'orange'
+}
+
+def classify_fruit(image: Image.Image) -> str:
+    inputs = processor(images=image, return_tensors="pt").to(DEVICE)
+    with torch.no_grad():
+        outputs = model(**inputs)
+        probabilities = F.softmax(outputs.logits, dim=-1)
+        confidence, predicted_idx = torch.max(probabilities, dim=-1)
+        confidence = confidence.item()
+        predicted_label = model.config.id2label[predicted_idx.item()].lower()
+
+    if confidence < 0.7:
+        return "unknown"
+
+    mapped_class = CLASS_MAPPING.get(predicted_label, None)
+    if mapped_class:
+        return mapped_class
+
+    for valid_class in VALID_CLASSES:
+        if valid_class in predicted_label:
+            return valid_class
+
+    return "unknown"
+
+@app.post("/classify")
+async def classify_image(file: UploadFile = File(...)):
+    try:
+        image_bytes = await file.read()
+        image = Image.open(BytesIO(image_bytes)).convert("RGB")
+        result = classify_fruit(image)
+        return JSONResponse(content={"prediction": result})
+    except Exception as e:
+        return JSONResponse(content={"prediction": "unknown", "error": str(e)}, status_code=500)