update model and code for handling TSM layers

.gitignore

@@ -1 +1,2 @@
-*.pyc*
+*.pyc*
+classified_damage_sites.*

app.py

@@ -68,14 +68,44 @@ def damage_classification(SEM_image,image_threshold, model1_threshold, model2_th
     ##
     logging.debug('---------------: prepare model 1 :=====================')
     images_model1 = utils.prepare_classifier_input(SEM_image, all_centroids, window_size=model1_windowsize)
+    from utils import debug_classification_input
+    debug_classification_input(images_model1)
     
     logging.debug('---------------: run model 1 :=====================')
     #y1_pred = model1.predict(np.asarray(images_model1, float))
-    y1_pred = model1(np.asarray(images_model1, float))
-
-    logging.debug('---------------: model1 threshold :=====================')
-    inclusions = y1_pred[:,0].reshape(len(y1_pred),1)
+    #y1_pred = model1(np.asarray(images_model1, float))
+    # logging.debug('---------------: model1 threshold :=====================')
+    # inclusions = y1_pred[:,0].reshape(len(y1_pred),1)
+    # inclusions = np.where(inclusions > model1_threshold)
+
+    batch_model1 = np.array(images_model1, dtype=np.float32)
+    logging.info(f"Model 1 input shape: {batch_model1.shape}")
+    # Get predictions from model 1
+    y1_pred_raw = model1(batch_model1)
+    logging.info(f"Model 1 raw output type: {type(y1_pred_raw)}")
+    
+    # Extract actual predictions from the model output
+    y1_pred = utils.extract_predictions_from_tfsm(y1_pred_raw)
+    logging.info(f"Model 1 predictions shape: {y1_pred.shape}")
+    logging.info(f"Model 1 predictions sample: {y1_pred[:3] if len(y1_pred) > 0 else 'Empty'}")
+    
+    logging.info('---------------: model1 threshold :=====================')
+    # Handle predictions based on their shape
+    if len(y1_pred.shape) == 2:
+        # Predictions are 2D: (batch_size, num_classes)
+        inclusions = y1_pred[:, 0]  # Get first column (inclusion probability)
+    elif len(y1_pred.shape) == 1:
+        # Predictions are 1D: (batch_size,)
+        inclusions = y1_pred
+    else:
+        raise ValueError(f"Unexpected prediction shape: {y1_pred.shape}")
+
+
+    logging.info('---------------: model1 threshold :=====================')
     inclusions = np.where(inclusions > model1_threshold)
+    logging.info('Inclusions found at indices:')
+    logging.info(inclusions)
+
 
     logging.debug('---------------: model 1 update dict :=====================')
     for i in range(len(inclusions[0])):
@@ -105,7 +135,19 @@ def damage_classification(SEM_image,image_threshold, model1_threshold, model2_th
     
     logging.debug('---------------: run model 2 :=====================')
     #y2_pred = model2.predict(np.asarray(images_model2, float))
-    y2_pred = model2(np.asarray(images_model2, float))
+    #y2_pred = model2(np.asarray(images_model2, float))
+    batch_model2 = np.array(images_model2, dtype=np.float32)
+    logging.info(f"Model 2 input shape: {batch_model2.shape}")
+    # Get predictions from model 2
+    y2_pred_raw = model2(batch_model2)
+    logging.info(f"Model 2 raw output type: {type(y2_pred_raw)}")
+    # Extract actual predictions from the model output
+    y2_pred = utils.extract_predictions_from_tfsm(y2_pred_raw)
+    logging.info(f"Model 2 predictions shape: {y2_pred.shape}")
+    logging.info(f"Model 2 predictions sample: {y2_pred[:3] if len(y2_pred) > 0 else 'Empty'}")
+    logging.info(y2_pred)
+
+    logging.debug('---------------: model2 threshold :=====================')   
 
     damage_index = np.asarray(y2_pred > model2_threshold).nonzero()
 

rwthmaterials_dp800_network2_damage.tgz

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c0b85a55ba6f970661ae81eb14063dd157538d310ca4aa3288e9f263f58b6749
-size 80796486
+oid sha256:9db591833c36177cb865e13a6332eaa15db98a3104860d9c3749214da6e06e0e
+size 3943261

rwthmaterials_dp800_network2_damage/fingerprint.pb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:42b75267132948545aba844b272093f741daea371271c5adca0122be1bfb91cf
-size 59
+oid sha256:ca1d16acc9246b42769f826cb6b63586f8d922ed0c58ce42f9ce64d8b2680725
+size 56

rwthmaterials_dp800_network2_damage/saved_model.pb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8644d5da2015e6a2ce6fd4181ede406e98171e7fa170ab256411a639b1bbb014
-size 4482600
+oid sha256:6c243d6b76e1fe71cbbe1bf7312f7148409b04d44dd796c0b1ea1219f1aaec06
+size 620656

rwthmaterials_dp800_network2_damage/variables/variables.data-00000-of-00001

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f65b218b3718aac76d545ebef2d6a6ef9b4c2d7e22e7885bc3dd91836563f5ca
-size 87475223
+oid sha256:a6faa8b229462f1cb0e7d57ccf3df66ab3935bff97d685a7154cd4ec381b5a34
+size 3977928

utils.py

@@ -124,11 +124,110 @@ def show_boxes(image : np.ndarray, damage_sites : dict, box_size = [250,250],
 
     return data
 
+##
+## orig
+##
+
+# ###
+# ### cut out small images from panorama, append colour information
+# ###
+# def prepare_classifier_input(panorama, centroids: list, window_size=[250, 250]) -> list: # Removed np.ndarray type hint for panorama
+#     """
+#     Extracts square image patches centered at each given centroid from a grayscale panoramic SEM image.
+    
+#     Each extracted patch is resized to the specified window size and converted into a 3-channel (RGB-like)
+#     normalized image suitable for use with classification neural networks that expect color input.
+
+#     Parameters
+#     ----------
+#     panorama : PIL.Image.Image or np.ndarray
+#         Input SEM image. Should be a 2D array (H, W) or a 3D array (H, W, 1) representing grayscale data,
+#         or a PIL Image object.
+    
+#     centroids : list of [int, int]
+#         List of (y, x) coordinates marking the centers of regions of interest. These are typically damage sites
+#         identified in preprocessing (e.g., clustering).
+
+#     window_size : list of int, optional
+#         Size [height, width] of each extracted image patch. Defaults to [250, 250].
+
+#     Returns
+#     -------
+#     list of np.ndarray
+#         List of extracted and normalized 3-channel image patches, each with shape (height, width, 3). Only
+#         centroids that allow full window extraction within image bounds are used.
+#     """
+#     logging.info(f"prepare_classifier_input: Input panorama type: {type(panorama)}") # Added logging
+
+#     # --- MINIMAL FIX START ---
+#     # Convert PIL Image to NumPy array if necessary
+#     if isinstance(panorama, Image.Image):
+#         # Convert to grayscale NumPy array as your original code expects this structure for processing
+#         if panorama.mode == 'RGB':
+#             panorama_array = np.array(panorama.convert('L'))
+#             logging.info("prepare_classifier_input: Converted RGB PIL Image to grayscale NumPy array.")
+#         else:
+#             panorama_array = np.array(panorama)
+#             logging.info("prepare_classifier_input: Converted PIL Image to grayscale NumPy array.")
+#     elif isinstance(panorama, np.ndarray):
+#         # Ensure it's treated as a grayscale array for consistency with original logic
+#         if panorama.ndim == 3 and panorama.shape[2] in [3, 4]: # RGB or RGBA NumPy array
+#             panorama_array = np.mean(panorama, axis=2).astype(panorama.dtype) # Convert to grayscale
+#             logging.info("prepare_classifier_input: Converted multi-channel NumPy array to grayscale.")
+#         else:
+#             panorama_array = panorama # Assume it's already grayscale 2D or (H,W,1)
+#             logging.info("prepare_classifier_input: Panorama is already a suitable NumPy array.")
+#     else:
+#         logging.error("prepare_classifier_input: Unsupported panorama format received. Expected PIL Image or NumPy array.")
+#         raise ValueError("Unsupported panorama format for classifier input.")
+
+#     # Now, ensure panorama_array has a channel dimension if it's 2D for consistency
+#     if panorama_array.ndim == 2:
+#         panorama_array = np.expand_dims(panorama_array, axis=-1)  # (H, W, 1)
+#         logging.info("prepare_classifier_input: Expanded 2D panorama to 3D (H,W,1).")
+#     # --- MINIMAL FIX END ---
+
+#     H, W, _ = panorama_array.shape # Use panorama_array here
+#     win_h, win_w = window_size
+#     images = []
+
+#     for (cy, cx) in centroids:
+#         # Ensure coordinates are integers
+#         cy, cx = int(round(cy)), int(round(cx))
+
+#         x1 = int(cx - win_w / 2)
+#         y1 = int(cy - win_h / 2)
+#         x2 = x1 + win_w
+#         y2 = y1 + win_h
+
+#         # Skip if patch would go out of bounds
+#         if x1 < 0 or y1 < 0 or x2 > W or y2 > H:
+#             logging.warning(f"prepare_classifier_input: Skipping centroid ({cy},{cx}) as patch is out of bounds.") # Added warning
+#             continue
+
+#         # Extract and normalize patch
+#         patch = panorama_array[y1:y2, x1:x2, 0].astype(np.float32) # Use panorama_array
+#         patch = patch * 2. / 255. - 1. # Keep your original normalization
+
+#         # Replicate grayscale channel to simulate RGB
+#         patch_color = np.repeat(patch[:, :, np.newaxis], 3, axis=2)
+#         images.append(patch_color)
+
+#     return images
 
 ###
-### cut out small images from panorama, append colour information
+### refactored
 ###
-def prepare_classifier_input(panorama, centroids: list, window_size=[250, 250]) -> list: # Removed np.ndarray type hint for panorama
+import numpy as np
+from PIL import Image
+import logging
+from typing import List, Union, Tuple
+
+def prepare_classifier_input(
+    panorama: Union[Image.Image, np.ndarray], 
+    centroids: List[Tuple[int, int]], 
+    window_size: List[int] = [250, 250]
+) -> List[np.ndarray]:
     """
     Extracts square image patches centered at each given centroid from a grayscale panoramic SEM image.
     
@@ -154,61 +253,337 @@ def prepare_classifier_input(panorama, centroids: list, window_size=[250, 250])
         List of extracted and normalized 3-channel image patches, each with shape (height, width, 3). Only
         centroids that allow full window extraction within image bounds are used.
     """
-    logging.info(f"prepare_classifier_input: Input panorama type: {type(panorama)}") # Added logging
+    logging.info(f"prepare_classifier_input: Input panorama type: {type(panorama)}")
 
-    # --- MINIMAL FIX START ---
-    # Convert PIL Image to NumPy array if necessary
-    if isinstance(panorama, Image.Image):
-        # Convert to grayscale NumPy array as your original code expects this structure for processing
-        if panorama.mode == 'RGB':
-            panorama_array = np.array(panorama.convert('L'))
-            logging.info("prepare_classifier_input: Converted RGB PIL Image to grayscale NumPy array.")
-        else:
-            panorama_array = np.array(panorama)
-            logging.info("prepare_classifier_input: Converted PIL Image to grayscale NumPy array.")
-    elif isinstance(panorama, np.ndarray):
-        # Ensure it's treated as a grayscale array for consistency with original logic
-        if panorama.ndim == 3 and panorama.shape[2] in [3, 4]: # RGB or RGBA NumPy array
-            panorama_array = np.mean(panorama, axis=2).astype(panorama.dtype) # Convert to grayscale
-            logging.info("prepare_classifier_input: Converted multi-channel NumPy array to grayscale.")
+    # Convert input to standardized NumPy array format
+    panorama_array = _convert_to_grayscale_array(panorama)
+    
+    # Ensure we have the correct dimensions
+    if panorama_array.ndim == 2:
+        H, W = panorama_array.shape
+        logging.info("prepare_classifier_input: Working with 2D grayscale array.")
+    elif panorama_array.ndim == 3:
+        H, W, C = panorama_array.shape
+        if C == 1:
+            # Squeeze the single channel dimension for easier processing
+            panorama_array = panorama_array.squeeze(axis=2)
+            H, W = panorama_array.shape
+            logging.info("prepare_classifier_input: Squeezed single channel dimension.")
         else:
-            panorama_array = panorama # Assume it's already grayscale 2D or (H,W,1)
-            logging.info("prepare_classifier_input: Panorama is already a suitable NumPy array.")
+            logging.error(f"prepare_classifier_input: Unexpected number of channels: {C}")
+            raise ValueError(f"Expected 1 channel, got {C}")
     else:
-        logging.error("prepare_classifier_input: Unsupported panorama format received. Expected PIL Image or NumPy array.")
-        raise ValueError("Unsupported panorama format for classifier input.")
+        logging.error(f"prepare_classifier_input: Unexpected array dimensions: {panorama_array.ndim}")
+        raise ValueError(f"Expected 2D or 3D array, got {panorama_array.ndim}D")
 
-    # Now, ensure panorama_array has a channel dimension if it's 2D for consistency
-    if panorama_array.ndim == 2:
-        panorama_array = np.expand_dims(panorama_array, axis=-1)  # (H, W, 1)
-        logging.info("prepare_classifier_input: Expanded 2D panorama to 3D (H,W,1).")
-    # --- MINIMAL FIX END ---
-
-    H, W, _ = panorama_array.shape # Use panorama_array here
     win_h, win_w = window_size
     images = []
+    
+    logging.info(f"prepare_classifier_input: Image dimensions: {H}x{W}, Window size: {win_h}x{win_w}")
+    logging.info(f"prepare_classifier_input: Processing {len(centroids)} centroids")
 
-    for (cy, cx) in centroids:
+    for i, (cy, cx) in enumerate(centroids):
         # Ensure coordinates are integers
         cy, cx = int(round(cy)), int(round(cx))
-
-        x1 = int(cx - win_w / 2)
-        y1 = int(cy - win_h / 2)
-        x2 = x1 + win_w
+        
+        # Calculate patch boundaries
+        half_h, half_w = win_h // 2, win_w // 2
+        y1 = cy - half_h
         y2 = y1 + win_h
+        x1 = cx - half_w
+        x2 = x1 + win_w
 
-        # Skip if patch would go out of bounds
-        if x1 < 0 or y1 < 0 or x2 > W or y2 > H:
-            logging.warning(f"prepare_classifier_input: Skipping centroid ({cy},{cx}) as patch is out of bounds.") # Added warning
+        # Check bounds more explicitly
+        if y1 < 0 or x1 < 0 or y2 > H or x2 > W:
+            logging.warning(
+                f"prepare_classifier_input: Skipping centroid {i+1}/{len(centroids)} "
+                f"at ({cy},{cx}) - patch bounds ({y1}:{y2}, {x1}:{x2}) exceed image bounds (0:{H}, 0:{W})"
+            )
             continue
 
-        # Extract and normalize patch
-        patch = panorama_array[y1:y2, x1:x2, 0].astype(np.float32) # Use panorama_array
-        patch = patch * 2. / 255. - 1. # Keep your original normalization
-
-        # Replicate grayscale channel to simulate RGB
-        patch_color = np.repeat(patch[:, :, np.newaxis], 3, axis=2)
-        images.append(patch_color)
+        try:
+            # Extract patch with explicit bounds checking
+            patch = panorama_array[y1:y2, x1:x2].astype(np.float32)
+            
+            # Verify patch dimensions
+            if patch.shape != (win_h, win_w):
+                logging.warning(
+                    f"prepare_classifier_input: Patch {i+1} has unexpected shape {patch.shape}, "
+                    f"expected ({win_h}, {win_w}). Skipping."
+                )
+                continue
+            
+            # Normalize patch: [0, 255] -> [-1, 1]
+            patch_normalized = (patch * 2.0 / 255.0) - 1.0
+            
+            # Convert to 3-channel RGB-like format
+            patch_rgb = np.stack([patch_normalized] * 3, axis=2)
+            
+            images.append(patch_rgb)
+            logging.debug(f"prepare_classifier_input: Successfully processed centroid {i+1} at ({cy},{cx})")
+            
+        except Exception as e:
+            logging.error(
+                f"prepare_classifier_input: Error processing centroid {i+1} at ({cy},{cx}): {e}"
+            )
+            continue
 
+    logging.info(f"prepare_classifier_input: Successfully extracted {len(images)} patches from {len(centroids)} centroids")
+    
+    # Add diagnostic information about the output
+    if images:
+        sample_shape = images[0].shape
+        sample_dtype = images[0].dtype
+        sample_min = images[0].min()
+        sample_max = images[0].max()
+        logging.info(f"prepare_classifier_input: Output patches - Shape: {sample_shape}, Dtype: {sample_dtype}, Range: [{sample_min:.3f}, {sample_max:.3f}]")
+        
+        # Verify all patches have consistent shapes
+        shapes = [img.shape for img in images]
+        if not all(shape == sample_shape for shape in shapes):
+            logging.warning("prepare_classifier_input: Inconsistent patch shapes detected!")
+            for i, shape in enumerate(shapes):
+                if shape != sample_shape:
+                    logging.warning(f"  Patch {i}: {shape} (expected {sample_shape})")
+    else:
+        logging.warning("prepare_classifier_input: No valid patches were extracted!")
+    
     return images
 
+
+def _convert_to_grayscale_array(panorama: Union[Image.Image, np.ndarray]) -> np.ndarray:
+    """
+    Helper function to convert various input formats to a standardized grayscale NumPy array.
+    
+    Parameters
+    ----------
+    panorama : PIL.Image.Image or np.ndarray
+        Input image in various formats
+        
+    Returns
+    -------
+    np.ndarray
+        Standardized grayscale array
+    """
+    if isinstance(panorama, Image.Image):
+        if panorama.mode in ['RGB', 'RGBA']:
+            # Convert to grayscale
+            panorama_array = np.array(panorama.convert('L'))
+            logging.info("_convert_to_grayscale_array: Converted RGB/RGBA PIL Image to grayscale.")
+        elif panorama.mode == 'L':
+            panorama_array = np.array(panorama)
+            logging.info("_convert_to_grayscale_array: Converted grayscale PIL Image to NumPy array.")
+        else:
+            # Handle other modes by converting to grayscale
+            panorama_array = np.array(panorama.convert('L'))
+            logging.info(f"_convert_to_grayscale_array: Converted PIL Image mode '{panorama.mode}' to grayscale.")
+            
+    elif isinstance(panorama, np.ndarray):
+        if panorama.ndim == 2:
+            # Already grayscale
+            panorama_array = panorama.copy()
+            logging.info("_convert_to_grayscale_array: Using existing 2D grayscale array.")
+        elif panorama.ndim == 3:
+            if panorama.shape[2] in [3, 4]:  # RGB or RGBA
+                # Convert to grayscale using luminance weights
+                if panorama.shape[2] == 3:  # RGB
+                    panorama_array = np.dot(panorama, [0.299, 0.587, 0.114]).astype(panorama.dtype)
+                else:  # RGBA
+                    panorama_array = np.dot(panorama[:, :, :3], [0.299, 0.587, 0.114]).astype(panorama.dtype)
+                logging.info("_convert_to_grayscale_array: Converted multi-channel NumPy array to grayscale using luminance weights.")
+            elif panorama.shape[2] == 1:
+                # Already single channel
+                panorama_array = panorama.copy()
+                logging.info("_convert_to_grayscale_array: Using existing single-channel array.")
+            else:
+                raise ValueError(f"Unsupported number of channels: {panorama.shape[2]}")
+        else:
+            raise ValueError(f"Unsupported array dimensions: {panorama.ndim}")
+    else:
+        raise ValueError(f"Unsupported panorama type: {type(panorama)}")
+    
+    return panorama_array
+
+
+##
+##  debug
+## 
+import numpy as np
+import logging
+from typing import List, Any
+
+def debug_classification_input(patches: List[np.ndarray], model: Any = None) -> None:
+    """
+    Debug function to help identify issues in the classification pipeline.
+    Call this right before your classification step.
+    
+    Parameters
+    ----------
+    patches : List[np.ndarray]
+        List of image patches from prepare_classifier_input
+    model : Any, optional
+        Your classification model (for additional debugging)
+    """
+    logging.info("=== CLASSIFICATION DEBUG INFO ===")
+    logging.info(f"Number of patches: {len(patches)}")
+    
+    if not patches:
+        logging.error("No patches provided for classification!")
+        return
+    
+    for i, patch in enumerate(patches):
+        logging.info(f"Patch {i}:")
+        logging.info(f"  Shape: {patch.shape}")
+        logging.info(f"  Dtype: {patch.dtype}")
+        logging.info(f"  Range: [{patch.min():.3f}, {patch.max():.3f}]")
+        logging.info(f"  Memory layout: {patch.flags}")
+        
+        # Check for common issues
+        if np.isnan(patch).any():
+            logging.warning(f"  Contains NaN values: {np.isnan(patch).sum()}")
+        if np.isinf(patch).any():
+            logging.warning(f"  Contains infinite values: {np.isinf(patch).sum()}")
+        
+        # Check if patch is contiguous (some models require this)
+        if not patch.flags.c_contiguous:
+            logging.warning(f"  Patch {i} is not C-contiguous")
+    
+    # Test conversion to common formats
+    try:
+        patches_array = np.array(patches)
+        logging.info(f"Stacked array shape: {patches_array.shape}")
+        logging.info(f"Stacked array dtype: {patches_array.dtype}")
+    except Exception as e:
+        logging.error(f"Failed to stack patches into array: {e}")
+    
+    # Test batch preparation (common source of slice errors)
+    try:
+        if len(patches) > 0:
+            # Common preprocessing steps that might cause issues
+            test_batch = np.stack(patches, axis=0)  # Shape: (batch_size, height, width, channels)
+            logging.info(f"Test batch shape: {test_batch.shape}")
+            
+            # Test various indexing operations that might cause slice errors
+            test_slice = test_batch[0]  # Should work
+            logging.info(f"Single item slice shape: {test_slice.shape}")
+            
+            test_batch_slice = test_batch[:]  # Should work
+            logging.info(f"Full batch slice shape: {test_batch_slice.shape}")
+            
+    except Exception as e:
+        logging.error(f"Error during batch preparation testing: {e}")
+        logging.error(f"Error type: {type(e)}")
+        import traceback
+        logging.error(f"Traceback: {traceback.format_exc()}")
+    
+    logging.info("=== END CLASSIFICATION DEBUG ===")
+
+
+def safe_classify_patches(patches: List[np.ndarray], classify_func, **kwargs) -> Any:
+    """
+    Wrapper function to safely run classification with better error handling.
+    
+    Parameters
+    ----------
+    patches : List[np.ndarray]
+        List of image patches
+    classify_func : callable
+        Your classification function
+    **kwargs
+        Additional arguments for classify_func
+        
+    Returns
+    -------
+    Any
+        Classification results or None if error occurred
+    """
+    try:
+        logging.info("Starting safe classification...")
+        
+        # Debug the input
+        debug_classification_input(patches)
+        
+        # Ensure patches are properly formatted
+        if not patches:
+            logging.error("No patches to classify")
+            return None
+        
+        # Make sure all patches are contiguous arrays
+        patches_clean = []
+        for i, patch in enumerate(patches):
+            if not patch.flags.c_contiguous:
+                patch_clean = np.ascontiguousarray(patch)
+                logging.info(f"Made patch {i} contiguous")
+            else:
+                patch_clean = patch
+            patches_clean.append(patch_clean)
+        
+        # Call the actual classification function
+        logging.info("Calling classification function...")
+        result = classify_func(patches_clean, **kwargs)
+        logging.info("Classification completed successfully")
+        
+        return result
+        
+    except Exception as e:
+        logging.error(f"Error in safe_classify_patches: {e}")
+        logging.error(f"Error type: {type(e)}")
+        import traceback
+        logging.error(f"Full traceback: {traceback.format_exc()}")
+        return None
+
+
+# Example usage function
+def example_usage():
+    """
+    Example of how to use the debug functions in your pipeline
+    """
+    # Your existing code that calls prepare_classifier_input
+    # patches = prepare_classifier_input(panorama, centroids, window_size)
+    
+    # Add debugging before classification
+    # debug_classification_input(patches)
+    
+    # Use safe wrapper for classification
+    # results = safe_classify_patches(patches, your_classify_function, model=your_model)
+    
+    pass
+
+
+########################################
+##
+##
+########################################
+def extract_predictions_from_tfsm(model_output):
+    """
+    Helper function to extract predictions from TFSMLayer output.
+    TFSMLayer often returns a dictionary with multiple outputs.
+    """
+    logging.info(f"Model output type: {type(model_output)}")
+    logging.info(f"Model output keys: {model_output.keys() if isinstance(model_output, dict) else 'Not a dict'}")
+    
+    if isinstance(model_output, dict):
+        # Try common output key names
+        possible_keys = ['output', 'predictions', 'dense', 'logits', 'probabilities']
+        
+        # First, log all available keys
+        available_keys = list(model_output.keys())
+        logging.info(f"Available output keys: {available_keys}")
+        
+        # Try to find the right output
+        for key in possible_keys:
+            if key in model_output:
+                logging.info(f"Using output key: {key}")
+                return model_output[key].numpy()
+        
+        # If no standard key found, use the first available key
+        if available_keys:
+            first_key = available_keys[0]
+            logging.info(f"Using first available key: {first_key}")
+            return model_output[first_key].numpy()
+        else:
+            raise ValueError("No output keys found in model response")
+    else:
+        # If it's not a dictionary, assume it's already the tensor we need
+        logging.info("Model output is not a dictionary, using directly")
+        return model_output.numpy() if hasattr(model_output, 'numpy') else np.array(model_output)