Update app.py

app.py

@@ -7,240 +7,172 @@ import tensorflow_hub as hub
 import numpy as np
 from PIL import Image
 import io
-from typing import Optional, Dict, Any, Union
+import os
 
-# Set up logging
+# Set up logging with more detailed format
 logging.basicConfig(
     level=logging.INFO,
     format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
 )
 logger = logging.getLogger(__name__)
 
-class MedicalDiagnosisModel:
-    def __init__(self, model_path: str):
-        self.model_path = model_path
-        self.model = self._load_model()
+# API key and user ID for on-demand
+api_key = 'KGSjxB1uptfSk8I8A7ciCuNT9Xa3qWC3'
+external_user_id = 'plugin-1717464304'
+
+def load_model():
+    try:
+        model_path = 'model_epoch_01.h5.keras'
         
-    def _load_model(self) -> Optional[tf_keras.Model]:
-        """Load the transfer learning model with proper error handling."""
-        try:
-            if not os.path.exists(self.model_path):
-                raise FileNotFoundError(f"Model file not found at {self.model_path}")
-                
-            logger.info(f"Loading model from {self.model_path}")
-            
-            # Define custom objects dictionary for transfer learning
-            custom_objects = {
-                'KerasLayer': hub.KerasLayer
-            }
-            
-            try:
-                logger.info("Attempting to load model with custom objects...")
-                with tf_keras.utils.custom_object_scope(custom_objects):
-                    model = tf_keras.models.load_model(self.model_path, compile=False)
-            except Exception as e:
-                logger.error(f"Failed to load with custom objects: {str(e)}")
-                logger.info("Attempting to load model without custom objects...")
-                model = tf_keras.models.load_model(self.model_path, compile=False)
-                
-            model.summary()
-            logger.info("Model loaded successfully")
-            return model
-            
-        except Exception as e:
-            logger.error(f"Error loading model: {str(e)}")
-            return None
-            
-    def preprocess_image(self, image: np.ndarray) -> np.ndarray:
-        """Preprocess the input image for model prediction."""
-        try:
-            logger.info(f"Input image shape: {image.shape}, dtype: {image.dtype}")
-            
-            # If image is RGBA, convert to RGB
-            if image.shape[-1] == 4:
-                logger.info("Converting RGBA to RGB")
-                # Convert to PIL Image and back to handle RGBA->RGB conversion
-                image = Image.fromarray(image).convert('RGB')
-                image = np.array(image)
-            
-            # Resize image
-            image = tf_keras.preprocessing.image.smart_resize(
-                image, (256, 256), interpolation='bilinear'
-            )
-            
-            # Ensure values are between 0 and 1
-            if image.max() > 1.0:
-                image = image / 255.0
-                
-            # Add batch dimension if not present
-            if len(image.shape) == 3:
-                image = np.expand_dims(image, axis=0)
-                
-            logger.info(f"Preprocessed image shape: {image.shape}")
-            return image
-            
-        except Exception as e:
-            logger.error(f"Error preprocessing image: {str(e)}")
-            raise
+        # Check if model file exists
+        if not os.path.exists(model_path):
+            raise FileNotFoundError(f"Model file not found at {model_path}")
             
-    def predict(self, image: np.ndarray) -> Dict[str, float]:
-        """Run model prediction and return results."""
-        try:
-            prediction = self.model.predict(image)
-            return {
-                "prediction": float(prediction[0][0]),
-                "confidence": float(prediction[0][0]) * 100
-            }
-        except Exception as e:
-            logger.error(f"Error during prediction: {str(e)}")
-            raise
-
-class MedicalDiagnosisAPI:
-    def __init__(self, api_key: str, user_id: str):
-        self.api_key = api_key
-        self.user_id = user_id
-        self.base_url = "https://api.example.com/v1"  # Replace with actual API URL
+        logger.info(f"Attempting to load model from {model_path}")
+        
+        # Define custom objects dictionary
+        custom_objects = {
+            'KerasLayer': hub.KerasLayer
+            # Add more custom objects if needed
+        }
         
