fixed rexvqa benchmark and added handling for image norm for tools

benchmarking/benchmarks/rexvqa_benchmark.py

@@ -46,10 +46,10 @@ class ReXVQABenchmark(Benchmark):
         self.image_dataset = None
         self.image_mapping = {}  # Maps study_id to image data
         
-        super().__init__(data_dir, **kwargs)
+        # Set images_dir BEFORE parent initialization to avoid AttributeError
+        self.images_dir = f"{data_dir}/images/deid_png"
         
-        # Set images_dir after parent initialization
-        self.images_dir = f"{self.data_dir}/images/deid_png"
+        super().__init__(data_dir, **kwargs)
 
     @staticmethod
     def download_rexgradient_images(output_dir: str = "benchmarking/data/rexvqa", repo_id: str = "rajpurkarlab/ReXGradient-160K", test_only: bool = True):

medrax/tools/classification/torchxrayvision.py

@@ -12,6 +12,8 @@ from langchain_core.callbacks import (
 )
 from langchain_core.tools import BaseTool
 
+from medrax.utils.utils import preprocess_medical_image
+
 
 class TorchXRayVisionInput(BaseModel):
     """Input for TorchXRayVision chest X-ray analysis tools. Only supports JPG or PNG images."""
@@ -76,7 +78,9 @@ class TorchXRayVisionClassifierTool(BaseTool):
             ValueError: If the image cannot be properly loaded or processed.
         """
         img = skimage.io.imread(image_path)
-        img = xrv.datasets.normalize(img, 255)
+        
+        # Use robust normalization that handles both 8-bit and 16-bit images
+        img = preprocess_medical_image(img, target_range=(-1024.0, 1024.0))
 
         if len(img.shape) > 2:
             img = img[:, :, 0]

medrax/tools/segmentation/segmentation.py

@@ -20,6 +20,8 @@ from langchain_core.callbacks import (
 )
 from langchain_core.tools import BaseTool
 
+from medrax.utils.utils import preprocess_medical_image
+
 
 class ChestXRaySegmentationInput(BaseModel):
     """Input schema for the Chest X-ray Segmentation Tool."""
@@ -246,7 +248,8 @@ class ChestXRaySegmentationTool(BaseTool):
             if len(original_img.shape) > 2:
                 original_img = original_img[:, :, 0]
 
-            img = xrv.datasets.normalize(original_img, 255)
+            # Use robust normalization that handles both 8-bit and 16-bit images
+            img = preprocess_medical_image(original_img)
             img = img[None, ...]
             img = self.transform(img)
             img = torch.from_numpy(img)

medrax/utils/utils.py

@@ -1,6 +1,90 @@
 import os
 import json
-from typing import Dict, List
+import numpy as np
+from typing import Dict, List, Union, Tuple
+
+
+def preprocess_medical_image(
+    image: np.ndarray, 
+    target_range: Tuple[float, float] = (0.0, 1.0),
+    clip_values: bool = True
+) -> np.ndarray:
+    """
+    Preprocess medical images by auto-detecting bit depth and normalizing appropriately.
+    
+    This function handles both 8-bit (0-255) and 16-bit (0-65535) images automatically,
+    normalizing them to the target range. It's designed for medical imaging tools that
+    expect consistent input ranges regardless of the original image bit depth.
+    
+    Args:
+        image (np.ndarray): Input image array (2D or 3D)
+        target_range (Tuple[float, float]): Target range for normalization (default: (0.0, 1.0))
+        clip_values (bool): Whether to clip values to target range (default: True)
+    
+    Returns:
+        np.ndarray: Normalized image in the target range
+        
+    Raises:
+        ValueError: If image is empty or has invalid values
+        ValueError: If target_range is invalid
+    """
+    if image.size == 0:
+        raise ValueError("Input image is empty")
+    
+    if len(target_range) != 2 or target_range[0] >= target_range[1]:
+        raise ValueError("target_range must be a tuple of (min, max) where min < max")
+    
+    # Convert to float for processing
+    image = image.astype(np.float32)
+    
+    # Auto-detect bit depth based on maximum value
+    max_val = np.max(image)
+    min_val = np.min(image)
+    
+    # Determine the expected maximum value based on bit depth
+    if max_val <= 255:
+        # 8-bit image
+        expected_max = 255.0
+    elif max_val <= 65535:
+        # 16-bit image
+        expected_max = 65535.0
+    else:
+        # Higher bit depth or already normalized, use actual max
+        expected_max = max_val
+    
+    # Normalize to 0-1 range first
+    if expected_max > 0:
+        image = (image - min_val) / (expected_max - min_val)
+    else:
+        # Handle edge case where image has no contrast
+        image = np.zeros_like(image)
+    
+    # Scale to target range
+    target_min, target_max = target_range
+    image = image * (target_max - target_min) + target_min
+    
+    # Clip values if requested
+    if clip_values:
+        image = np.clip(image, target_min, target_max)
+    
+    return image
+
+
+def normalize_medical_image_for_torchxrayvision(image: np.ndarray) -> np.ndarray:
+    """
+    Normalize medical images specifically for TorchXRayVision models.
+    
+    This function is a convenience wrapper around preprocess_medical_image
+    that normalizes images to the -1024 to 1024 range expected by TorchXRayVision models.
+    This range corresponds to the Hounsfield Unit scale adapted for X-ray images.
+    
+    Args:
+        image (np.ndarray): Input image array (2D or 3D)
+    
+    Returns:
+        np.ndarray: Normalized image in -1024 to 1024 range
+    """
+    return preprocess_medical_image(image, target_range=(-1024.0, 1024.0))
 
 
 def load_prompts_from_file(file_path: str) -> Dict[str, str]: