lakshayC007/PressureUlcerClassifierDocker

1import io2import os3import gc4import torch5import torch.nn as nn6import torchvision.models as models7import torchvision.transforms as transforms8from PIL import Image9from fastapi import FastAPI, File, UploadFile, HTTPException10from fastapi.middleware.cors import CORSMiddleware11from pydantic import BaseModel12from typing import Dict, Any13 14# Initialize FastAPI app15app = FastAPI(16    title="Pressure Ulcer Classifier API",17    description="API for classifying pressure ulcer images into 6 categories",18    version="1.0.0"19)20 21# Add CORS middleware22app.add_middleware(23    CORSMiddleware,24    allow_origins=["*"],  # Allow all origins25    allow_credentials=True,26    allow_methods=["*"],  # Allow all methods27    allow_headers=["*"],  # Allow all headers28)29 30# Configuration31IMAGE_SIZE = 25632DEVICE = torch.device("cpu")  # Force CPU to reduce memory usage33model = None34class_names = None35 36# Define transforms for image preprocessing37transform = transforms.Compose([38    transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),39    transforms.ToTensor(),40    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])41])42 43# Create model architecture44def create_model(num_classes=6):45    # Using ResNet18 architecture46    model = models.resnet18(weights=None)47    num_ftrs = model.fc.in_features48    model.fc = nn.Linear(num_ftrs, num_classes)49    return model50 51# Lazy loading - only load the model when needed52def get_model():53    global model, class_names54    55    if model is None:56        try:57            print("Loading model...")58            model = create_model()59            checkpoint = torch.load('final_cross_validated_model.pth', map_location=DEVICE)60            61            # Free memory62            torch.cuda.empty_cache()63            gc.collect()64            65            if 'model_state_dict' in checkpoint:66                model.load_state_dict(checkpoint['model_state_dict'])67                class_names = checkpoint.get('class_names', ['he', 's1', 's2', 's3', 's4', 'un'])68            else:69                model.load_state_dict(checkpoint)70                class_names = ['he', 's1', 's2', 's3', 's4', 'un']71            72            # Free memory from checkpoint73            del checkpoint74            torch.cuda.empty_cache()75            gc.collect()76            77            model.to(DEVICE)78            model.eval()79            print("Model loaded successfully")80        except Exception as e:81            print(f"Error loading model: {e}")82            raise83    84    return model, class_names85 86# Response model87class PredictionResponse(BaseModel):88    prediction: str89    confidence: float90    class_confidences: Dict[str, float]91 92# Root endpoint93@app.get("/")94async def root():95    return {96        "message": "Pressure Ulcer Classification API is running",97        "endpoints": {98            "/predict": "POST - Upload an image for classification",99            "/health": "GET - Health check endpoint"100        }101    }102 103# Health check endpoint104@app.get("/health")105async def health_check():106    return {"status": "healthy"}107 108# Prediction endpoint109@app.post("/predict", response_model=PredictionResponse)110async def predict(file: UploadFile = File(...)):111    try:112        # Read image file113        contents = await file.read()114        115        # Open and preprocess the image116        img = Image.open(io.BytesIO(contents)).convert('RGB')117        118        # Get model119        model, class_names = get_model()120        121        # Preprocess image122        img_tensor = transform(img).unsqueeze(0).to(DEVICE)123        124        # Make prediction125        with torch.no_grad():126            outputs = model(img_tensor)127            _, preds = torch.max(outputs, 1)128            probability = torch.nn.functional.softmax(outputs, dim=1)[0]129        130        # Get prediction details131        pred_class = class_names[preds[0]]132        pred_confidence = float(probability[preds[0]].item())133        134        # Create confidence scores for all classes135        class_confidences = {136            class_name: float(probability[i].item()) 137            for i, class_name in enumerate(class_names)138        }139        140        # Free memory141        del img_tensor, outputs, preds, probability142        gc.collect()143        144        # Return prediction result145        return {146            "prediction": pred_class,147            "confidence": pred_confidence,148            "class_confidences": class_confidences149        }150        151    except Exception as e:152        raise HTTPException(status_code=500, detail=str(e))153 154# Run the server155if __name__ == "__main__":156    import uvicorn157       # In app.py158    port = int(os.environ.get("PORT", 7860))  # Use 7860 not 8000159    uvicorn.run("app:app", host="0.0.0.0", port=port, reload=False)