-    def create_chat_session(self) -> str:
-        """Create a new chat session and return session ID."""
+        # Try loading with different configurations
         try:
-            response = requests.post(
-                f"{self.base_url}/sessions",
-                headers={
-                    "Authorization": f"Bearer {self.api_key}",
-                    "X-User-ID": self.user_id
-                }
-            )
-            response.raise_for_status()
-            return response.json()["session_id"]
+            logger.info("Attempting to load model with custom objects...")
+            with tf.keras.utils.custom_object_scope(custom_objects):
+                model = tf_keras.models.load_model(model_path, compile=False)
         except Exception as e:
-            logger.error(f"Error creating chat session: {str(e)}")
-            raise
-            
-    def submit_query(self, session_id: str, patient_info: str, 
-                    image_analysis: Optional[str] = None) -> Dict[str, Any]:
-        """Submit a query to the API and return the response."""
-        try:
-            payload = {
-                "patient_info": patient_info,
-                "image_analysis": image_analysis
-            }
+            logger.error(f"Failed to load with custom objects: {str(e)}")
+            logger.info("Attempting to load model without custom objects...")
+            model = tf_keras.models.load_model(model_path, compile=False)
+        
+        # Verify model loaded correctly
+        if model is None:
+            raise ValueError("Model loading returned None")
             
-            response = requests.post(
-                f"{self.base_url}/sessions/{session_id}/query",
-                headers={
-                    "Authorization": f"Bearer {self.api_key}",
-                    "X-User-ID": self.user_id
-                },
-                json=payload
-            )
-            response.raise_for_status()
-            return response.json()
-        except Exception as e:
-            logger.error(f"Error submitting query: {str(e)}")
-            raise
+        # Print model summary for debugging
+        model.summary()
+        logger.info("Model loaded successfully")
+        return model
+        
+    except Exception as e:
+        logger.error(f"Error loading model: {str(e)}")
+        logger.error(f"Model loading failed with exception type: {type(e)}")
+        raise
 
-def extract_json_from_answer(answer: str) -> Dict[str, Any]:
-    """Extract and parse JSON from the API response."""
+# Initialize the model globally
+try:
+    logger.info("Initializing model...")
+    model = load_model()
+    logger.info("Model initialization completed")
+except Exception as e:
+    logger.error(f"Failed to initialize model: {str(e)}")
+    model = None
+    
+def preprocess_image(image):
     try:
-        # Find JSON content between triple backticks if present
-        if "```json" in answer and "```" in answer:
-            json_str = answer.split("```json")[1].split("```")[0].strip()
-        else:
-            json_str = answer.strip()
-            
-        return json.loads(json_str)
+        # Log image shape and type for debugging
+        logger.info(f"Input image shape: {image.shape}, dtype: {image.dtype}")
+        
+        image = image.convert('rgb')
+        image = image.resize((256, 256, 3))
+        image = np.array(image)
+        
+        # Normalize pixel values
+        image = image / 255.0
+        
+        # Add batch dimension
+        image = np.expand_dims(image, axis=0)
+        logger.info(f"Final preprocessed image shape: {image.shape}")
+        
+        return image
+        
     except Exception as e:
-        logger.error(f"Error extracting JSON from answer: {str(e)}")
+        logger.error(f"Error preprocessing image: {str(e)}")
         raise
 
