Update app.py

app.py

@@ -6,11 +6,9 @@ from huggingface_hub import hf_hub_download
 from PIL import Image
 import torch.nn.functional as F
 import json
-import traceback
 
 # Model repositories
 BIOMEDCLIP_REPO = "AssanaliAidarkhan/Biomedclip"
-QWEN_RAG_REPO = "AssanaliAidarkhan/qwen-medical-rag"
 
 # Global variables
 biomedclip_model = None
@@ -18,7 +16,6 @@ biomedclip_processor = None
 biomedclip_id2label = {}
 qwen_model = None
 qwen_tokenizer = None
-medical_knowledge = []
 
 class CLIPClassifier(nn.Module):
     def __init__(self, clip_model, num_classes):
@@ -32,12 +29,10 @@ class CLIPClassifier(nn.Module):
         return {'logits': logits}
 
 def load_biomedclip():
-    """Load BiomedCLIP model"""
+    """Load BiomedCLIP (we know this works)"""
     global biomedclip_model, biomedclip_processor, biomedclip_id2label
     
     try:
-        print("🔄 Loading BiomedCLIP...")
-        
         model_path = hf_hub_download(repo_id=BIOMEDCLIP_REPO, filename="pytorch_model.bin")
         checkpoint = torch.load(model_path, map_location='cpu')
         
@@ -54,133 +49,42 @@ def load_biomedclip():
         
         print("✅ BiomedCLIP loaded!")
         return True
-        
     except Exception as e:
         print(f"❌ BiomedCLIP error: {e}")
         return False
 
-def load_qwen():
-    """Load Qwen model and medical knowledge"""
-    global qwen_model, qwen_tokenizer, medical_knowledge
+def load_qwen_simple():
+    """Load Qwen with minimal setup"""
+    global qwen_model, qwen_tokenizer
     
     try:
-        print("🔄 Loading Qwen...")
+        print("🔄 Loading Qwen (simple)...")
         
-        # Load medical knowledge
-        try:
-            knowledge_path = hf_hub_download(repo_id=QWEN_RAG_REPO, filename="medical_knowledge.json")
-            with open(knowledge_path, 'r', encoding='utf-8') as f:
-                medical_knowledge = json.load(f)
-            print(f"✅ Loaded {len(medical_knowledge)} medical docs")
-        except:
-            # Fallback knowledge base
-            medical_knowledge = [
-                {
-                    "category": "partial_acl_injury",
-                    "advice": "Partial ACL injury detected. Recommend: rest, ice therapy, physical therapy consultation, avoid pivoting activities. Follow-up MRI in 6-8 weeks to assess healing progress."
-                },
-                {
-                    "category": "complete_acl_tear", 
-                    "advice": "Complete ACL tear detected. Urgent orthopedic consultation required. Likely surgical reconstruction needed, especially for active patients. Immediate immobilization recommended."
-                },
-                {
-                    "category": "acl_sprain",
-                    "advice": "ACL sprain detected. Conservative management with RICE protocol. Physical therapy for strengthening. Return to activity when pain-free and strength restored."
-                },
-                {
-                    "category": "normal",
-                    "advice": "ACL appears normal on MRI. Continue regular activities. If symptoms persist, consider clinical examination for other possible causes."
-                }
-            ]
-            print("⚠️ Using fallback medical knowledge")
+        # Load Qwen directly
+        qwen_tokenizer = AutoTokenizer.from_pretrained(
+            "Qwen/Qwen1.5-0.5B-Chat", 
+            trust_remote_code=True
+        )
         
-        # Load Qwen
-        qwen_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen1.5-0.5B-Chat", trust_remote_code=True)
         qwen_model = AutoModelForCausalLM.from_pretrained(
             "Qwen/Qwen1.5-0.5B-Chat",
             torch_dtype=torch.float32,
             trust_remote_code=True
         )
-        qwen_model.eval()
         
         print("✅ Qwen loaded!")
         return True
         
     except Exception as e:
         print(f"❌ Qwen error: {e}")
-        print(traceback.format_exc())
         return False
 
