Add pneumonia detection app with Grad-CAM

app.py

@@ -102,17 +102,21 @@ def query_medgemma(message, history, image=None):
         # Multimodal format: Send image as base64 in the content
         image_base64 = image_to_base64(image)
         
-        # TGI multimodal format for MedGemma
+        # TGI multimodal format for MedGemma with better stopping
         payload = {
             "inputs": {
                 "text": message,
                 "image": image_base64
             },
             "parameters": {
-                "max_new_tokens": 500,
-                "temperature": 0.7,
+                "max_new_tokens": 200,  # Reduced for faster response
+                "temperature": 0.3,     # Lower temperature for more focused responses
                 "do_sample": True,
-                "return_full_text": False
+                "return_full_text": False,
+                "stop": ["<|im_end|>", "</s>", "<|endoftext|>", "\n\n", "El personal", "Llamé", "Hice"],
+                "repetition_penalty": 1.2,  # Prevent repetition
+                "top_p": 0.8,               # Nucleus sampling
+                "top_k": 40                 # Limit token choices
             }
         }
         
@@ -121,23 +125,30 @@ def query_medgemma(message, history, image=None):
             "inputs": message,
             "image": image_base64,
             "parameters": {
-                "max_new_tokens": 500,
-                "temperature": 0.7,
+                "max_new_tokens": 200,
+                "temperature": 0.3,
                 "do_sample": True,
-                "return_full_text": False
+                "return_full_text": False,
+                "stop": ["<|im_end|>", "</s>", "<|endoftext|>", "\n\n"],
+                "repetition_penalty": 1.2,
+                "top_p": 0.8,
+                "top_k": 40
             }
         }
         
     else:
-        # Text-only format
+        # Text-only format with better parameters
         payload = {
             "inputs": message,
             "parameters": {
-                "max_new_tokens": 500,
-                "temperature": 0.7,
+                "max_new_tokens": 200,
+                "temperature": 0.3,
                 "do_sample": True,
                 "return_full_text": False,
-                "stop": ["<|im_end|>", "</s>"]
+                "stop": ["<|im_end|>", "</s>", "<|endoftext|>", "\n\n", "El personal", "Llamé", "Hice"],
+                "repetition_penalty": 1.2,
+                "top_p": 0.8,
+                "top_k": 40
             }
         }
         payload_alt = None
@@ -187,15 +198,26 @@ def query_medgemma(message, history, image=None):
 def medical_chat(message, history, uploaded_image):
     """Handle medical chat with context from pneumonia detection"""
     
-    # Add context about pneumonia detection if there's an image
-    context_message = message
+    # Always add medical context to keep the model focused
     if uploaded_image is not None:
-        context_message = f"""Como asistente médico especializado en radiología, analiza esta imagen de rayos X y responde: {message}
-        
-Contexto: Esta es una radiografía de tórax que puede mostrar signos de neumonía. Proporciona información médica precisa pero recuerda que siempre se debe consultar a un profesional médico."""
+        context_message = f"""Eres MedGemma, un asistente médico especializado en radiología. Analiza esta radiografía de tórax y responde la siguiente pregunta médica de forma precisa y profesional:
+
+Pregunta: {message}
+
+Contexto: Esta es una radiografía de tórax para detectar neumonía. Enfócate en proporcionar información médica relevante, síntomas, diagnósticos o recomendaciones. Siempre menciona que se debe consultar a un profesional médico."""
+    else:
+        context_message = f"""Eres MedGemma, un asistente médico especializado. Responde la siguiente pregunta médica de forma precisa y profesional:
+
+Pregunta: {message}
+
+Instrucciones: Proporciona información médica precisa y relevante. Si no es una pregunta médica, redirige hacia temas de salud relacionados. Siempre menciona que se debe consultar a un profesional médico para diagnósticos definitivos."""
     
     response = query_medgemma(context_message, history, uploaded_image)
     
+    # Clean the response if it contains non-medical content
+    if any(keyword in response.lower() for keyword in ["requests", "python", "código", "api", "import", "status code"]):
+        response = "Como asistente médico, me especializo en temas de salud y medicina. ¿Podrías hacer una pregunta relacionada con medicina, síntomas, diagnósticos o radiografías? Estoy aquí para ayudarte con consultas médicas."
+    
     # Add the exchange to history
     history.append([message, response])
     return history, ""