Update app.py

app.py

@@ -0,0 +1,174 @@
+# app.py
+
+import gradio as gr
+import torch
+from PIL import Image
+import numpy as np
+import json
+from huggingface_hub import hf_hub_download
+
+# Import necessary model libraries
+import segmentation_models_pytorch as smp
+import timm
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+from torchvision import transforms
+
+# --- 1. SETUP: Download and Load all models and data ---
+
+print("--> Initializing application and downloading models...")
+DEVICE = "cpu"
+
+# --- Download and Load Segmentation Model (UNet) ---
+try:
+    SEG_REPO_ID = "sheikh987/unet-isic2018"
+    SEG_MODEL_FILENAME = "unet_full_data_best_model.pth"
+    print(f"--> Downloading segmentation model from: {SEG_REPO_ID}")
+    seg_model_path = hf_hub_download(repo_id=SEG_REPO_ID, filename=SEG_MODEL_FILENAME)
+    
+    segmentation_model = smp.Unet(encoder_name="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 successfully.")
+
+except Exception as e:
+    print(f"!!! ERROR loading segmentation model: {e}")
+    raise gr.Error("Failed to load the segmentation model. Check repository name and file paths.")
+
+
+# --- Download and Load Classification Model (EfficientNet) ---
+try:
+    CLASS_REPO_ID = "sheikh987/efficientnet-B4"
+    CLASS_MODEL_FILENAME = "efficientnet_b4_augmented_best.pth"
+    print(f"--> Downloading classification model from: {CLASS_REPO_ID}")
+    class_model_path = hf_hub_download(repo_id=CLASS_REPO_ID, filename=CLASS_MODEL_FILENAME)
+
+    NUM_CLASSES = 7
+    classification_model = timm.create_model('efficientnet_b3', pretrained=False, num_classes=NUM_CLASSES).to(DEVICE)
+    classification_model.load_state_dict(torch.load(class_model_path, map_location=DEVICE))
+    classification_model.eval()
+    print("    Classification model loaded successfully.")
+
+except Exception as e:
+    print(f"!!! ERROR loading classification model: {e}")
+    raise gr.Error("Failed to load the classification model. Check repository name and file paths.")
+
+
+# --- Load Knowledge Base and Labels ---
+try:
+    with open('knowledge_base.json', 'r') as f:
+        knowledge_base = json.load(f)
+    print("--> Knowledge base loaded.")
+except FileNotFoundError:
+    raise gr.Error("knowledge_base.json not found. Make sure it has been uploaded to the Space.")
+
+idx_to_class_abbr = {0: 'MEL', 1: 'NV', 2: 'BCC', 3: 'AKIEC', 4: 'BKL', 5: 'DF', 6: 'VASC'}
+
+# --- Define Image Transforms ---
+transform_segment = A.Compose([
+    A.Resize(height=256, width=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("\n--> Application ready to accept requests.")
+
+
+# --- 2. DEFINE THE FULL PIPELINE FUNCTION (UPDATED) ---
+
+def full_pipeline(input_image):
+    if input_image is None:
+        return None, None, "Please upload an image."
+
+    image_np = np.array(input_image.convert("RGB"))
+
+    # STAGE 1: SEGMENTATION
+    augmented_seg = transform_segment(image=image_np)
+    seg_input_tensor = augmented_seg['image'].unsqueeze(0).to(DEVICE)
+    with torch.no_grad():
+        seg_logits = segmentation_model(seg_input_tensor)
+    seg_mask = (torch.sigmoid(seg_logits) > 0.5).float().squeeze().cpu().numpy()
+    
+    if seg_mask.sum() < 200:
+        return None, None, "Analysis Failed: No lesion could be clearly identified."
+
+    # STAGE 2: CROP and CLASSIFY
+    rows = np.any(seg_mask, axis=1)
+    cols = np.any(seg_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_pil = Image.fromarray(image_np[rmin:rmax, cmin:cmax])
+
+    class_input_tensor = transform_classify(cropped_image_pil).unsqueeze(0).to(DEVICE)
+    with torch.no_grad():
+        class_logits = classification_model(class_input_tensor)
+        probabilities = torch.nn.functional.softmax(class_logits, dim=1)
+    
+    confidence, predicted_idx = torch.max(probabilities, 1)
+    confidence_percent = confidence.item() * 100
+    
+    # SAFETY NET
+    CONFIDENCE_THRESHOLD = 50.0
+    if confidence_percent < CONFIDENCE_THRESHOLD:
+        inconclusive_text = (
+            f"**Analysis Inconclusive**\n\n"
+            f"The AI model's confidence ({confidence_percent:.2f}%) is below the required threshold of {CONFIDENCE_THRESHOLD}%.\n\n"
+            "This can happen if the image is blurry, has poor lighting, or shows a condition the model was not trained on.\n\n"
+            "**--- IMPORTANT DISCLAIMER ---**\n"
+            "This is NOT a diagnosis. Please consult a qualified dermatologist for an accurate assessment."
+        )
+        mask_display = Image.fromarray((seg_mask * 255).astype(np.uint8))
+        return mask_display, cropped_image_pil, inconclusive_text
+
+    # --- STAGE 3: LOOKUP and FORMAT (UPDATED) ---
+    predicted_abbr = idx_to_class_abbr[predicted_idx.item()]
+    info = knowledge_base.get(predicted_abbr, {})
+    
+    # Format the 'causes' and 'treatments' lists into clean, bulleted strings
+    causes_list = info.get('causes', ['Specific causes not listed.'])
+    causes_text = "\n".join([f"• {c}" for c in causes_list])
+
+    treatments_list = info.get('common_treatments', ['No specific treatments listed.'])
+    treatments_text = "\n".join([f"• {t}" for t in treatments_list])
+
+    # Build the final output text using all the information
+    info_text = (
+        f"**Predicted Condition:** {info.get('full_name', 'N/A')} ({predicted_abbr})\n"
+        f"**Confidence:** {confidence_percent:.2f}%\n\n"
+        f"**Description:**\n{info.get('description', 'No description available.')}\n\n"
+        f"**Common Causes:**\n{causes_text}\n\n"
+        f"**Common Treatments:**\n{treatments_text}\n\n"
+        f"**--- IMPORTANT DISCLAIMER ---**\n{info.get('disclaimer', '')}"
+    )
+    
+    mask_display = Image.fromarray((seg_mask * 255).astype(np.uint8))
+    
+    return mask_display, cropped_image_pil, info_text
+
+
+# --- 3. CREATE AND LAUNCH THE GRADIO INTERFACE ---
+iface = gr.Interface(
+    fn=full_pipeline,
+    inputs=gr.Image(type="pil", label="Upload Skin Image"),
+    outputs=[
+        gr.Image(type="pil", label="Segmentation Mask"),
+        gr.Image(type="pil", label="Cropped Lesion"),
+        gr.Markdown(label="Analysis Result")
+    ],
+    title="AI Skin Lesion Analyzer",
+    description="This tool performs a two-stage analysis on a skin lesion image. **Stage 1:** A UNet model segments the lesion. **Stage 2:** An EfficientNet model classifies the segmented lesion. \n\n**DISCLAIMER:** This is an educational tool and is NOT a substitute for professional medical advice. Always consult a qualified dermatologist for any health concerns.",
+    allow_flagging="never"
+)
+
+if __name__ == "__main__":
+    iface.launch()