Update main.py

main.py

@@ -1,38 +1,43 @@
-import gradio as gr
-from transformers import AutoImageProcessor, AutoModelForImageClassification
-from PIL import Image
+import io
 import torch
+from PIL import Image
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from transformers import AutoImageProcessor, AutoModelForImageClassification
 
-# Load model and processor
+# Log model loading
+print("🚀 Starting model download...")
 model_id = "sheikh987/Skin_Cancer-Image_Classification"
 processor = AutoImageProcessor.from_pretrained(model_id)
 model = AutoModelForImageClassification.from_pretrained(model_id)
+print("✅ Model loaded successfully.")
+
+# FastAPI app
+app = FastAPI(title="Skin Cancer Classifier API")
 
-# Prediction function
-def classify_image(img):
-    inputs = processor(images=img, return_tensors="pt")
+# Health check endpoint
+@app.get("/status")
+def status():
+    return {"status": "ok", "model": model_id}
+
+# Image classification endpoint
+@app.post("/predict/")
+async def predict(file: UploadFile = File(...)):
+    if not file.content_type.startswith("image/"):
+        raise HTTPException(status_code=400, detail="Invalid image file")
+
+    try:
+        image_bytes = await file.read()
+        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+    except Exception:
+        raise HTTPException(status_code=400, detail="Could not decode image")
+
+    inputs = processor(images=image, return_tensors="pt")
     with torch.no_grad():
         outputs = model(**inputs)
+
     logits = outputs.logits
-    
-    # Get the top prediction
-    predicted_class_idx = logits.argmax(-1).item()
-    predicted_class = model.config.id2label[predicted_class_idx]
-    
-    # Get the confidence score for the top prediction
-    confidence = torch.nn.functional.softmax(logits, dim=-1)[0][predicted_class_idx].item()
-    
-    # Return the data in the format Gradio's Label component expects
-    return {predicted_class: confidence}
-
-# Gradio interface
-interface = gr.Interface(
-    fn=classify_image,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Label(num_top_classes=5), # Show top 5 classes
-    title="Skin Cancer Image Classifier",
-    description="Upload an image of skin lesion to classify."
-)
-
-# Launch the app and the API
-interface.launch()
+    idx = logits.argmax(-1).item()
+    label = model.config.id2label[idx]
+    confidence = torch.nn.functional.softmax(logits, dim=-1)[0][idx].item()
+
+    return {"label": label, "confidence": round(confidence, 4)}