Update app.py

app.py

@@ -1,42 +1,100 @@
+import gradio as gr
 import torch
-import timm
+from PIL import Image
+import numpy as np
+import json
 from huggingface_hub import hf_hub_download
-import gradio as gr
+import timm
+from torchvision import transforms
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Repo iyo filename sax ah
+# --- Download and Load Classification Model (EfficientNet-B3) ---
 CLASS_REPO_ID = "sheikh987/efficientnet-b3-skin"
 CLASS_MODEL_FILENAME = "efficientnet_b3_skin_model.pth"
+NUM_CLASSES = 7
 
-# Download model-ka Hugging Face hub
 try:
     class_model_path = hf_hub_download(
         repo_id=CLASS_REPO_ID,
         filename=CLASS_MODEL_FILENAME,
-        cache_dir="/tmp"  # important in Spaces
+        cache_dir="/tmp"
     )
-    print(f"✅ Model downloaded to: {class_model_path}")
-except Exception as e:
-    raise gr.Error(f"Failed to download model: {e}")
-
-# Auto-detect state_dict vs full model
-try:
     checkpoint = torch.load(class_model_path, map_location=DEVICE)
-    
+
+    # Auto-detect state_dict vs full model
     if isinstance(checkpoint, dict):
-        # state_dict case
         classification_model = timm.create_model(
-            'efficientnet_b3', pretrained=False, num_classes=7
+            'efficientnet_b3', pretrained=False, num_classes=NUM_CLASSES
         ).to(DEVICE)
-        classification_model.load_state_dict(checkpoint)
-        print("✅ Loaded as state_dict")
+        classification_model.load_state_dict(checkpoint, strict=False)
     else:
-        # full model case
         classification_model = checkpoint.to(DEVICE)
-        print("✅ Loaded as full model")
 
     classification_model.eval()
-    print("✅ Classification model is ready.")
+    print("✅ Classification model loaded successfully.")
+
 except Exception as e:
-    raise gr.Error(f"Failed to load classification model: {e}")
+    raise gr.Error(f"Failed to load the classification model: {e}")
+
+
+# --- Load Knowledge Base ---
+try:
+    with open('knowledge_base.json', 'r') as f:
+        knowledge_base = json.load(f)
+except FileNotFoundError:
+    raise gr.Error("knowledge_base.json not found. Upload it to the Space.")
+
+idx_to_class_abbr = {0: 'MEL', 1: 'NV', 2: 'BCC', 3: 'AKIEC', 4: 'BKL', 5: 'DF', 6: 'VASC'}
+
+# --- Image Transform ---
+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]),
+])
+
+
+# --- Pipeline Function ---
+def classify_image(input_image):
+    if input_image is None:
+        return None, "Please upload an image."
+
+    class_input_tensor = transform_classify(input_image).unsqueeze(0).to(DEVICE)
+
+    with torch.no_grad():
+        logits = classification_model(class_input_tensor)
+        probs = torch.nn.functional.softmax(logits, dim=1)
+
+    confidence, predicted_idx = torch.max(probs, 1)
+    confidence_percent = confidence.item() * 100
+    predicted_abbr = idx_to_class_abbr[predicted_idx.item()]
+    info = knowledge_base.get(predicted_abbr, {})
+
+    # Build output
+    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" + "\n".join([f"• {c}" for c in info.get('causes', ['N/A'])]) + "\n\n"
+        f"**Common Treatments:**\n" + "\n".join([f"• {t}" for t in info.get('common_treatments', ['N/A'])]) + "\n\n"
+        f"**--- IMPORTANT DISCLAIMER ---**\n{info.get('disclaimer', '')}"
+    )
+
+    return input_image, info_text
+
+
+# --- Gradio Interface ---
+iface = gr.Interface(
+    fn=classify_image,
+    inputs=gr.Image(type="pil", label="Upload Skin Image"),
+    outputs=[gr.Image(type="pil", label="Input Image"),
+             gr.Markdown(label="Analysis Result")],
+    title="AI Skin Lesion Classifier",
+    description="Upload a skin lesion image and the AI EfficientNet-B3 model will classify it.\n\n"
+                "**DISCLAIMER:** This is NOT a diagnosis. Always consult a qualified dermatologist."
+)
+
+if __name__ == "__main__":
+    iface.launch()