Update app.py

app.py

@@ -43,6 +43,10 @@ class MedicalImageAnalysisSystem:
             model_kwargs={"low_cpu_mem_usage": True}
         )
         
+        # Set models to evaluation mode
+        self.tumor_classifier_model.eval()
+        self.radius_model.eval()
+        
         print("System initialization complete!")
 
     def generate_synthetic_metadata(self):
@@ -64,46 +68,51 @@ class MedicalImageAnalysisSystem:
             image = image.convert('RGB')
         return image.resize((224, 224))
 
-    @torch.no_grad()  # Disable gradient computation for inference
+    @torch.no_grad()
     def predict_tumor_presence(self, processed_image):
         inputs = self.tumor_classifier_processor(processed_image, return_tensors="pt")
-        inputs = {k: v.to(self.device) for k, v in inputs.items()}  # Move inputs to correct device
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
         outputs = self.tumor_classifier_model(**inputs)
         predictions = torch.softmax(outputs.logits, dim=-1)
-        probs = predictions[0].cpu().tolist()  # Move back to CPU for numpy operations
-        return {
-            "non-tumor": float(probs[0]),
-            "tumor": float(probs[1])
-        }
+        probs = predictions[0].cpu().tolist()
+        # Return just the predicted class instead of probabilities
+        return "tumor" if probs[1] > probs[0] else "non-tumor"
 
-    @torch.no_grad()  # Disable gradient computation for inference
+    @torch.no_grad()
     def predict_tumor_radius(self, processed_image):
         inputs = self.radius_processor(processed_image, return_tensors="pt")
-        inputs = {k: v.to(self.device) for k, v in inputs.items()}  # Move inputs to correct device
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
         outputs = self.radius_model(**inputs)
         predictions = outputs.logits.softmax(dim=-1)
         predicted_label = predictions.argmax().item()
-        confidence = predictions[0][predicted_label].cpu().item()  # Move back to CPU
-        
         class_names = ["no-tumor", "0.5", "1.0", "1.5"]
-        return {
-            "radius": class_names[predicted_label],
-            "confidence": float(confidence)
-        }
+        # Return just the radius without confidence
+        return class_names[predicted_label]
 
     def generate_llm_interpretation(self, tumor_presence, tumor_radius, metadata):
-        prompt = f"""<|system|>You are a medical AI assistant. Be concise but thorough.</s>
-        <|user|>Analyze these results:
-        Tumor Detection: {json.dumps(tumor_presence)}
-        Tumor Radius: {json.dumps(tumor_radius)}
-        Patient: {metadata['age']}y/o {metadata['gender']}, {metadata['smoking_status']}, {metadata['drinking_status']}
-        Medications: {', '.join(metadata['medications']) if metadata['medications'] else 'None'}
-        Provide: 1. Key findings 2. Risk assessment 3. Recommendations</s>
+        prompt = f"""<|system|>You are a medical AI assistant. Provide a clear and concise medical interpretation.</s>
+        <|user|>Analyze the following medical findings:
+
+        Image Analysis:
+        - Tumor Detection: {tumor_presence}
+        - Tumor Size: {tumor_radius} cm
+
+        Patient Profile:
+        - Age: {metadata['age']} years
+        - Gender: {metadata['gender']}
+        - Smoking: {metadata['smoking_status']}
+        - Alcohol: {metadata['drinking_status']}
+        - Current Medications: {', '.join(metadata['medications']) if metadata['medications'] else 'None'}
+
+        Provide a brief:
+        1. Key findings
+        2. Clinical recommendations
+        3. Follow-up plan</s>
         <|assistant|>"""
 
         response = self.llm(
             prompt,
-            max_new_tokens=300,  # Reduced for faster response
+            max_new_tokens=300,
             temperature=0.7,
             do_sample=True,
             top_p=0.95,
@@ -114,51 +123,44 @@ class MedicalImageAnalysisSystem:
 
     def analyze_image(self, image):
         try:
-            # Add progress updates
-            yield "Processing image..."
+            # Process image and generate metadata
             processed_image = self.process_image(image)
-            
-            yield "Generating patient metadata..."
             metadata = self.generate_synthetic_metadata()
             
-            yield "Analyzing tumor presence..."
+            # Get predictions
             tumor_presence = self.predict_tumor_presence(processed_image)
-            
-            yield "Analyzing tumor radius..."
             tumor_radius = self.predict_tumor_radius(processed_image)
             
-            yield "Generating medical interpretation..."
+            # Generate interpretation
             interpretation = self.generate_llm_interpretation(
                 tumor_presence,
                 tumor_radius,
                 metadata
             )
             
-            # Final results
-            result = {
+            # Format results
+            results = {
                 "metadata": metadata,
                 "tumor_presence": tumor_presence,
                 "tumor_radius": tumor_radius,
                 "interpretation": interpretation
             }
             
-            yield self.format_results(result)
+            return self.format_results(results)
             
         except Exception as e:
-            yield f"Error: {str(e)}"
+            return f"Error: {str(e)}"
 
     def format_results(self, results):
         return f"""
-        Patient Metadata:
+        Patient Information:
         {json.dumps(results['metadata'], indent=2)}
         
-        Tumor Presence Analysis:
-        {json.dumps(results['tumor_presence'], indent=2)}
-        
-        Tumor Radius Analysis:
-        {json.dumps(results['tumor_radius'], indent=2)}
+        Image Analysis Results:
+        - Tumor Detection: {results['tumor_presence']}
+        - Tumor Size: {results['tumor_radius']} cm
         
-        Medical Interpretation and Recommendations:
+        Medical Assessment:
         {results['interpretation']}
         """
 
@@ -175,8 +177,7 @@ def create_interface():
         ],
         title="Medical Image Analysis System",
         description="Upload a medical image for tumor analysis and recommendations.",
-        theme=gr.themes.Base(),
-        flagging=False
+        theme=gr.themes.Base()
     )
     
     return iface
@@ -184,5 +185,5 @@ def create_interface():
 if __name__ == "__main__":
     print("Starting application...")
     iface = create_interface()
-    iface.queue()  # Enable queuing for better handling of multiple requests
+    iface.queue()
     iface.launch(debug=True, share=True)