-def find_medical_advice(classification_result):
-    """Find relevant medical advice based on classification"""
-    
-    # Extract the main condition from classification
-    condition = classification_result.lower()
-    
-    # Find matching advice
-    for doc in medical_knowledge:
-        if doc['category'].lower() in condition or any(tag in condition for tag in doc.get('tags', [])):
-            return doc.get('advice', 'No specific advice available for this condition.')
-    
-    # Generic advice if no match
-    return "Consult with a medical professional for proper evaluation and treatment recommendations."
-
-def generate_qwen_advice(classification_result, medical_context):
-    """Generate advice using Qwen"""
-    global qwen_model, qwen_tokenizer
-    
-    if qwen_model is None:
-        return "❌ Qwen model not available"
-    
-    try:
-        # Create medical prompt
-        prompt = f"""Medical Image Analysis Result: {classification_result}
-
-Relevant Medical Knowledge: {medical_context}
-
-Based on this MRI classification, provide clinical recommendations including:
-1. Immediate actions needed
-2. Treatment options
-3. Follow-up requirements
-4. Patient advice
-
-Response:"""
-        
-        # Tokenize and generate
-        inputs = qwen_tokenizer(prompt, return_tensors="pt", max_length=500, truncation=True)
-        
-        with torch.no_grad():
-            outputs = qwen_model.generate(
-                inputs.input_ids,
-                max_new_tokens=150,
-                temperature=0.7,
-                do_sample=True,
-                pad_token_id=qwen_tokenizer.eos_token_id
-            )
-        
-        # Decode response
-        generated_ids = outputs[0][len(inputs.input_ids[0]):]
-        advice = qwen_tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
-        
-        return advice
-        
-    except Exception as e:
-        return f"❌ Qwen generation error: {str(e)}"
-
-def complete_analysis(image):
-    """Complete pipeline: Classification + Medical Advice"""
-    
-    if image is None:
-        return "❌ Please upload an MRI scan", ""
-    
-    print("🏥 Starting complete analysis...")
+def classify_mri(image):
+    """Classify MRI (working code)"""
+    if biomedclip_model is None or image is None:
+        return None
     
-    # Step 1: Classify with BiomedCLIP
     try:
-        # Classification code (same as your working debug version)
-        if biomedclip_model is None:
-            return "❌ BiomedCLIP not loaded", ""
-        
         if image.mode != 'RGB':
             image = image.convert('RGB')
         
@@ -202,28 +106,89 @@ def complete_analysis(image):
             class_name = f"Class_{class_idx}"
         
         confidence = top_prob.item() * 100
-        classification_result = f"{class_name} (confidence: {confidence:.1f}%)"
         
-        print(f"✅ Classification: {classification_result}")
+        return class_name, confidence
         
     except Exception as e:
-        return f"❌ Classification error: {e}", ""
+        print(f"Classification error: {e}")
+        return None, None
+
+def generate_simple_advice(class_name, confidence):
+    """Generate advice using Qwen (simple approach)"""
+    global qwen_model, qwen_tokenizer
+    
+    if qwen_model is None:
+        return "❌ Qwen model not loaded"
     
-    # Step 2: Get medical advice
     try:
-        # Find relevant medical context
-        medical_context = find_medical_advice(class_name)
+        print(f"🔄 Generating advice for: {class_name}")
         
-        # Generate advice with Qwen
-        qwen_advice = generate_qwen_advice(classification_result, medical_context)
+        # Simple medical knowledge lookup
+        advice_map = {
+            "partial_acl_injury": "Partial ACL injury detected. Recommendations: Rest and avoid pivoting activities. Apply ice for 15-20 minutes several times daily. Consider physical therapy consultation. Follow-up MRI in 6-8 weeks to monitor healing.",
+            "complete_acl_tear": "Complete ACL tear detected. Urgent orthopedic consultation required. Likely surgical reconstruction needed. Immediate immobilization and avoid weight-bearing activities.",
+            "acl_sprain": "ACL sprain detected. Conservative treatment with RICE protocol (Rest, Ice, Compression, Elevation). Physical therapy for strengthening. Gradual return to activities.",
+            "normal": "ACL appears normal. Continue regular activities. If symptoms persist, consider clinical examination for other causes."
+        }
         
-        print("✅ Medical advice generated")
+        # Get base advice
+        base_advice = advice_map.get(class_name.lower(), "Consult medical professional for evaluation.")
+        
+        # Create simple prompt for Qwen
+        simple_prompt = f"Medical diagnosis: {class_name} with {confidence:.1f}% confidence. Provide brief clinical advice:"
+        
+        # Tokenize
+        inputs = qwen_tokenizer(simple_prompt, return_tensors="pt")
+        
+        # Generate
+        with torch.no_grad():
+            outputs = qwen_model.generate(
+                inputs.input_ids,
+                max_new_tokens=100,
+                temperature=0.8,
+                do_sample=True,
+                pad_token_id=qwen_tokenizer.eos_token_id
+            )
+        
+        # Decode
+        full_output = qwen_tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Extract just the generated part
+        if simple_prompt in full_output:
+            generated_advice = full_output.replace(simple_prompt, "").strip()
+        else:
+            generated_advice = full_output
+        
+        # Combine base advice with Qwen advice
+        if generated_advice and len(generated_advice) > 10:
+            combined_advice = f"**Clinical Guidelines:** {base_advice}\n\n**AI Analysis:** {generated_advice}"
+        else:
+            combined_advice = base_advice
+        
+        print(f"✅ Generated advice: {generated_advice[:50]}...")
+        return combined_advice
         
     except Exception as e:
