Update processing_cxrmate2.py

processing_cxrmate2.py

@@ -6,6 +6,7 @@ from typing import Dict, List, Union
 
 import cv2
 import numpy as np
+import pydicom
 import requests
 import torch
 import torch.nn.functional as F
@@ -298,18 +299,29 @@ class CXRMate2Processor(transformers.ProcessorMixin):
             for j in range(max_images):
                 if j < len(images[i]):
 
+                    image_np = None
+
                     if isinstance(images[i][j], bytes):
                         image = Image.open(io.BytesIO(images[i][j]))
 
                     elif isinstance(images[i][j], str):
-                        response = requests.get(images[i][j], stream=True)
-                        image = Image.open(BytesIO(response.content))
+                        if images[i][j].endswith('.dcm'):
+                            assert self.mimic_cxr_normalisation, 'MIMIC-CXR normalisation must be True when using DICOM images.'
+                            ds = pydicom.dcmread(images[i][j])
+                            image_np = ds.pixel_array.astype(float)
+                    
+                        else:
+                            response = requests.get(images[i][j], stream=True)
+                            image = Image.open(BytesIO(response.content))
                     
                     elif isinstance(images[i][j], Image.Image):
                         image = images[i][j]
 
                     if self.mimic_cxr_normalisation:
-                        image_np = np.array(image.convert('L'), dtype=np.float32)
+
+                        # MIMIC-CXR normalisation:
+                        if image_np is None:
+                            image_np = np.array(image.convert('L'), dtype=np.float32)
                         assert image_np.ndim == 2
                         min_val = image_np.min()
                         denom = image_np.max() - min_val