Create main.py

main.py

@@ -0,0 +1,175 @@
+import torch
+import torch.nn as nn
+import yaml
+from torchvision import models, transforms
+from PIL import Image
+import gradio as gr
+import os
+import base64
+import io
+import time
+import threading
+from typing import List, Dict, Union, Tuple, Optional
+
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+
+class Base64Image(BaseModel):
+    image_data: str
+
+class BatchBase64Images(BaseModel):
+    image_data_list: List[str]
+    use_gpu: bool = True
+
+CONFIG_PATH: str = os.getenv('CONFIG_PATH', 'staging_config.yaml')
+CHECKPOINT_FILENAME: str = os.getenv('CHECKPOINT_PATH', 'model.pt')
+
+model_lock: threading.Lock = threading.Lock()
+
+def get_model(model_name: str, num_classes: int) -> nn.Module:
+    model: Optional[nn.Module] = None
+    if model_name == "efficientnet_b0":
+        model = models.efficientnet_b0(weights=None)
+        num_ftrs: int = model.classifier[1].in_features
+        model.classifier[1] = nn.Linear(num_ftrs, num_classes)
+    else:
+        raise ValueError(f"Model '{model_name}' not supported.")
+    return model
+
+def load_checkpoint(checkpoint_path: str, device: torch.device) -> Tuple[nn.Module, Dict[int, str]]:
+    if not os.path.exists(checkpoint_path):
+        raise FileNotFoundError(f"Checkpoint file not found at: {checkpoint_path}")
+    checkpoint: dict = torch.load(checkpoint_path, map_location=device)
+    model_name: str = checkpoint['model_name']
+    class_to_idx: Dict[str, int] = checkpoint['class_to_idx']
+    model_output_size: int = 1 if len(class_to_idx) == 2 else len(class_to_idx)
+    model: nn.Module = get_model(model_name, num_classes=model_output_size)
+    model.load_state_dict(checkpoint['state_dict'])
+    model.to(device)
+    model.eval()
+    idx_to_class: Dict[int, str] = {v: k for k, v in class_to_idx.items()}
+    return model, idx_to_class
+
+try:
+    with open(CONFIG_PATH, 'r') as f:
+        config: dict = yaml.safe_load(f)
+except FileNotFoundError:
+    raise RuntimeError(f"ERROR: Config file not found at '{CONFIG_PATH}'. Make sure it's uploaded to the Space.")
+
+if torch.cuda.is_available():
+    gpu_device: torch.device = torch.device("cuda")
+    gpu_model: nn.Module
+    IDX_TO_CLASS: Dict[int, str]
+    gpu_model, IDX_TO_CLASS = load_checkpoint(CHECKPOINT_FILENAME, gpu_device)
+    print(f"GPU model loaded successfully on {gpu_device}")
+else:
+    gpu_device: Optional[torch.device] = None
+    gpu_model: Optional[nn.Module] = None
+    print("No GPU available")
+
+cpu_device: torch.device = torch.device("cpu")
+cpu_model: nn.Module
+IDX_TO_CLASS: Dict[int, str]
+cpu_model, IDX_TO_CLASS = load_checkpoint(CHECKPOINT_FILENAME, cpu_device)
+print(f"CPU model loaded successfully")
+print(f"Class mapping: {IDX_TO_CLASS}")
+
+IMG_SIZE: int = config['data_params']['image_size']
+inference_transform: transforms.Compose = transforms.Compose([
+    transforms.Resize((IMG_SIZE, IMG_SIZE)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
+def base64_to_pil(base64_str: str) -> Image.Image:
+    try:
+        if "base64," in base64_str:
+            base64_str = base64_str.split("base64,")[1]
+        image_data: bytes = base64.b64decode(base64_str)
+        return Image.open(io.BytesIO(image_data))
+    except Exception as e:
+        raise ValueError(f"Invalid base64 string: {e}")
+
+def predict_batch(pil_images: List[Image.Image], use_gpu: bool) -> List[Dict[str, Union[dict, float]]]:
+    device: torch.device = gpu_device if (use_gpu and gpu_device) else cpu_device
+    model: nn.Module = gpu_model if (use_gpu and gpu_model) else cpu_model
+    image_tensors: List[torch.Tensor] = []
+    for img in pil_images:
+        if img.mode != "RGB":
+            img = img.convert("RGB")
+        image_tensors.append(inference_transform(img))
+    batch_tensor: torch.Tensor = torch.stack(image_tensors).to(device)
+    with model_lock, torch.no_grad():
+        start_time: float = time.time()
+        output: torch.Tensor = model(batch_tensor)
+        batch_time: float = time.time() - start_time
+    results: List[Dict[str, Union[dict, float]]] = []
+    probs: Union[List[float], float] = torch.sigmoid(output).squeeze().tolist()
+    class_0_name: str = IDX_TO_CLASS.get(0, "Class 0")
+    class_1_name: str = IDX_TO_CLASS.get(1, "Class 1")
+    if isinstance(probs, float):
+        probs = [probs]
+    for i, prob in enumerate(probs):
+        prediction: Dict[str, float] = {
+            class_0_name: 1 - prob,
+            class_1_name: prob
+        }
+        metadata: Dict[str, Union[str, float]] = {
+            "device": "gpu" if use_gpu and gpu_device else "cpu",
+            "inference_ms": batch_time * 1000 / len(pil_images),
+            "image_size": f"{pil_images[i].width}x{pil_images[i].height}"
+        }
+        results.append({
+            "prediction": prediction,
+            "metadata": metadata
+        })
+    return results
+
+app: FastAPI = FastAPI(
+    title="Image Classifier API",
+    description="A FastAPI server with a Gradio UI for image classification. Supports batch processing",
+)
+
+app.add_middleware(CORSMiddleware, 
+                   allow_origins=["*"], 
+                   allow_credentials=True, 
+                   allow_methods=["*"], 
+                   allow_headers=["*"])
+
+@app.post("/predict", response_model=dict)
+async def predict_api(request: Base64Image, use_gpu: bool = True) -> dict:
+    try:
+        pil_image: Image.Image = base64_to_pil(request.image_data)
+        result: Dict[str, Union[dict, float]] = predict_batch([pil_image], use_gpu)[0]
+        return result
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+@app.post("/batch_predict", response_model=List[dict])
+async def batch_predict_api(request: BatchBase64Images) -> List[dict]:
+    try:
+        pil_images: List[Image.Image] = []
+        for base64_str in request.image_data_list:
+            pil_images.append(base64_to_pil(base64_str))
+        results: List[Dict[str, Union[dict, float]]] = predict_batch(pil_images, request.use_gpu)
+        return results
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+def predict_from_pil(pil_image: Image.Image) -> Optional[dict]:
+    if pil_image is None:
+        return None
+    result: Dict[str, Union[dict, float]] = predict_batch([pil_image], use_gpu=True)[0]
+    return result["prediction"]
+
+gradio_iface: gr.Interface = gr.Interface(
+    fn=predict_from_pil,
+    inputs=gr.Image(type="pil", label="Upload Image"),
+    outputs=gr.Label(num_top_classes=2, label="Predictions"),
+    title="Image Classifier",
+    description="Upload an image to see its classification. The API is available at the /docs endpoint.",
+    allow_flagging="never" 
+)
+
+app = gr.mount_gradio_app(app, gradio_iface, path="/")