text: import torch
  from torch.utils.data import Dataset
  import torchvision.transforms as transforms
  from PIL import Image
  import numpy as np
  import os
  import torch.nn.functional as F
  import utils
  import random
  
  
  class ListDataset(Dataset):
      def __init__(self, list_path, img_dir="images", labels_dir="labels",  img_extensions=[".JPG"], img_size=608, train=True, bbox_minsize=0.01, brightness_range=0.25, contrast_range=0.25, hue_range=0.05, saturation_range=0.25, cross_offset=0.2):
          with open(list_path, "r") as file:
              self.img_files = file.read().splitlines()
  
          self.label_files = []
          for path in self.img_files:
              path = path.replace(img_dir, labels_dir)
              for ext in img_extensions:
                  path = path.replace(ext, ".txt")
  
              self.label_files.append(path)
  
          self.img_size = img_size
          self.to_tensor = transforms.ToTensor()
  
          self.train = train
  
          self.bbox_minsize = bbox_minsize
  
          self.brightness_range = brightness_range
          self.contrast_range = contrast_range
          self.hue_range = hue_range
          self.saturation_range = saturation_range
  
          self.cross_offset = cross_offset
  
      def __getitem__(self, index):
  
          img_path = self.img_files[index % len(self.img_files)].rstrip()
          label_path = self.label_files[index % len(self.img_files)].rstrip()
  
  
          # Getting image
          img = Image.open(img_path).convert('RGB')
          width, height = img.size
  
          if os.path.exists(label_path):
              boxes = torch.from_numpy(np.loadtxt(label_path).reshape(-1, 5))
          else:
              print(label_path)
          
          # RESIZING
          if width > height:
              ratio = height/width
              t_width = self.img_size
              t_height = int(ratio * self.img_size)
          else:
              ratio = width/height
              t_width = int(ratio * self.img_size)
              t_height = self.img_size
  
          img = transforms.functional.resize(img, (t_height, t_width))
  
          # IF TRAIN APPLY BRIGHTNESS CONTRAST HUE SATURTATION
          if self.train:
              brightness_rnd = random.uniform(1 - self.brightness_range, 1 + self.brightness_range)
              contrast_rnd = random.uniform(1 - self.contrast_range, 1 + self.contrast_range)
              hue_rnd = random.uniform(-self.hue_range, self.hue_range)
              saturation_rnd = random.uniform(1 - self.saturation_range, 1 + self.saturation_range)
  
              img = transforms.functional.adjust_brightness(img, brightness_rnd)
              img = transforms.functional.adjust_contrast(img, contrast_rnd)
              img = transforms.functional.adjust_hue(img, hue_rnd)
              img = transforms.functional.adjust_saturation(img, saturation_rnd)
  
  
          # CONVERTING TO TENSOR
          tensor_img = transforms.functional.to_tensor(img)
  
          # Handle grayscaled images
          if len(tensor_img.shape) != 3:
              tensor_img = tensor_img.unsqueeze(0)
              tensor_img = tensor_img.expand((3, img.shape[1:]))
  
          # !!!WARNING IN PIL IT'S WIDTH HEIGHT, WHEN IN PYTORCH IT IS HEIGHT WIDTH
  
          # Apply augmentations for train it would be mosaic
          if self.train:
              mossaic_img = torch.zeros(3, self.img_size, self.img_size)
  
              # FINDING CROSS POINT
              cross_x = int(random.uniform(self.img_size * self.cross_offset, self.img_size * (1 - self.cross_offset)))
              cross_y = int(random.uniform(self.img_size * self.cross_offset, self.img_size * (1 - self.cross_offset)))
  
              fragment_img, fragment_bbox = self.get_mosaic(0, cross_x, cross_y, tensor_img, boxes)
              mossaic_img[:, 0:cross_y, 0:cross_x] = fragment_img
              boxes = fragment_bbox
              
              for n in range(1, 4):
                  raw_fragment_img, raw_fragment_bbox = self.get_img_for_mosaic(brightness_rnd, contrast_rnd, hue_rnd, saturation_rnd)
                  fragment_img, fragment_bbox = self.get_mosaic(n, cross_x, cross_y, raw_fragment_img, raw_fragment_bbox)
                  boxes = torch.cat([boxes, fragment_bbox])
  
                  if n == 1:
                      mossaic_img[:, 0 : cross_y, cross_x : self.img_size] = fragment_img
                  elif n == 2:
                      mossaic_img[:, cross_y : self.img_size, 0 : cross_x] = fragment_img
                  elif n == 3:
                      mossaic_img[:, cross_y : self.img_size, cross_x : self.img_size] = fragment_img
  
              #Set mossaic to return tensor
              tensor_img = mossaic_img
  
  
          # For validation it would be letterbox
          else:
              xyxy_bboxes = utils.xywh2xyxy(boxes[:, 1:])
  
              #IMG
              padding = abs((t_width - t_height))//2
              padded_img = torch.zeros(3, self.img_size, self.img_size)
              if t_width > t_height:
                  padded_img[:, padding:padding+t_height] = tensor_img
              else:
                  padded_img[:, :, padding:padding+t_width] = tensor_img
  
              tensor_img = padded_img
              
              relative_padding = padding/self.img_size
              #BOXES
              if t_width > t_height:
                  #Change y's relative position
                  xyxy_bboxes[:, 1] *= ratio
                  xyxy_bboxes[:, 3] *= ratio
                  xyxy_bboxes[:, 1] += relative_padding
                  xyxy_bboxes[:, 3] += relative_padding
              else:#x's
                  xyxy_bboxes[:, 0] *= ratio
                  xyxy_bboxes[:, 2] *= ratio
                  xyxy_bboxes[:, 0] += relative_padding
                  xyxy_bboxes[:, 2] += relative_padding
  
              boxes[:, 1:] = utils.xyxy2xywh(xyxy_bboxes)
  
  
          
          
          targets = torch.zeros((len(boxes), 6))
          targets[:, 1:] = boxes
  
  
          return img_path, tensor_img, targets
  
      def get_img_for_mosaic(self, brightness_rnd, contrast_rnd, hue_rnd, saturation_rnd):
          random_index = random.randrange(0, len(self.img_files))
          img_path = self.img_files[random_index].rstrip()
          label_path = self.label_files[random_index].rstrip()
  
          
  
          # Getting image
          img = Image.open(img_path).convert('RGB')
          width, height = img.size
  
          if os.path.exists(label_path):
              boxes = torch.from_numpy(np.loadtxt(label_path).reshape(-1, 5))
  
          #RESIZING
          if width > height:
              ratio = height/width
              t_width = self.img_size
              t_height = int(ratio * self.img_size) 
              
          else:
              ratio = width/height
              t_width = int(ratio * self.img_size)
              t_height = self.img_size
  
          img = transforms.functional.resize(img, (t_height, t_width))
  
          img = transforms.functional.adjust_brightness(img, brightness_rnd)
          img = transforms.functional.adjust_contrast(img, contrast_rnd)
          img = transforms.functional.adjust_hue(img, hue_rnd)
          img = transforms.functional.adjust_saturation(img, saturation_rnd)
  
          #CONVERTING TO TENSOR
          tensor_img = transforms.functional.to_tensor(img)
  
          # Handle grayscaled images
          if len(tensor_img.shape) != 3:
              tensor_img = tensor_img.unsqueeze(0)
              tensor_img = tensor_img.expand((3, img.shape[1:]))
  
          return tensor_img, boxes
  
  
      # N is spatial parameter if 0 TOP LEFT, if 1 TOP RIGHT, if 2 BOTTOM LEFT, if 3 BOTTOM RIGHT
      def get_mosaic(self, n, cross_x, cross_y, tensor_img, boxes):
          t_height = tensor_img.shape[1]
          t_width = tensor_img.shape[2]
  
          xyxy_bboxes = utils.xywh2xyxy(boxes[:, 1:])
  
          relative_cross_x = cross_x / self.img_size
          relative_cross_y = cross_y / self.img_size
  
          #CALCULATING TARGET WIDTH AND HEIGHT OF PICTURE
          if n == 0:
              width_of_nth_pic = cross_x 
              height_of_nth_pic = cross_y
          elif n == 1:
              width_of_nth_pic = self.img_size - cross_x
              height_of_nth_pic = cross_y
          elif n == 2:
              width_of_nth_pic = cross_x
              height_of_nth_pic = self.img_size - cross_y
          elif n == 3:
              width_of_nth_pic = self.img_size - cross_x
              height_of_nth_pic = self.img_size - cross_y
  
          # self.img_size - width_of_1st_pic
          # selg.img_size - height_of_1st_pic  
  
  
          # CHOOSING TOP LEFT CORNER (doing offset to have more than fex pixels in bbox :-) )
          cut_x1 = random.randint(0, int(t_width * 0.33))
          cut_y1 = random.randint(0, int(t_height * 0.33))
  
  
          # Now we should find which axis should we randomly enlarge (this we do by finding out which ratio is bigger); cross x is basically width of the top left picture
          if (t_width - cut_x1) / width_of_nth_pic < (t_height - cut_y1) / height_of_nth_pic:
              cut_x2 = random.randint(cut_x1 + int(t_width * 0.67), t_width)
              cut_y2 = int(cut_y1 + (cut_x2-cut_x1)/width_of_nth_pic*height_of_nth_pic)
  
          else:
              cut_y2 = random.randint(cut_y1 + int(t_height * 0.67), t_height)
              cut_x2 = int(cut_x1 + (cut_y2-cut_y1)/height_of_nth_pic*width_of_nth_pic)
          
          # RESIZING AND INSERTING (TO DO 2D interpolation wants 4 dimensions, so I add and remove one by using None and squeeze)
          tensor_img = F.interpolate(tensor_img[:, cut_y1:cut_y2,  cut_x1:cut_x2][None], (height_of_nth_pic, width_of_nth_pic)).squeeze()
          
          # BBOX
          relative_cut_x1 = cut_x1 / t_width
          relative_cut_y1 = cut_y1 / t_height
          relative_cropped_width = (cut_x2 - cut_x1) / t_width
          relative_cropped_height = (cut_y2 - cut_y1) / t_height
  
          # SHIFTING TO CUTTED IMG SO X1 Y1 WILL 0
          xyxy_bboxes[:, 0] = xyxy_bboxes[:, 0] - relative_cut_x1
          xyxy_bboxes[:, 1] = xyxy_bboxes[:, 1] - relative_cut_y1
          xyxy_bboxes[:, 2] = xyxy_bboxes[:, 2] - relative_cut_x1
          xyxy_bboxes[:, 3] = xyxy_bboxes[:, 3] - relative_cut_y1
  
          # RESIZING TO CUTTED IMG SO X2 WILL BE 1
          xyxy_bboxes[:, 0] /= relative_cropped_width
          xyxy_bboxes[:, 1] /= relative_cropped_height
          xyxy_bboxes[:, 2] /= relative_cropped_width
          xyxy_bboxes[:, 3] /= relative_cropped_height
  
          # CLAMPING BOUNDING BOXES, SO THEY DO NOT OVERCOME OUTSIDE THE IMAGE
          xyxy_bboxes[:, 0].clamp_(0, 1)
          xyxy_bboxes[:, 1].clamp_(0, 1)
          xyxy_bboxes[:, 2].clamp_(0, 1)
          xyxy_bboxes[:, 3].clamp_(0, 1)
  
          # FILTER TO THROUGH OUT ALL SMALL BBOXES
          filter_minbbox = (xyxy_bboxes[:, 2] - xyxy_bboxes[:, 0] > self.bbox_minsize) & (xyxy_bboxes[:, 3] - xyxy_bboxes[:, 1] > self.bbox_minsize)
  
          # RESIZING TO MOSAIC
          if n == 0:
              xyxy_bboxes[:, 0] *= relative_cross_x #
              xyxy_bboxes[:, 1] *= relative_cross_y #(1 - relative_cross_y)
              xyxy_bboxes[:, 2] *= relative_cross_x #
              xyxy_bboxes[:, 3] *= relative_cross_y #(1 - relative_cross_y)
          elif n==1:
              xyxy_bboxes[:, 0] *= (1 - relative_cross_x) 
              xyxy_bboxes[:, 1] *= relative_cross_y
              xyxy_bboxes[:, 2] *= (1 - relative_cross_x)
              xyxy_bboxes[:, 3] *= relative_cross_y
          elif n==2:
              xyxy_bboxes[:, 0] *= relative_cross_x
              xyxy_bboxes[:, 1] *= (1 - relative_cross_y)
              xyxy_bboxes[:, 2] *= relative_cross_x
              xyxy_bboxes[:, 3] *= (1 - relative_cross_y)
          elif n==3:
              xyxy_bboxes[:, 0] *= (1 - relative_cross_x)
              xyxy_bboxes[:, 1] *= (1 - relative_cross_y)
              xyxy_bboxes[:, 2] *= (1 - relative_cross_x)
              xyxy_bboxes[:, 3] *= (1 - relative_cross_y)
  
          # RESIZING TO MOSAIC
          if n == 0:
              xyxy_bboxes[:, 0] = xyxy_bboxes[:, 0] # + relative_cross_x
              xyxy_bboxes[:, 1] = xyxy_bboxes[:, 1] # + relative_cross_y
              xyxy_bboxes[:, 2] = xyxy_bboxes[:, 2] # + relative_cross_x
              xyxy_bboxes[:, 3] = xyxy_bboxes[:, 3] # + relative_cross_y
          elif n==1:
              xyxy_bboxes[:, 0] = xyxy_bboxes[:, 0] + relative_cross_x
              xyxy_bboxes[:, 1] = xyxy_bboxes[:, 1]
              xyxy_bboxes[:, 2] = xyxy_bboxes[:, 2] + relative_cross_x
              xyxy_bboxes[:, 3] = xyxy_bboxes[:, 3]
          elif n==2:
              xyxy_bboxes[:, 0] = xyxy_bboxes[:, 0]
              xyxy_bboxes[:, 1] = xyxy_bboxes[:, 1] + relative_cross_y
              xyxy_bboxes[:, 2] = xyxy_bboxes[:, 2]
              xyxy_bboxes[:, 3] = xyxy_bboxes[:, 3] + relative_cross_y
          elif n==3:
              xyxy_bboxes[:, 0] = xyxy_bboxes[:, 0] + relative_cross_x
              xyxy_bboxes[:, 1] = xyxy_bboxes[:, 1] + relative_cross_y
              xyxy_bboxes[:, 2] = xyxy_bboxes[:, 2] + relative_cross_x
              xyxy_bboxes[:, 3] = xyxy_bboxes[:, 3] + relative_cross_y
          
          boxes = boxes[filter_minbbox]
          boxes[:, 1:] = utils.xyxy2xywh(xyxy_bboxes)[filter_minbbox]
  
          return tensor_img, boxes
  
      def collate_fn(self, batch):
          paths, imgs, targets = list(zip(*batch))
          # Remove empty placeholder targets
          targets = [boxes for boxes in targets if boxes is not None]
          # Add sample index to targets
          for i, boxes in enumerate(targets):
              boxes[:, 0] = i
          targets = torch.cat(targets, 0)
  
          return paths, torch.stack(imgs), targets
  
      def __len__(self):
          return len(self.img_files)