Create main.py

main.py

@@ -0,0 +1,107 @@
+# main.py
+
+import base64, io, json, numpy as np, torch
+from fastapi import FastAPI, HTTPException
+from huggingface_hub import hf_hub_download
+from PIL import Image
+from pydantic import BaseModel
+from torchvision import transforms
+import timm, segmentation_models_pytorch as smp, albumentations as A
+from albumentations.pytorch import ToTensorV2
+
+# --- 1. SETUP ---
+app = FastAPI(title="AI Skin Lesion Analyzer API")
+
+DEVICE = "cpu"
+segmentation_model, classification_model, knowledge_base = None, None, None
+idx_to_class_abbr = {0: 'MEL', 1: 'NV', 2: 'BCC', 3: 'AKIEC', 4: 'BKL', 5: 'DF', 6: 'VASC'}
+transform_segment, transform_classify = None, None
+
+class ImageRequest(BaseModel):
+    image_base64: str
+
+@app.on_event("startup")
+def load_assets():
+    global segmentation_model, classification_model, knowledge_base, transform_segment, transform_classify
+    print("--> API starting up: Downloading models...")
+    
+    try:
+        seg_model_path = hf_hub_download("sheikh987/unet-isic2018", "unet_full_data_best_model.pth")
+        segmentation_model = smp.Unet("resnet34", encoder_weights=None, in_channels=3, classes=1).to(DEVICE)
+        segmentation_model.load_state_dict(torch.load(seg_model_path, map_location=DEVICE))
+        segmentation_model.eval()
+        print("    Segmentation model loaded.")
+    except Exception as e:
+        print(f"!!! FATAL: Could not load segmentation model: {e}")
+
+    try:
+        class_model_path = hf_hub_download("sheikh987/efficientnet-isic", "efficientnet_augmented_best.pth")
+        classification_model = timm.create_model('efficientnet_b3', pretrained=False, num_classes=7).to(DEVICE)
+        classification_model.load_state_dict(torch.load(class_model_path, map_location=DEVICE))
+        classification_model.eval()
+        print("    Classification model loaded.")
+    except Exception as e:
+        print(f"!!! FATAL: Could not load classification model: {e}")
+
+    with open('knowledge_base.json', 'r') as f:
+        knowledge_base = json.load(f)
+    print("    Knowledge base loaded.")
+        
+    transform_segment = A.Compose([A.Resize(256, 256), A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], max_pixel_value=255.0), ToTensorV2()])
+    transform_classify = transforms.Compose([transforms.Resize((300, 300)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
+    print("--> API ready.")
+
+def process_image(image_np):
+    # STAGE 1: SEGMENTATION
+    aug = transform_segment(image=image_np)
+    tensor = aug['image'].unsqueeze(0).to(DEVICE)
+    with torch.no_grad():
+        logits = segmentation_model(tensor)
+    mask = (torch.sigmoid(logits) > 0.5).float().squeeze().cpu().numpy()
+
+    if mask.sum() < 200: return None
+
+    # STAGE 2: CROP
+    rows, cols = np.any(mask, axis=1), np.any(mask, axis=0)
+    rmin, rmax = np.where(rows)[0][[0, -1]]
+    cmin, cmax = np.where(cols)[0][[0, -1]]
+    padding = 15
+    rmin, rmax = max(0, rmin - padding), min(image_np.shape[0], rmax + padding)
+    cmin, cmax = max(0, cmin - padding), min(image_np.shape[1], cmax + padding)
+    cropped = Image.fromarray(image_np[rmin:rmax, cmin:cmax])
+
+    # STAGE 3: CLASSIFY
+    tensor = transform_classify(cropped).unsqueeze(0).to(DEVICE)
+    with torch.no_grad():
+        logits = classification_model(tensor)
+        probs = torch.nn.functional.softmax(logits, dim=1)
+    
+    conf, idx = torch.max(probs, 1)
+    return idx.item(), conf.item()
+
+# --- 2. API ENDPOINTS ---
+@app.post("/analyze/")
+async def analyze_image(request: ImageRequest):
+    if not all([segmentation_model, classification_model]):
+        raise HTTPException(status_code=503, detail="Models are not ready.")
+    try:
+        img_data = base64.b64decode(request.image_base64)
+        img_np = np.array(Image.open(io.BytesIO(img_data)).convert("RGB"))
+    except:
+        raise HTTPException(status_code=400, detail="Invalid base64 image.")
+
+    analysis = process_image(img_np)
+    if analysis is None:
+        return {"status": "Failed", "message": "No lesion could be identified."}
+
+    pred_idx, confidence = analysis
+    if confidence < 0.75:
+         return {"status": "Inconclusive", "message": f"Model confidence ({confidence*100:.2f}%) is below the 75% threshold."}
+
+    pred_abbr = idx_to_class_abbr[pred_idx]
+    info = knowledge_base.get(pred_abbr, {})
+    return {"status": "Success", "prediction": info, "abbreviation": pred_abbr, "confidence": f"{confidence*100:.2f}%"}
+
+@app.get("/")
+def root():
+    return {"message": "AI Skin Lesion Analyzer API is running."}