-class MedicalDiagnosisApp:
-    def __init__(self, model_path: str, api_key: str, user_id: str):
-        self.model = MedicalDiagnosisModel(model_path)
-        self.api = MedicalDiagnosisAPI(api_key, user_id)
-        
-    def process_request(self, patient_info: str, 
-                       image: Optional[np.ndarray]) -> str:
-        """Process a medical diagnosis request."""
-        try:
-            if self.model.model is None:
-                return json.dumps({
-                    "error": "Model initialization failed",
-                    "status": "error"
-                }, indent=2)
-                
-            # Process image if provided
-            image_analysis = None
-            if image is not None:
-                processed_image = self.model.preprocess_image(image)
-                image_analysis = self.model.predict(processed_image)
-                logger.info(f"Image analysis results: {image_analysis}")
-                
-            # Create chat session and submit query
-            session_id = self.api.create_chat_session()
-            llm_response = self.api.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")
-                
-            json_data = extract_json_from_answer(llm_response['data']['answer'])
-            return json.dumps(json_data, indent=2)
-            
-        except Exception as e:
-            logger.error(f"Error processing request: {str(e)}")
+def gradio_interface(patient_info, image):
+    try:
+        if model is None:
+            logger.error("Model is not initialized")
             return json.dumps({
-                "error": str(e),
-                "status": "error",
-                "details": "Check the application logs for more information"
+                "error": "Model initialization failed. Please check the logs for details.",
+                "status": "error"
             }, indent=2)
+            
+        # Process image if provided
+        image_analysis = None
+        if image is not None:
+            logger.info("Processing uploaded image")
+            # Preprocess image
+            processed_image = preprocess_image(image)
+            
+            # Get model prediction
+            logger.info("Running model prediction")
+            prediction = model.predict(processed_image)
+            logger.info(f"Raw prediction shape: {prediction.shape}")
+            logger.info(f"Prediction: {prediction}")
+            # Format prediction results
+            image_analysis = {
+                "prediction": float(prediction[0][0]),
+                "confidence": float(prediction[0][0]) * 100
+            }
+            logger.info(f"Image analysis results: {image_analysis}")
+        
+        # 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'])
+        
+        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),
+            "status": "error",
+            "details": "Check the application logs for more information"
+        }, indent=2)
 
-def create_gradio_interface() -> gr.Interface:
-    """Create and configure the Gradio interface."""
-    app = MedicalDiagnosisApp(
-        model_path='model_epoch_01.h5.keras',
-        api_key='KGSjxB1uptfSk8I8A7ciCuNT9Xa3qWC3',
-        user_id='plugin-1717464304'
-    )
-    
-    return gr.Interface(
-        fn=app.process_request,
-        inputs=[
-            gr.Textbox(
-                label="Patient Information",
-                placeholder="Enter patient details including: symptoms, medical history, current medications, age, gender, and any relevant test results...",
-                lines=5,
-                max_lines=10
-            ),
-            gr.Image(
-                label="Medical Image",
-                type="numpy",
-                interactive=True
-            )
-        ],
-        outputs=gr.Textbox(
-            label="Medical Analysis",
-            placeholder="JSON analysis will appear here...",
-            lines=15
+# Gradio interface
+iface = gr.Interface(
+    fn=gradio_interface,
+    inputs=[
+        gr.Textbox(
+            label="Patient Information",
+            placeholder="Enter patient details including: symptoms, medical history, current medications, age, gender, and any relevant test results...",
+            lines=5,
+            max_lines=10
         ),
-        title="Medical Diagnosis Assistant",
-        description="Enter patient information and optionally upload a medical image for analysis."
-    )
+        gr.Image(
+            label="Medical Image",
+            type="numpy",
+            interactive=True
+        )
+    ],
+    outputs=gr.Textbox(
+        label="Medical Analysis",
+        placeholder="JSON analysis will appear here...",
+        lines=15
+    ),
+    title="Medical Diagnosis Assistant",
+    description="Enter patient information and optionally upload a medical image for analysis."
+)
 
 if __name__ == "__main__":
-    # Log version information
-    logger.info(f"TF-Keras version: {tf_keras.__version__}")
+    # Add version information logging
+    logger.info(f"TensorFlow Keras version: {tf_keras.__version__}")
     logger.info(f"TensorFlow Hub version: {hub.__version__}")
     logger.info(f"Gradio version: {gr.__version__}")
     
-    # Create and launch the interface
-    iface = create_gradio_interface()
     iface.launch(
         server_name="0.0.0.0",
-        share=True,  # Enable public link
         debug=True
     )