Update app.py

app.py

@@ -39,12 +39,16 @@ 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 ===
     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)
+    # === Define the correct model architecture (B4) to prevent size mismatch ===
+    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))
     classification_model.eval()
     print("    Classification model loaded successfully.")
@@ -56,11 +60,14 @@ except Exception as e:
 
 # --- Load Knowledge Base and Labels ---
 try:
-    with open('knowledge_base.json', 'r') as f:
+    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)
-    print("--> Knowledge base loaded.")
-except FileNotFoundError:
-    raise gr.Error("knowledge_base.json not found. Make sure it has been uploaded to the Space.")
+    print("--> Knowledge base loaded from Hub.")
+except Exception as e:
+    print(f"!!! ERROR loading knowledge_base.json from Hub: {e}")
+    raise gr.Error("knowledge_base.json not found in the Hub repo. Make sure it has been uploaded.")
+
 
 idx_to_class_abbr = {0: 'MEL', 1: 'NV', 2: 'BCC', 3: 'AKIEC', 4: 'BKL', 5: 'DF', 6: 'VASC'}
 
@@ -70,8 +77,9 @@ 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) ===
 transform_classify = transforms.Compose([
-    transforms.Resize((300, 300)),
+    transforms.Resize((380, 380)),
     transforms.ToTensor(),
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
 ])
@@ -79,7 +87,7 @@ transform_classify = transforms.Compose([
 print("\n--> Application ready to accept requests.")
 
 
-# --- 2. DEFINE THE FULL PIPELINE FUNCTION (UPDATED) ---
+# --- 2. DEFINE THE FULL PIPELINE FUNCTION ---
 
 def full_pipeline(input_image):
     if input_image is None:
@@ -117,35 +125,29 @@ def full_pipeline(input_image):
     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."
+            f"The AI model's confidence ({confidence_percent:.2f}%) is below the required threshold of {CONFIDENCE_THRESHOLD}%."
         )
         mask_display = Image.fromarray((seg_mask * 255).astype(np.uint8))
         return mask_display, cropped_image_pil, inconclusive_text
 
-    # --- STAGE 3: LOOKUP and FORMAT (UPDATED) ---
+    # STAGE 3: LOOKUP and FORMAT
     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_list = info.get('causes', [])
     causes_text = "\n".join([f"• {c}" for c in causes_list])
 
-    treatments_list = info.get('common_treatments', ['No specific treatments listed.'])
+    treatments_list = info.get('common_treatments', [])
     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"**Description:**\n{info.get('description', 'N/A')}\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', '')}"
@@ -166,7 +168,7 @@ iface = gr.Interface(
         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.",
+    description="A two-stage AI tool for skin lesion analysis.",
     allow_flagging="never"
 )