Malay911/facemask_detection

1import os2import io3import base644import torch5import torchvision6import numpy as np7import cv28from typing import List, Dict9from contextlib import asynccontextmanager10from PIL import Image11from pydantic import BaseModel12from torchvision import transforms13from fastapi import FastAPI, UploadFile, File, HTTPException14from fastapi.middleware.cors import CORSMiddleware15 16# ---------- Constants & Configuration ----------17 18# Path configuration19BASE_DIR = os.path.dirname(os.path.abspath(__file__))20MODEL_PATH = os.path.join(BASE_DIR, "model", "mask_detector.pth")21 22# Class label mapping: 1=With Mask, 2=Without Mask, 3=Incorrect Mask (0 is background)23CLASS_LABELS = {24    1: "With Mask",25    2: "Without Mask",26    3: "Incorrect Mask"27}28 29# Allowed image MIME types30ALLOWED_TYPES = {"image/jpeg", "image/png", "image/jpg", "image/webp", "image/bmp"}31 32# Global model reference33model = None34 35# ---------- Pydantic Schemas ----------36 37class PredictionResponse(BaseModel):38    """Response schema for single mask detection prediction."""39    prediction: str40    confidence: float41 42class HealthResponse(BaseModel):43    """Response schema for health check endpoint."""44    status: str45 46class Base64Request(BaseModel):47    """Request schema for base64 image prediction."""48    image: str49 50class ErrorResponse(BaseModel):51    """Response schema for error messages."""52    detail: str53 54class FaceResult(BaseModel):55    """Single face detection result with mask classification."""56    label: str57    confidence: float58    bbox: List[int] # [x, y, width, height]59 60class DetectionResponse(BaseModel):61    """Response schema for multi-face detection endpoint."""62    faces_detected: int63    results: List[FaceResult]64    annotated_image: str65 66# ---------- Utility Functions ----------67 68def validate_content_type(content_type: str) -> bool:69    """Check if the uploaded file has a valid image MIME type."""70    return content_type in ALLOWED_TYPES71 72def preprocess_image(image: Image.Image):73    """Preprocess a PIL image for model inference."""74    image = image.convert("RGB")75    tensor = transforms.ToTensor()(image)76    return [tensor]77 78def decode_base64_image(base64_string: str) -> Image.Image:79    """Decode a base64-encoded image string to a PIL Image."""80    if "," in base64_string:81        base64_string = base64_string.split(",", 1)[1]82    try:83        image_bytes = base64.b64decode(base64_string)84        image = Image.open(io.BytesIO(image_bytes))85        return image86    except Exception as e:87        raise ValueError(f"Invalid base64 image: {str(e)}")88 89# ---------- Model Loader ----------90 91def load_model():92    """Load the Faster R-CNN model from disk."""93    global model94    if not os.path.exists(MODEL_PATH):95        raise RuntimeError(f"Model file not found at: {MODEL_PATH}")96    try:97        model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False, min_size=512, max_size=512)98        num_classes = 499        in_features = model.roi_heads.box_predictor.cls_score.in_features100        model.roi_heads.box_predictor = torchvision.models.detection.faster_rcnn.FastRCNNPredictor(in_features, num_classes)101        model.load_state_dict(torch.load(MODEL_PATH, map_location=torch.device("cpu")))102        model.eval()103        print(f"Model loaded successfully from: {MODEL_PATH}")104        return model105    except Exception as e:106        raise RuntimeError(f"Failed to load model: {str(e)}")107 108def get_model():109    """Get the loaded model instance."""110    if model is None:111        raise RuntimeError("Model has not been loaded. Call load_model() first.")112    return model113 114# ---------- Inference Logic ----------115 116def predict(image: Image.Image) -> dict:117    """Run mask detection inference on a PIL image (legacy single-face)."""118    m = get_model()119    input_tensors = preprocess_image(image)120    with torch.no_grad():121        preds = m(input_tensors)122    pred = preds[0]123    pred_labels = pred["labels"].cpu().numpy()124    pred_scores = pred["scores"].cpu().numpy()125    if len(pred_scores) == 0:126        return {"prediction": "No Face Detected", "confidence": 0.0}127    best_idx = pred_scores.argmax()128    best_score = float(pred_scores[best_idx])129    if best_score < 0.5:130        return {"prediction": "No Face Detected", "confidence": best_score}131    label = CLASS_LABELS.get(int(pred_labels[best_idx]), "Unknown")132    return {"prediction": label, "confidence": round(best_score, 4)}133 134def detect_and_classify(image_bytes: bytes) -> dict:135    """Multi-face detection with bounding boxes and annotation."""136    try:137        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")138    except Exception:139        raise ValueError("Could not decode image.")140    frame = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)141    (img_h, img_w) = frame.shape[:2]142    m = get_model()143    input_tensors = preprocess_image(image)144    with torch.no_grad():145        preds = m(input_tensors)146    pred = preds[0]147    pred_boxes = pred["boxes"].cpu().numpy()148    pred_labels = pred["labels"].cpu().numpy()149    pred_scores = pred["scores"].cpu().numpy()150    results = []151    faces_detected = 0152    for bbox, label_idx, score in zip(pred_boxes, pred_labels, pred_scores):153        if score < 0.5:154            continue155        faces_detected += 1156        label = CLASS_LABELS.get(int(label_idx), "Unknown")157        confidence = float(score)158        x1, y1, x2, y2 = map(int, bbox)159        results.append({160            "label": label,161            "confidence": round(confidence, 4),162            "bbox": [x1, y1, x2 - x1, y2 - y1]163        })164        color = (0, 200, 80) if label == "With Mask" else (0, 0, 230) if label == "Without Mask" else (0, 165, 255)165        thickness = max(2, int(min(img_h, img_w) / 250))166        cv2.rectangle(frame, (x1, y1), (x2, y2), color, thickness)167        font_scale = max(0.45, min(img_h, img_w) / 800)168        font_thickness = max(1, int(font_scale * 1.8))169        label_text = f"{label} ({int(confidence * 100)}%)"170        (text_w, text_h), baseline = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, font_thickness)171        pad = int(text_h * 0.35)172        label_y = max(y1 - pad, text_h + pad)173        cv2.rectangle(frame, (x1, label_y - text_h - pad), (x1 + text_w + pad * 2, label_y + pad), color, -1)174        cv2.putText(frame, label_text, (x1 + pad, label_y), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), font_thickness, cv2.LINE_AA)175    _, buffer = cv2.imencode(".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, 95])176    annotated_b64 = base64.b64encode(buffer).decode("utf-8")177    return {178        "faces_detected": faces_detected,179        "results": results,180        "annotated_image": f"data:image/jpeg;base64,{annotated_b64}"181    }182 183# ---------- FastAPI Application ----------184 185@asynccontextmanager186async def lifespan(app: FastAPI):187    """Load models on startup."""188    try:189        load_model()190        print("Face Mask Detection API is ready!")191    except RuntimeError as e:192        print(f"CRITICAL: {e}")193        raise194    yield195    print("Shutting down API...")196 197app = FastAPI(198    title="Face Mask Detection API",199    description="AI-powered face mask compliance monitoring with multi-face detection, analytics, and visualization.",200    version="2.1.0",201    lifespan=lifespan,202)203 204app.add_middleware(205    CORSMiddleware,206    allow_origins=["*"],207    allow_credentials=True,208    allow_methods=["*"],209    allow_headers=["*"],210)211 212@app.get("/health", response_model=HealthResponse, tags=["Health"])213async def health_check():214    return {"status": "API running"}215 216@app.post("/predict", response_model=PredictionResponse, tags=["Prediction"])217async def predict_image(file: UploadFile = File(...)):218    if not validate_content_type(file.content_type):219        raise HTTPException(status_code=400, detail=f"Invalid file type: {file.content_type}")220    contents = await file.read()221    if len(contents) == 0:222        raise HTTPException(status_code=400, detail="Uploaded file is empty.")223    try:224        image = Image.open(io.BytesIO(contents))225        return predict(image)226    except Exception as e:227        raise HTTPException(status_code=500, detail=f"Prediction failed: {str(e)}")228 229@app.post("/predict/base64", response_model=PredictionResponse, tags=["Prediction"])230async def predict_base64(request: Base64Request):231    if not request.image or len(request.image.strip()) == 0:232        raise HTTPException(status_code=400, detail="Base64 image string is empty.")233    try:234        image = decode_base64_image(request.image)235        return predict(image)236    except Exception as e:237        raise HTTPException(status_code=500, detail=f"Prediction failed: {str(e)}")238 239@app.post("/detect", response_model=DetectionResponse, tags=["Detection"])240async def detect_faces(file: UploadFile = File(...)):241    if not validate_content_type(file.content_type):242        raise HTTPException(status_code=400, detail=f"Invalid file type: {file.content_type}")243    contents = await file.read()244    if len(contents) == 0:245        raise HTTPException(status_code=400, detail="Uploaded file is empty.")246    try:247        detection_result = detect_and_classify(contents)248        return detection_result249    except Exception as e:250        raise HTTPException(status_code=500, detail=f"Detection failed: {str(e)}")251 252