Update app.py

app.py

@@ -39,14 +39,14 @@ except Exception as e:
 # --- Download and Load Classification Model (EfficientNet) ---
 try:
     CLASS_REPO_ID = "sheikh987/efficientnet-B4"
-    # === Use the filename you want, which now exists in the repo ===
+    # This filename MUST exactly match the file in your Hugging Face repo
     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
-    # === Define the correct model architecture (B4) to prevent size mismatch ===
+    # The architecture ('efficientnet_b4') MUST match the model weights
     classification_model = timm.create_model('efficientnet_b4', pretrained=False, num_classes=NUM_CLASSES).to(DEVICE)
     
     classification_model.load_state_dict(torch.load(class_model_path, map_location=DEVICE))
@@ -60,6 +60,7 @@ except Exception as e:
 
 # --- Load Knowledge Base and Labels ---
 try:
+    # It's best practice to also host the knowledge base in your repo
     knowledge_base_path = hf_hub_download(repo_id=CLASS_REPO_ID, filename="knowledge_base.json")
     with open(knowledge_base_path, 'r') as f:
         knowledge_base = json.load(f)
@@ -77,7 +78,7 @@ transform_segment = A.Compose([
     A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225], max_pixel_value=255.0),
     ToTensorV2(),
 ])
-# === Use the correct image size for EfficientNet-B4 (380x380) ===
+# The image size (380x380) MUST match the model architecture
 transform_classify = transforms.Compose([
     transforms.Resize((380, 380)),
     transforms.ToTensor(),
@@ -138,16 +139,16 @@ def full_pipeline(input_image):
     predicted_abbr = idx_to_class_abbr[predicted_idx.item()]
     info = knowledge_base.get(predicted_abbr, {})
     
-    causes_list = info.get('causes', [])
+    causes_list = info.get('causes', ['N/A'])
     causes_text = "\n".join([f"• {c}" for c in causes_list])
 
-    treatments_list = info.get('common_treatments', [])
+    treatments_list = info.get('common_treatments', ['N/A'])
     treatments_text = "\n".join([f"• {t}" for t in treatments_list])
 
     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', 'N/A')}\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', '')}"
@@ -168,7 +169,7 @@ iface = gr.Interface(
         gr.Markdown(label="Analysis Result")
     ],
     title="AI Skin Lesion Analyzer",
-    description="A two-stage AI tool for skin lesion analysis.",
+    description="A two-stage AI tool for skin lesion analysis. **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"
 )