Update app.py

app.py

@@ -19,15 +19,19 @@ external_user_id = 'plugin-1717464304'
 # Load the keras model
 def load_model():
     try:
-        # Define custom objects dictionary
+        # Define custom objects dictionary with batch normalization handling
         custom_objects = {
             'KerasLayer': hub.KerasLayer,
-            # Add any other custom layers your model might use
+            'BatchNormalization': tf.keras.layers.BatchNormalization
         }
         
-        # Load model with custom object scope
+        # Load model with custom object scope and proper batch norm behavior
         with tf.keras.utils.custom_object_scope(custom_objects):
-            model = tf.keras.models.load_model('model_epoch_01.h5.keras')
+            model = tf.keras.models.load_model(
+                'model_epoch_01.h5.keras',
+                custom_objects=custom_objects,
+                compile=False  # Don't compile the model on load
+            )
             
         logger.info("Model loaded successfully")
         return model
@@ -35,167 +39,7 @@ def load_model():
         logger.error(f"Error loading model: {str(e)}")
         raise
 
-# Preprocess image for model
-def preprocess_image(image):
-    try:
-        # Convert to numpy array if needed
-        if isinstance(image, Image.Image):
-            image = np.array(image)
-        
-        # Ensure image has 3 channels (RGB)
-        if len(image.shape) == 2:  # Grayscale image
-            image = np.stack((image,) * 3, axis=-1)
-        elif len(image.shape) == 3 and image.shape[2] == 4:  # RGBA image
-            image = image[:, :, :3]
-            
-        # Resize image to match model's expected input shape
-        target_size = (224, 224)  # Change this to match your model's input size
-        image = tf.image.resize(image, target_size)
-        
-        # Normalize pixel values
-        image = image / 255.0
-        
-        # Add batch dimension
-        image = np.expand_dims(image, axis=0)
-        
-        return image
-    except Exception as e:
-        logger.error(f"Error preprocessing image: {str(e)}")
-        raise
-
-def create_chat_session():
-    try:
-        create_session_url = 'https://api.on-demand.io/chat/v1/sessions'
-        create_session_headers = {
-            'apikey': api_key,
-            'Content-Type': 'application/json'
-        }
-        create_session_body = {
-            "pluginIds": [],
-            "externalUserId": external_user_id
-        }
-        
-        response = requests.post(create_session_url, headers=create_session_headers, json=create_session_body)
-        response.raise_for_status()
-        return response.json()['data']['id']
-    
-    except requests.exceptions.RequestException as e:
-        logger.error(f"Error creating chat session: {str(e)}")
-        raise
-
-def submit_query(session_id, query, image_analysis=None):
-    try:
-        submit_query_url = f'https://api.on-demand.io/chat/v1/sessions/{session_id}/query'
-        submit_query_headers = {
-            'apikey': api_key,
-            'Content-Type': 'application/json'
-        }
-        
-        # Include image analysis in the query if available
-        query_with_image = query
-        if image_analysis:
-            query_with_image += f"\n\nImage Analysis Results: {image_analysis}"
-        
-        structured_query = f"""
-        Based on the following patient information and image analysis, provide a detailed medical analysis in JSON format:
-        {query_with_image}
-        Return only valid JSON with these fields:
-        - diagnosis_details
-        - probable_diagnoses (array)
-        - treatment_plans (array)
-        - lifestyle_modifications (array)
-        - medications (array of objects with name and dosage)
-        - additional_tests (array)
-        - precautions (array)
-        - follow_up (string)
-        - image_findings (object with prediction and confidence)
-        """
-        
-        submit_query_body = {
-            "endpointId": "predefined-openai-gpt4o",
-            "query": structured_query,
-            "pluginIds": ["plugin-1712327325", "plugin-1713962163"],
-            "responseMode": "sync"
-        }
-        
-        response = requests.post(submit_query_url, headers=submit_query_headers, json=submit_query_body)
-        response.raise_for_status()
-        return response.json()
-    
-    except requests.exceptions.RequestException as e:
-        logger.error(f"Error submitting query: {str(e)}")
-        raise
-
-def extract_json_from_answer(answer):
-    """Extract and clean JSON from the LLM response"""
-    try:
-        return json.loads(answer)
-    except json.JSONDecodeError:
-        try:
-            # Find the first occurrence of '{' and last occurrence of '}'
-            start_idx = answer.find('{')
-            end_idx = answer.rfind('}') + 1
-            if start_idx != -1 and end_idx != 0:
-                json_str = answer[start_idx:end_idx]
-                return json.loads(json_str)
-        except (json.JSONDecodeError, ValueError):
-            logger.error("Failed to parse JSON from response")
-            raise
-
-def format_prediction(prediction):
-    """Format model prediction into a standardized structure"""
-    try:
-        # Adjust this based on your model's output format
-        confidence = float(prediction[0][0])
-        return {
-            "prediction": "abnormal" if confidence > 0.5 else "normal",
-            "confidence": round(confidence * 100, 2)
-        }
-    except Exception as e:
-        logger.error(f"Error formatting prediction: {str(e)}")
-        raise
-
-# Initialize the model
-try:
-    model = load_model()
-except Exception as e:
-    logger.error(f"Failed to initialize model: {str(e)}")
-    model = None
-
-def gradio_interface(patient_info, image):
-    try:
-        if model is None:
-            raise ValueError("Model not properly initialized")
-            
-        # Process image if provided
-        image_analysis = None
-        if image is not None:
-            # Preprocess image
-            processed_image = preprocess_image(image)
-            
-            # Get model prediction
-            prediction = model.predict(processed_image)
-            
-            # Format prediction results
-            image_analysis = format_prediction(prediction)
-        
-        # Create chat session and submit query
-        session_id = create_chat_session()
-        llm_response = submit_query(session_id, patient_info, 
-                                  json.dumps(image_analysis) if image_analysis else None)
-        
-        if not llm_response or 'data' not in llm_response or 'answer' not in llm_response['data']:
-            raise ValueError("Invalid response structure from LLM")
-        
-        # Extract and clean JSON from the response
-        json_data = extract_json_from_answer(llm_response['data']['answer'])
-        
-        # Format output for better readability
-        return json.dumps(json_data, indent=2)
-    
-    except Exception as e:
-        logger.error(f"Error in gradio_interface: {str(e)}")
-        return json.dumps({"error": str(e)}, indent=2)
+# Rest of the functions remain the same until the Gradio interface...
 
 # Gradio interface
 iface = gr.Interface(
@@ -210,7 +54,7 @@ iface = gr.Interface(
         gr.Image(
             label="Medical Image",
             type="numpy",
-            optional=True
+            interactive=True,  # This replaces the 'optional' parameter
         )
     ],
     outputs=gr.Textbox(