targeting 80% recall

app.py

@@ -9,7 +9,7 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # 2. Load Fine-Tuned Model and Processor
 model_path = "actorcritic/twak"  # <-- Change this back to just the repo name
-#model_path = "/home/shohog/Documents/twok"
+model_path = "/home/shohog/Documents/twok"
 
 try:
     # Add subfolder="model" here
@@ -48,8 +48,28 @@ disease_info = {
     "Vitiligo": "শ্বেতী (Vitiligo) হলো ত্বকের এমন একটি অবস্থা যেখানে ত্বক থেকে মেলানিন রঞ্জক পদার্থ নষ্ট হয়ে সাদা ছোপের সৃষ্টি হয়।"
 }
 
-# Confidence threshold
-CONFIDENCE_THRESHOLD = 80.0  # percent
+
+# 15 Original Folders (Index 6 is the 'Others/Healthy' folder)
+ORIGINAL_CLASSES =[
+    "Acne", "Arsenic", "Atopic_Dermatitis", "Candidal_Intertrigo", 
+    "Contact_Dermatitis", "Eczema", "Healthy_or_Others", "Psoriasis", 
+    "Scabies", "Seborrheic_Dermatitis", "Steroid_Modified_Tinea", 
+    "Tinea_Corporis", "Tinea_Cruris", "Tinea_Faciei", "Vitiligo"
+]
+
+# Optimal Thresholds per Output Neuron (Targeting 80% Recall)
+CLASS_THRESHOLDS = {
+    0: 0.89, 1: 0.88, 2: 0.34, 3: 0.04, 4: 0.66, 5: 0.48, 
+    6: 0.28, 7: 0.73, 8: 0.54, 9: 0.44, 10: 0.68, 11: 0.83, 
+    12: 0.85, 13: 0.95
+}
+
+def get_original_class_idx(neuron_idx):
+    """Maps the 14 output neurons back to the original 15 folder indices."""
+    if neuron_idx < 6:
+        return neuron_idx
+    else:
+        return neuron_idx + 1  # Shift back to account for the skipped Class 6
 
 # 4. Prediction Function
 def predict(image):
@@ -68,50 +88,66 @@ def predict(image):
 
     logits = outputs.logits
     scores = torch.sigmoid(logits)[0]
+    print(logits, scores)
+
+    detected_diseases =[]
 
-    # Get top prediction
-    top_idx = torch.argmax(scores).item()
-    confidence = scores[top_idx].item() * 100
-
-    # Check confidence threshold
-    if confidence < CONFIDENCE_THRESHOLD:
-        # Show "Unknown" with the closest guess in tooltip
-        guessed_class = ft_model.config.id2label.get(top_idx, f"Class {top_idx}")
-        html_output = f"""
-        <div class="result-container">
-            <span class="disease-name unknown">UNK</span>
+    # Check each of the 14 neurons against its specific threshold
+    for neuron_idx in range(14):
+        prob = scores[neuron_idx].item()
+        if prob >= CLASS_THRESHOLDS[neuron_idx]:
+            orig_idx = get_original_class_idx(neuron_idx)
+            detected_diseases.append((orig_idx, prob))
+
+    html_output = '<div class="result-container">'
+
+    # SCENARIO A: No thresholds were crossed -> It's Class 6 (Healthy/Others)
+    if len(detected_diseases) == 0:
+        predicted_class = "Healthy / Others"
+        info_text = "কোনো চর্মরোগ শনাক্ত হয়নি। ত্বক সুস্থ অথবা এটি অন্য কোনো সাধারণ অবস্থা হতে পারে।"
+        
+        # Added flex-direction: column here to stack them vertically
+        html_output += f"""
+        <div style="display: flex; flex-direction: column; align-items: center; justify-content: center; gap: 8px;">
+            <span class="disease-name" style="color: #10b981;">{predicted_class}</span>
             <div class="tooltip">
-                <span class="info-icon">?</span>
+                <span class="info-icon">i</span>
                 <div class="tooltiptext">
-                    <strong>Low Confidence Prediction</strong><br><br>
-                    The model is not confident enough for a reliable diagnosis.<br><br>
-                    Closest match: <b>{guessed_class}</b> ({confidence:.1f}%)<br><br>
-                    <i>Please consult a healthcare professional for proper diagnosis.</i>
+                    <strong>{predicted_class}</strong><br><br>
+                    {info_text}
                 </div>
             </div>
         </div>
         """
-        return html_output
-
-    predicted_class = ft_model.config.id2label.get(top_idx, f"Class {top_idx}")
-    info_text = disease_info.get(predicted_class, "Fundamental details for this condition are not available.")
-
-    # Build HTML output with tooltip for confident predictions
-    html_output = f"""
-    <div class="result-container">
-        <span class="disease-name">{predicted_class}</span>
-        <div class="tooltip">
-            <span class="info-icon">i</span>
-            <div class="tooltiptext">
-                <strong>{predicted_class}</strong><br><br>
-                {info_text}
+        
+    # SCENARIO B: One or more diseases detected
+    else:
+        # Sort by probability (highest confidence first)
+        detected_diseases.sort(key=lambda x: x[1], reverse=True)
+        
+        for orig_idx, prob in detected_diseases:
+            predicted_class = ORIGINAL_CLASSES[orig_idx]
+            info_text = disease_info.get(predicted_class, "Fundamental details for this condition are not available.")
+            confidence = prob * 100
+            print(confidence)
+            
+            # Added flex-direction: column here to stack them vertically
+            html_output += f"""
+            <div style="display: flex; flex-direction: column; align-items: center; justify-content: center; gap: 8px; margin-bottom: 15px; width: 100%;">
+                <span class="disease-name">{predicted_class} <span style="font-size: 16px; color: #6b7280;">({confidence:.1f}%)</span></span>
+                <div class="tooltip">
+                    <span class="info-icon">i</span>
+                    <div class="tooltiptext">
+                        <strong>{predicted_class}</strong><br><br>
+                        {info_text}
+                    </div>
+                </div>
             </div>
-        </div>
-    </div>
-    """
+            """
+            
+    html_output += '</div>'
     return html_output
 
-
 # 5. Custom CSS
 custom_css = """
 .result-container {
@@ -263,4 +299,4 @@ with gr.Blocks() as demo:
 
     predict_btn.click(fn=predict, inputs=image_input, outputs=output_html)
 
-demo.launch(css=custom_css, theme=gr.themes.Default())
+demo.launch(css=custom_css, theme=gr.themes.Default())