|
|
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="/") |
