Update app.py

app.py

@@ -3,6 +3,7 @@ import gradio as gr
 import logging
 import json
 import tensorflow as tf
+import tensorflow_hub as hub
 import numpy as np
 from PIL import Image
 import io
@@ -18,7 +19,16 @@ external_user_id = 'plugin-1717464304'
 # Load the keras model
 def load_model():
     try:
-        model = tf.keras.models.load_model('model_epoch_01.h5.keras')
+        # Define custom objects dictionary
+        custom_objects = {
+            'KerasLayer': hub.KerasLayer,
+            # Add any other custom layers your model might use
+        }
+        
+        # Load model with custom object scope
+        with tf.keras.utils.custom_object_scope(custom_objects):
+            model = tf.keras.models.load_model('model_epoch_01.h5.keras')
+            
         logger.info("Model loaded successfully")
         return model
     except Exception as e:
@@ -32,8 +42,13 @@ def preprocess_image(image):
         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
-        # Note: Adjust these dimensions to match your model's requirements
         target_size = (224, 224)  # Change this to match your model's input size
         image = tf.image.resize(image, target_size)
         
@@ -83,9 +98,7 @@ def submit_query(session_id, query, image_analysis=None):
         
         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)
@@ -119,6 +132,7 @@ def extract_json_from_answer(answer):
         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:
@@ -128,11 +142,31 @@ def extract_json_from_answer(answer):
             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
-model = load_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:
@@ -143,11 +177,7 @@ def gradio_interface(patient_info, image):
             prediction = model.predict(processed_image)
             
             # Format prediction results
-            # Note: Adjust this based on your model's output format
-            image_analysis = {
-                "prediction": float(prediction[0][0]),  # Adjust indexing based on your model's output
-                "confidence": float(prediction[0][0]) * 100  # Convert to percentage
-            }
+            image_analysis = format_prediction(prediction)
         
         # Create chat session and submit query
         session_id = create_chat_session()
@@ -155,17 +185,17 @@ def gradio_interface(patient_info, image):
                                   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")
+            raise ValueError("Invalid response structure from LLM")
         
         # Extract and clean JSON from the response
         json_data = extract_json_from_answer(llm_response['data']['answer'])
         
-        # Return clean JSON string
-        return json.dumps(json_data)
+        # 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)})
+        return json.dumps({"error": str(e)}, indent=2)
 
 # Gradio interface
 iface = gr.Interface(