unifying the process

app.py

@@ -4,9 +4,25 @@ import tensorflow as tf
 from PIL import Image
 import efficientnet.tfkeras as efn 
 import random 
+import torch
+from open_clip import create_model_and_transforms, get_tokenizer
 
 # ==========================================
-# 1. MRI Model Setup (Your Existing Model)
+# 1. Modality Router Setup (BiomedCLIP)
+# ==========================================
+print("Loading BiomedCLIP Router...")
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+clip_model_name = 'hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224'
+clip_model, _, clip_preprocess = create_model_and_transforms(clip_model_name)
+clip_model = clip_model.to(device)
+clip_tokenizer = get_tokenizer(clip_model_name)
+
+# Define the text embeddings for routing
+router_labels = ['an MRI brain scan', 'a chest X-ray']
+text_tokens = clip_tokenizer(router_labels).to(device)
+
+# ==========================================
+# 2. MRI Model Setup 
 # ==========================================
 print("Loading MRI model...")
 mri_model = tf.keras.models.load_model("mri.keras")
@@ -16,7 +32,6 @@ def predict_mri(image):
     if image is None:
         return None
     
-    # Preprocess the MRI
     img = Image.fromarray(image).convert('L') 
     img = img.resize((168, 168))
     
@@ -24,28 +39,21 @@ def predict_mri(image):
     img_array = np.expand_dims(img_array, axis=-1)
     img_array = np.expand_dims(img_array, axis=0)
     
-    # Predict
     predictions = mri_model.predict(img_array)[0]
     
-    # Apply the 3% to 7% random reduction
     confidences = {}
     for i in range(len(mri_class_names)):
         original_conf = float(predictions[i])
         random_drop = random.uniform(0.03, 0.07) 
-        
-        # Ensure it doesn't drop below 0
         adjusted_conf = max(0.0, original_conf - random_drop)
-        
-        # Rounding to 4 decimal places
         confidences[mri_class_names[i]] = round(adjusted_conf, 4)
         
     return confidences
 
 # ==========================================
-# 2. X-Ray Model Setup (Using original EfficientNet library)
+# 3. X-Ray Model Setup 
 # ==========================================
 print("Building X-Ray model architecture...")
-
 xray_class_names = [
     'Cardiomegaly', 'Emphysema', 'Effusion', 'Hernia', 'Infiltration', 
     'Mass', 'Nodule', 'Atelectasis', 'Pneumothorax', 'Pleural_Thickening', 
@@ -76,7 +84,6 @@ def predict_xray(image):
     if image is None:
         return None
     
-    # Preprocess the X-Ray input
     img = Image.fromarray(image).convert('RGB')
     img = img.resize((128, 128))
     
@@ -85,54 +92,76 @@ def predict_xray(image):
     
     img_array = efn.preprocess_input(img_array)
     
-    # Predict
     predictions = xray_model.predict(img_array)[0]
     
-    # Apply the 3% to 7% random reduction
     confidences = {}
     for i in range(len(xray_class_names)):
         original_conf = float(predictions[i])
         random_drop = random.uniform(0.03, 0.07) 
-        
-        # Ensure it doesn't drop below 0
         adjusted_conf = max(0.0, original_conf - random_drop)
-        
-        # Rounding to 4 decimal places
         confidences[xray_class_names[i]] = round(adjusted_conf, 4)
         
     return confidences
 
 # ==========================================
-# 3. Define the Gradio Interface with Tabs
+# 4. Master Routing Function
+# ==========================================
+def process_scan(image):
+    if image is None:
+        return "No image provided.", None
+
+    # Step A: Preprocess for CLIP
+    img_pil = Image.fromarray(image).convert('RGB')
+    img_tensor = clip_preprocess(img_pil).unsqueeze(0).to(device)
+
+    # Step B: Calculate Modality Probabilities
+    with torch.no_grad():
+        image_features = clip_model.encode_image(img_tensor)
+        text_features = clip_model.encode_text(text_tokens)
+        
+        image_features /= image_features.norm(dim=-1, keepdim=True)
+        text_features /= text_features.norm(dim=-1, keepdim=True)
+        
+        text_probs = (100.0 * image_features @ text_features.T).softmax(dim=-1)[0]
+
+    mri_prob = text_probs[0].item()
+    xray_prob = text_probs[1].item()
+
+    # Step C: Route to Specific Model
+    if mri_prob > xray_prob:
+        modality_status = f"🧠 Modality Detected: MRI Brain Scan (Confidence: {mri_prob:.1%})"
+        diagnostic_results = predict_mri(image)
+        # We only want top 1 for MRI based on your previous UI setup
+        top_k = 1 
+    else:
+        modality_status = f"🩻 Modality Detected: Chest X-Ray (Confidence: {xray_prob:.1%})"
+        diagnostic_results = predict_xray(image)
+        # We want top 2 for X-Ray based on your previous UI setup
+        top_k = 2
+
+    return modality_status, diagnostic_results
+
+# ==========================================
+# 5. Define the Unified Gradio Interface
 # ==========================================
 with gr.Blocks(title="Medical Scan Classification") as interface:
     gr.Markdown("# 🩺 Medical Scan Classifier")
-    gr.Markdown("Upload an **MRI Brain Scan** or a **Chest X-Ray** into the respective tabs below for AI-powered classification.")
+    gr.Markdown("Upload **any** scan (MRI Brain Scan or Chest X-Ray). The system will automatically detect the modality and route it to the appropriate diagnostic model.")
     
-    with gr.Tabs():
-        # --- TAB 1: MRI ---
-        with gr.TabItem("MRI Brain Scan"):
-            with gr.Row():
-                with gr.Column():
-                    mri_input = gr.Image(label="Upload MRI Brain Scan")
-                    mri_button = gr.Button("Classify MRI", variant="primary")
-                with gr.Column():
-                    # CHANGE APPLIED HERE: num_top_classes changed to 1
-                    mri_output = gr.Label(num_top_classes=1, label="Top Predicted Condition")
-            
-            mri_button.click(fn=predict_mri, inputs=mri_input, outputs=mri_output)
+    with gr.Row():
+        with gr.Column():
+            scan_input = gr.Image(label="Upload Medical Scan")
+            analyze_button = gr.Button("Analyze Scan", variant="primary")
+        
+        with gr.Column():
+            modality_output = gr.Textbox(label="Detection Routing Status", interactive=False)
+            diagnostic_output = gr.Label(label="Predicted Conditions")
             
-        # --- TAB 2: X-Ray ---
-        with gr.TabItem("Chest X-Ray"):
-            with gr.Row():
-                with gr.Column():
-                    xray_input = gr.Image(label="Upload Chest X-Ray")
-                    xray_button = gr.Button("Classify X-Ray", variant="primary")
-                with gr.Column():
-                    xray_output = gr.Label(num_top_classes=2, label="Top 2 Predicted Conditions")
-                    
-            xray_button.click(fn=predict_xray, inputs=xray_input, outputs=xray_output)
-
-# Launch the app
+    analyze_button.click(
+        fn=process_scan, 
+        inputs=scan_input, 
+        outputs=[modality_output, diagnostic_output]
+    )
+
 if __name__ == "__main__":
     interface.launch()