-        qwen_advice = f"❌ Advice generation error: {e}"
-        print(f"❌ Advice error: {e}")
+        print(f"❌ Advice generation error: {e}")
+        # Fallback to basic advice
+        return advice_map.get(class_name.lower(), "Consult medical professional for evaluation.")
+
+def complete_pipeline(image):
+    """Complete analysis pipeline"""
+    
+    if image is None:
+        return "❌ Please upload an MRI scan", ""
+    
+    # Step 1: Classification
+    class_name, confidence = classify_mri(image)
     
-    # Format results
+    if class_name is None:
+        return "❌ Classification failed", ""
+    
+    # Step 2: Medical advice
+    medical_advice = generate_simple_advice(class_name, confidence)
+    
+    # Format outputs
     classification_text = f"""
 # 🔬 **MRI Classification**
 
@@ -232,66 +197,46 @@ def complete_analysis(image):
 
 ## 📊 **Confidence:**
 **{confidence:.1f}%**
-
-## 📈 **Details:**
-- Class index: {class_idx}
-- Total classes: {len(biomedclip_id2label)}
-- Model: BiomedCLIP
     """
     
     advice_text = f"""
-# 🏥 **Clinical Recommendations**
-
-## 💡 **AI-Generated Advice:**
+# 🏥 **Medical Recommendations**
 
-{qwen_advice}
-
-## 📚 **Medical Knowledge:**
-
-{medical_context}
+{medical_advice}
 
 ---
-⚠️ **Disclaimer:** For educational purposes only. Always consult medical professionals.
+⚠️ **Disclaimer:** For educational purposes only. Consult medical professionals.
     """
     
     return classification_text, advice_text
 
 # Load models
-print("🚀 Initializing models...")
 biomedclip_loaded = load_biomedclip()
-qwen_loaded = load_qwen()
+qwen_loaded = load_qwen_simple()
 
 # Create interface
 with gr.Blocks(title="Medical AI Pipeline") as app:
     
-    gr.Markdown("# 🏥 Complete Medical AI Analysis Pipeline")
-    gr.Markdown("**BiomedCLIP** (Image Classification) + **Qwen** (Medical Advice)")
+    gr.Markdown("# 🏥 Medical AI Analysis Pipeline")
+    gr.Markdown("**BiomedCLIP** (Classification) + **Qwen** (Medical Advice)")
     
-    # Status
-    status_text = f"BiomedCLIP: {'✅' if biomedclip_loaded else '❌'} | Qwen: {'✅' if qwen_loaded else '❌'}"
-    gr.Markdown(f"**Status:** {status_text}")
+    status = f"Status: BiomedCLIP {'✅' if biomedclip_loaded else '❌'} | Qwen {'✅' if qwen_loaded else '❌'}"
+    gr.Markdown(f"**{status}**")
     
     with gr.Row():
         with gr.Column():
-            image_input = gr.Image(type="pil", label="📸 Upload MRI Scan", height=400)
-            analyze_btn = gr.Button("🔬 Complete Analysis", variant="primary", size="lg")
-            clear_btn = gr.Button("🗑️ Clear")
+            image_input = gr.Image(type="pil", label="📸 Upload MRI Scan")
+            analyze_btn = gr.Button("🔬 Complete Analysis", variant="primary")
         
         with gr.Column():
             classification_output = gr.Markdown(label="🔬 Classification")
-            advice_output = gr.Markdown(label="🏥 Medical Advice")
+            advice_output = gr.Markdown(label="🏥 Medical Advice") 
     
-    # Button actions
     analyze_btn.click(
-        fn=complete_analysis,
+        fn=complete_pipeline,
         inputs=image_input,
         outputs=[classification_output, advice_output]
     )
-    
-    clear_btn.click(
-        fn=lambda: [None, "", ""],
-        outputs=[image_input, classification_output, advice_output]
-    )
 
 if __name__ == "__main__":
     app.launch()