Update processing_cxrmate2.py

processing_cxrmate2.py

@@ -4,6 +4,7 @@ import random
 from io import BytesIO
 from typing import Dict, List, Union
 
+import cv2
 import numpy as np
 import requests
 import torch
@@ -15,7 +16,6 @@ from transformers.feature_extraction_utils import BatchFeature
 from transformers.image_utils import ImageInput
 from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
 from utils import compute_time_delta
-import cv2
 
 try:
     from .dataset import CXRMate2Dataset
@@ -277,7 +277,10 @@ class CXRMate2Processor(transformers.ProcessorMixin):
                 batch['time_deltas_mask'][i].append(torch.tensor([0.0], dtype=torch.float32))
 
         # Map the image time delta values using the time delta map:
-        image_time_deltas = [[self.time_delta_map(compute_time_delta(j, k)) if j is not None else float('nan') for j in i] for i, k in zip(image_datetime, study_datetime, strict=True)]
+        if study_datetime is not None:
+            image_time_deltas = [[self.time_delta_map(compute_time_delta(j, k)) if j is not None else float('nan') for j in i] for i, k in zip(image_datetime, study_datetime, strict=True)] 
+        else:
+            image_time_deltas = [[float('nan') for _ in range(len(i))] for i in images]
 
         # Randomly select max_train_images_per_study if the number of images for a study exceeds max_train_images_per_study.
         for i in range(len(images)):
@@ -632,4 +635,4 @@ class CXRMate2Processor(transformers.ProcessorMixin):
         return batch
 
     def wrap_dataset(self, dataset):
-        return CXRMate2Dataset(dataset)
+        return CXRMate2Dataset(dataset)