code: import os
  import copy
  import PIL
  import torch
  import torchvision
  import numpy as np
  import math
  
  import logging
  from typing import List
  
  LOG = logging.getLogger(__name__)
  
  
  def define_path(use_jaad=True, use_pie=True, use_titan=True):
      """
      Define the correct paths to datasets'annotations and images
      """
      all_anns_paths = {'JAAD': {'anns': '../../DATA/annotations/JAAD/JAAD_DATA.pkl', 
                   'split': '../../DATA/annotations/JAAD/splits'},
           'PIE': {'anns': '../../DATA/annotations/PIE/PIE_DATA.pkl'},
           'TITAN': {'anns': '../../DATA/annotations/TITAN/titan_0_4',
                     'split':'../../DATA/annotations/TITAN/splits' }
           }
      all_image_dir = {'JAAD': '../../DATA/JAAD/images/',
                       'PIE': '../../DATA/PIE/images/',
                       'TITAN': '../../DATA/TITAN/images_anonymized/'
                       }
      anns_paths = {}
      image_dir = {}
      if use_jaad:
          anns_paths['JAAD'] = all_anns_paths['JAAD']
          image_dir['JAAD'] = all_image_dir['JAAD']
      if use_pie:
          anns_paths['PIE'] = all_anns_paths['PIE']
          image_dir['PIE'] = all_image_dir['PIE']
      if use_titan:
          anns_paths['TITAN'] = all_anns_paths['TITAN']
          image_dir['TITAN'] = all_image_dir['TITAN']
  
      return anns_paths, image_dir
      
  
  class ImageList(torch.utils.data.Dataset):
      """
      Basic dataloader for images
      """
  
      def __init__(self, image_paths, preprocess=None):
          self.image_paths = image_paths
          self.preprocess = preprocess
  
      def __getitem__(self, index):
          image_path = self.image_paths[index]
          with open(image_path, 'rb') as f:
              image = PIL.Image.open(f).convert('RGB')
          if self.preprocess is not None:
              image = self.preprocess(image)
  
          return image
  
      def __len__(self):
          return len(self.image_paths)
  
  
  class MultiLoader:
      # Class for loading data from mulitple datasets
      last_task_index = None
      
  
      def __init__(self, loaders: List[torch.utils.data.DataLoader], 
                   weights=None,  n_batches=None):
                   
          self.loaders = loaders
          self._weights = weights
  
          if self._weights is None:
              self._weights = [1.0 / len(loaders) for _ in range(len(loaders))]
          elif len(self._weights) == len(loaders) - 1:
              self._weights.append(1.0 - sum(self._weights))
          elif len(self._weights) == len(loaders):
              pass
          else:
              raise Exception('invalid dataset weights: {}'.format(self._weights))
          assert all(w > 0.0 for w in self._weights)
          sum_w = sum(self._weights)
          # normalize weights between datasets
          self._weights = [w / sum_w for w in self._weights]
          LOG.info('dataset weights: %s', self._weights)
          # set the total number of batches in one epoch
          self.n_batches = int(min(len(l) / w for l, w in zip(loaders, self._weights)))
          if n_batches is not None:
              self.n_batches = min(self.n_batches, n_batches)
  
      def __iter__(self):
          loader_iters = [iter(l) for l in self.loaders]
          # counter of loaded batches for each dataset
          n_loaded = [0 for _ in self.loaders]
          while True:
              # select loader for one iteration
              loader_index = int(np.argmin([n / w for n, w in zip(n_loaded, self._weights)]))
              next_batch = next(loader_iters[loader_index], None)
              if next_batch is None:
                  break
              n_loaded[loader_index] += 1
              MultiLoader.last_task_index = loader_index
              # generator
              yield next_batch
              # termination
              if sum(n_loaded) >= self.n_batches:
                  break
  
      def __len__(self):
          return self.n_batches
          
  
  class FrameDataset(torch.utils.data.Dataset):
  
      def __init__(self, samples, image_dir, preprocess=None):
          self.samples = samples
          self.image_dir = image_dir
          self.preprocess = preprocess
  
      def __getitem__(self, index):
          ids = list(self.samples.keys())
          idx = ids[index]
          frame = self.samples[idx]['frame']
          bbox = copy.deepcopy(self.samples[idx]['bbox'])
          source = self.samples[idx]["source"]
          anns = {'bbox': bbox, 'source': source}
          TTE = self.samples[idx]["TTE"]
          if 'trans_label' in list(self.samples[idx].keys()):
              label = self.samples[idx]['trans_label']
          else:
              label = None
          if 'behavior' in list(self.samples[idx].keys()):
              behavior = self.samples[idx]['behavior']
          else:
              behavior = [-1,-1,-1,-1] # no behavior annotations
          if 'attributes' in list(self.samples[idx].keys()):
              attributes = self.samples[idx]['attributes'] # scene attributes
          else:
              attributes = [-1,-1,-1,-1,-1,-1]
          image_path = None
          # image paths
          if source == "JAAD":
              vid = self.samples[idx]['video_number']
              image_path = os.path.join(self.image_dir['JAAD'], vid, '{:05d}.png'.format(frame))
          elif source == "PIE":
              vid = self.samples[idx]['video_number']
              sid = self.samples[idx]['set_number']
              image_path = os.path.join(self.image_dir['PIE'], sid, vid, '{:05d}.png'.format(frame))
          elif source == "TITAN":
              vid = self.samples[idx]['video_number']
              image_path = os.path.join(self.image_dir['TITAN'], vid, 'images', '{:06}.png'.format(frame))
  
          with open(image_path, 'rb') as f:
              img = PIL.Image.open(f).convert('RGB')
          if self.preprocess is not None:
              img, anns = self.preprocess(img, anns)
          img_tensor = torchvision.transforms.ToTensor()(img)
          if label is not None:
              label = torch.tensor(label)
              label = label.to(torch.float32)
              
          if math.isnan(TTE):
              pass
          else:
              TTE = round(self.samples[idx]["TTE"],2)
          TTE = torch.tensor(TTE).to(torch.float32)
          attributes = torch.tensor(attributes).to(torch.float32)
          sample = {'image': img_tensor, 'bbox': anns['bbox'], 'id': idx,
                     'label': label, 'source': source, 'TTE': TTE,
                     'attributes': attributes, 'behavior': behavior
                     }
  
          return sample
  
      def __len__(self):
          return len(self.samples.keys())
  
  
  class SequenceDataset(torch.utils.data.Dataset):
      """
      Basic dataloader for loading sequence/history samples
      """
  
      def __init__(self, samples, image_dir, preprocess=None):
          """
          :params: samples: transition history samples(dict)
                  image_dir: root dir for images extracted from video clips
                  preprocess: optional preprocessing on image tensors and annotations
          """
          self.samples = samples
          self.image_dir = image_dir
          self.preprocess = preprocess
  
      def __getitem__(self, index):
          ids = list(self.samples.keys())
          idx = ids[index]
          frames = self.samples[idx]['frame']
          bbox = copy.deepcopy(self.samples[idx]['bbox'])
          source = self.samples[idx]["source"]
          action = self.samples[idx]['action']
          TTE = round(self.samples[idx]["TTE"],2)
          if 'trans_label' in list(self.samples[idx].keys()):
              label = self.samples[idx]['trans_label']
          else:
              label = None
          bbox_new= []
          image_path = None
          # image paths
          img_tensors = []
          for i in range(len(frames)):
              anns = {'bbox': bbox[i], 'source': source}
              if source == "JAAD":
                  vid = self.samples[idx]['video_number']
                  image_path = os.path.join(self.image_dir['JAAD'], vid, '{:05d}.png'.format(frames[i]))
              elif source == "PIE":
                  vid = self.samples[idx]['video_number']
                  sid = self.samples[idx]['set_number']
                  image_path = os.path.join(self.image_dir['PIE'], sid, vid, '{:05d}.png'.format(frames[i]))
              elif source == "TITAN":
                  vid = self.samples[idx]['video_number']
                  image_path = os.path.join(self.image_dir['TITAN'], vid, 'images', '{:06}.png'.format(frames[i]))
              with open(image_path, 'rb') as f:
                  img = PIL.Image.open(f).convert('RGB')
              if self.preprocess is not None:
                  img, anns = self.preprocess(img, anns)
              img_tensors.append(torchvision.transforms.ToTensor()(img))
              bbox_new.append(anns['bbox'])
          img_tensors = torch.stack(img_tensors)
          if label is not None:
              label = torch.tensor(label)
              label = label.to(torch.float32)
          sample = {'image': img_tensors, 'bbox': bbox_new, 'action': action, 'id': idx, 'label': label, 'source': source, 'TTE': TTE }
  
          return sample
  
      def __len__(self):
          return len(self.samples.keys())
  
  
  class PaddedSequenceDataset(torch.utils.data.Dataset):
      """
      Dataloader for loading sequence/history samples,
      all sequences are padded to unify the length
      """
  
      def __init__(self, samples, image_dir, padded_length=10, preprocess=None, hflip_p=0.0):
          """
          :params: samples: transition history samples(dict)
                  image_dir: root dir for images extracted from video clips
                  padded_length: length of each sequence after padded
                  preprocess: optional preprocessing on image tensors and annotations
          """
          self.samples = samples
          self.image_dir = image_dir
          self.preprocess = preprocess
          self.padded_length = padded_length
          self.hflip_p = hflip_p
  
      def __getitem__(self, index):
          ids = list(self.samples.keys())
          idx = ids[index]
          frames = self.samples[idx]['frame']
          bbox = copy.deepcopy(self.samples[idx]['bbox'])
          source = self.samples[idx]["source"]
          action = self.samples[idx]['action']
          TTE = self.samples[idx]["TTE"]
          if source == "PIE":
             set_number = self.samples[idx]['set_number']
          else:
             set_number = None
          if 'trans_label' in list(self.samples[idx].keys()):
              label = self.samples[idx]['trans_label']
          else:
              label = None
          if 'behavior' in list(self.samples[idx].keys()):
              behavior = self.samples[idx]['behavior']
          else:
              behavior = [-1,-1,-1,-1]
          if 'attributes' in list(self.samples[idx].keys()):
              attributes = self.samples[idx]['attributes']
          else:
              attributes = [-1,-1,-1,-1,-1,-1]
          bbox_new = []
          bbox_ped_new = []
          image_path = None
          # image paths
          img_tensors = []
          hflip = True if float(torch.rand(1).item()) < self.hflip_p else False
          for i in range(len(frames)):
              anns = {'bbox': bbox[i], 'source': source}
              if source == "JAAD":
                  vid = self.samples[idx]['video_number']
                  image_path = os.path.join(self.image_dir['JAAD'], vid, '{:05d}.png'.format(frames[i]))
              elif source == "PIE":
                  vid = self.samples[idx]['video_number']
                  sid = self.samples[idx]['set_number']
                  image_path = os.path.join(self.image_dir['PIE'], sid, vid, '{:05d}.png'.format(frames[i]))
              elif source == "TITAN":
                  vid = self.samples[idx]['video_number']
                  image_path = os.path.join(self.image_dir['TITAN'], vid, 'images', '{:06}.png'.format(frames[i]))
              with open(image_path, 'rb') as f:
                  img = PIL.Image.open(f).convert('RGB')
              if hflip:
                  img = img.transpose(PIL.Image.FLIP_LEFT_RIGHT)
                  w, h = img.size
                  x_max = w - anns['bbox'][0]
                  x_min = w - anns['bbox'][2]
                  anns['bbox'][0] = x_min
                  anns['bbox'][2] = x_max
              anns['bbox_ped'] =  copy.deepcopy(anns['bbox'])
              if self.preprocess is not None:
                  img, anns = self.preprocess(img, anns)
              img_tensors.append(torchvision.transforms.ToTensor()(img))
              bbox_new.append(anns['bbox'])
              bbox_ped_new.append(anns['bbox_ped'])
      
          img_tensors = torch.stack(img_tensors)
          imgs_size = img_tensors.size()
          img_tensors_padded = torch.zeros((self.padded_length, imgs_size[1], imgs_size[2], imgs_size[3]))
          img_tensors_padded[:imgs_size[0], :, :, :] = img_tensors
          bbox_new_padded = copy.deepcopy(bbox_new)
          bbox_ped_new_padded = copy.deepcopy(bbox_ped_new)
          action_padded = copy.deepcopy(action)
          behavior_padded = copy.deepcopy(behavior)
          for i in range(imgs_size[0],self.padded_length):
              bbox_new_padded.append([0,0,0,0])
              bbox_ped_new_padded.append([0,0,0,0])
              action_padded.append(-1)
              behavior_padded.append([-1,-1,-1,-1])
          # seq_len = torch.squeeze(torch.LongTensor(imgs_size[0]))
          seq_len = imgs_size[0]
          if label is not None:
              label = torch.tensor(label)
              label = label.to(torch.float32)
          TTE_tag = -1
          if math.isnan(TTE):
              pass
          else:
              TTE = round(self.samples[idx]["TTE"],2)
          TTE = torch.tensor(TTE).to(torch.float32)
          TTE_tag = torch.tensor(TTE_tag)
          TTE_tag = TTE_tag.to(torch.float32)
          attributes = torch.tensor(attributes).to(torch.float32)
          sample = {'image': img_tensors_padded, 'bbox': bbox_new_padded, 'bbox_ped': bbox_ped_new_padded,
                    'seq_length': seq_len, 'action': action_padded, 'id': idx, 'label': label,
                    'source': source, 'TTE': TTE,  
                    'behavior': behavior_padded, 'attributes': attributes}
  
          return sample
  
      def __len__(self):
          return len(self.samples.keys())
          
  
  
apis: torch.zeros torch.rand torch.stack torch.tensor