text: """
  Functions used to process and augment video data prior to passing into a model to train. 
  Additionally also processing all bounding boxes in a video according to the transformations performed on the video.
  
  Usage:
      In a custom dataset class:
      from preprocessing_transforms import *
  
  clip: Input to __call__ of each transform is a list of PIL images
  
  All functions have an example in the TestPreproc class at the bottom of this file
  """
  
  import torch
  import torchvision
  from torchvision.transforms import functional as F
  from PIL import Image
  from PIL import ImageChops
  import cv2
  from scipy import ndimage
  import numpy as np
  from abc import ABCMeta
  from math import floor, ceil
  
  class PreprocTransform(object):
      """
      Abstract class for preprocessing transforms that contains methods to convert clips to PIL images.
      """
      __metaclass__ = ABCMeta
      def __init__(self, **kwargs):
          pass
  
      def __init__(self, *args, **kwargs):
          pass
  
      def _to_pil(self, clip):
          # Must be of type uint8 if images have multiple channels, int16, int32, or float32 if there is only one channel
          if isinstance(clip[0], np.ndarray):
              if 'float' in str(clip[0].dtype):
                  clip = np.array(clip).astype('float32')
              if 'int64' == str(clip[0].dtype):
                  clip = np.array(clip).astype('int32')
              if clip[0].ndim == 3:
                  clip = np.array(clip).astype('uint8')
  
          output=[]
          for frame in clip:
              output.append(F.to_pil_image(frame))
          
          return output
  
  
      def _to_numpy(self, clip):
          output = []
          if isinstance(clip[0], torch.Tensor):
              if isinstance(clip, torch.Tensor):
                  output = clip.numpy()
              else:
                  for frame in clip:
                      f_shape = frame.shape
                      # Convert from torch's C, H, W to numpy H, W, C
                      frame = frame.numpy().reshape(f_shape[1], f_shape[2], f_shape[0])
  
                      output.append(frame)
              
  
          elif isinstance(clip[0], Image.Image):
              for frame in clip:
                  output.append(np.array(frame))
  
          else:
              output = clip 
  
          output = np.array(output)
  
          #if output.max() > 1.0:
          #    output = output/255.
  
          return output
  
  
      def _to_tensor(self, clip):
          """
          torchvision converts PIL images and numpy arrays that are uint8 0 to 255 to float 0 to 1
          Converts numpy arrays that are float to float tensor
          """
              
          if isinstance(clip[0], torch.Tensor):
              return clip
  
          output = []
          for frame in clip:
              output.append(F.to_tensor(frame))
          
          return output
  
  
  
  
  
  class ResizeClip(PreprocTransform):
      def __init__(self, *args, **kwargs):
          super(ResizeClip, self).__init__(*args, **kwargs)
          self.size_h, self.size_w = kwargs['resize_shape']
          
      def resize_bbox(self, xmin, ymin, xmax, ymax, img_shape, resize_shape):
          # Resize a bounding box within a frame relative to the amount that the frame was resized
      
          img_h = img_shape[0]
          img_w = img_shape[1]
      
          res_h = resize_shape[0]
          res_w = resize_shape[1]
      
          frac_h = res_h/float(img_h)
          frac_w = res_w/float(img_w)
      
          xmin_new = int(xmin * frac_w)
          xmax_new = int(xmax * frac_w)
      
          ymin_new = int(ymin * frac_h)
          ymax_new = int(ymax * frac_h)
      
          return xmin_new, ymin_new, xmax_new, ymax_new 
  
      def resize_pt_coords(self, x, y, img_shape, resize_shape):
          # Get relative position for point coords within a frame, after it's resized
  
          img_h = img_shape[0]
          img_w = img_shape[1]
      
          res_h = resize_shape[0]
          res_w = resize_shape[1]
      
          frac_h = res_h/float(img_h)
          frac_w = res_w/float(img_w)
  
          x_new = (x * frac_w).astype(int)
          y_new = (y * frac_h).astype(int)
  
          return x_new, y_new
  
      def __call__(self, clip, bbox=[]):
  
          clip = self._to_numpy(clip)
          out_clip = []
          out_bbox = []
          for frame_ind in range(len(clip)):
              frame = clip[frame_ind]
  
              proc_frame = cv2.resize(frame, (self.size_w, self.size_h))
              out_clip.append(proc_frame)
              if bbox!=[]:
                  temp_bbox = np.zeros(bbox[frame_ind].shape)-1 
                  for class_ind, box in enumerate(bbox[frame_ind]):
                      if np.array_equal(box,-1*np.ones(box.shape)): #only annotated objects
                          continue
  
                      if box.shape[-1] == 2: #Operate on point coordinates
                          proc_bbox = np.stack(self.resize_pt_coords(box[:,0], box[:,1], frame.shape, (self.size_h, self.size_w)),1)
                      else: #Operate on bounding box
                          xmin, ymin, xmax, ymax = box
                          proc_bbox = self.resize_bbox(xmin, ymin, xmax, ymax, frame.shape, (self.size_h, self.size_w))
  
                      temp_bbox[class_ind,:] = proc_bbox
                  out_bbox.append(temp_bbox)
  
          out_clip = np.array(out_clip)
  
          assert(out_clip.shape[1:3] == (self.size_h, self.size_w)), 'Proc frame: {} Crop h,w: {},{}'.format(out_clip.shape,self.size_h,self.size_w)
  
          if bbox!=[]:
              return out_clip, np.array(out_bbox)
          else:
              return out_clip
  
  
  class CropClip(PreprocTransform):
      def __init__(self, xmin=None, xmax=None, ymin=None, ymax=None, *args, **kwargs):
          super(CropClip, self).__init__(*args, **kwargs)
          self.crop_xmin = xmin
          self.crop_xmax = xmax
          self.crop_ymin = ymin
          self.crop_ymax = ymax
  
          self.crop_h, self.crop_w = kwargs['crop_shape']
  
  
      def _update_bbox(self, xmin, xmax, ymin, ymax, update_crop_shape=False):
          '''
              Args:
                  xmin (Float, shape []):
                  xmax (Float, shape []):
                  ymin (Float, shape []):
                  ymax (Float, shape []):
                  update_crop_shape (Boolean): Update expected crop shape along with bbox update call 
          '''
          self.crop_xmin = xmin
          self.crop_xmax = xmax
          self.crop_ymin = ymin
          self.crop_ymax = ymax
  
          if update_crop_shape:
              self.crop_h = ymax - ymin
              self.crop_w = xmax - xmin 
  
      def crop_bbox(self, xmin, ymin, xmax, ymax, crop_xmin, crop_ymin, crop_xmax, crop_ymax):
          if (xmin >= crop_xmax) or (xmax <= crop_xmin) or (ymin >= crop_ymax) or (ymax <= crop_ymin):
              return -1, -1, -1, -1
      
          if ymax > crop_ymax:
              ymax_new = crop_ymax
          else:
              ymax_new = ymax
      
          if xmax > crop_xmax:
              xmax_new = crop_xmax
          else:
              xmax_new = xmax
          
          if ymin < crop_ymin:
              ymin_new = crop_ymin
          else:
              ymin_new = ymin 
      
          if xmin < crop_xmin:
              xmin_new = crop_xmin
          else:
              xmin_new = xmin 
      
          return xmin_new-crop_xmin, ymin_new-crop_ymin, xmax_new-crop_xmin, ymax_new-crop_ymin
  
      def crop_coords(self, x, y, crop_xmin, crop_ymin, crop_xmax, crop_ymax):
          if np.any(x >= crop_xmax) or np.any(x <= crop_xmin) or np.any(y >= crop_ymax) or np.any(y <= crop_ymin):
              return -1*np.ones(x.shape), -1*np.ones(y.shape)
  
          x_new = np.clip(x, crop_xmin, crop_xmax)
          y_new = np.clip(y, crop_ymin, crop_ymax)
  
          return x_new-crop_xmin, y_new-crop_ymin 
    
      def __call__(self, clip, bbox=[]):
          out_clip = []
          out_bbox = []
  
          for frame_ind in range(len(clip)):
              frame = clip[frame_ind]
              proc_frame = np.array(frame[self.crop_ymin:self.crop_ymax, self.crop_xmin:self.crop_xmax]) 
              out_clip.append(proc_frame)
  
              assert(proc_frame.shape[:2] == (self.crop_h, self.crop_w)), 'Frame shape: {}, Proc frame: {} Crop h,w: {},{}'.format(frame.shape, proc_frame.shape,self.crop_h,self.crop_w)
  
              if bbox!=[]:
                  temp_bbox = np.zeros(bbox[frame_ind].shape)-1 
                  for class_ind, box in enumerate(bbox[frame_ind]):
                      if np.array_equal(box,-1*np.ones(box.shape)): #only annotated objects
                          continue
  
                      if box.shape[-1] == 2: #Operate on point coordinates
                          proc_bbox = np.stack(self.crop_coords(box[:,0], box[:,1], self.crop_xmin, self.crop_ymin, self.crop_xmax, self.crop_ymax), 1)
                      else: #Operate on bounding box
                          xmin, ymin, xmax, ymax = box
                          proc_bbox = self.crop_bbox(xmin, ymin, xmax, ymax, self.crop_xmin, self.crop_ymin, self.crop_xmax, self.crop_ymax)
                      temp_bbox[class_ind,:] = proc_bbox
                  out_bbox.append(temp_bbox)
  
          if bbox!=[]:
              return np.array(out_clip), np.array(out_bbox)
          else:
              return np.array(out_clip)
  
  
  class RandomCropClip(PreprocTransform):
      def __init__(self, *args, **kwargs):
          super(RandomCropClip, self).__init__(*args, **kwargs)
          self.crop_h, self.crop_w = kwargs['crop_shape']
  
          self.crop_transform = CropClip(0, 0, self.crop_w, self.crop_h, **kwargs)
  
          self.xmin = None
          self.xmax = None
          self.ymin = None
          self.ymax = None
  
  
      def _update_random_sample(self, frame_h, frame_w):
          if frame_w == self.crop_w:
              self.xmin = 0
          else:
              self.xmin = np.random.randint(0, frame_w-self.crop_w)   
  
          self.xmax = self.xmin + self.crop_w
  
          if frame_h == self.crop_h:
              self.ymin = 0
          else:
              self.ymin = np.random.randint(0, frame_h-self.crop_h)
          
          self.ymax = self.ymin + self.crop_h
  
      def get_random_sample(self):
          return self.xmin, self.xmax, self.ymin, self.ymax
          
      def __call__(self, clip, bbox=[]):
          frame_shape = clip[0].shape
          self._update_random_sample(frame_shape[0], frame_shape[1])
          self.crop_transform._update_bbox(self.xmin, self.xmax, self.ymin, self.ymax) 
          proc_clip = self.crop_transform(clip, bbox)
          if isinstance(proc_clip, tuple):
              assert(proc_clip[0].shape[1:3] == (self.crop_h, self.crop_w)), 'Proc frame: {} Crop h,w: {},{}'.format(proc_clip[0].shape,self.crop_h,self.crop_w)
          else:
              assert(proc_clip.shape[1:3] == (self.crop_h, self.crop_w)), 'Proc frame: {} Crop h,w: {},{}'.format(proc_clip.shape,self.crop_h,self.crop_w)
          return proc_clip 
  
  class CenterCropClip(PreprocTransform):
      def __init__(self, *args, **kwargs):
          super(CenterCropClip, self).__init__(*args, **kwargs)
          self.crop_h, self.crop_w = kwargs['crop_shape']
  
          self.crop_transform = CropClip(0, 0, self.crop_w, self.crop_h, **kwargs)
  
      def _calculate_center(self, frame_h, frame_w):
          xmin = int(frame_w/2 - self.crop_w/2)
          xmax = int(frame_w/2 + self.crop_w/2)
          ymin = int(frame_h/2 - self.crop_h/2)
          ymax = int(frame_h/2 + self.crop_h/2)
          return xmin, xmax, ymin, ymax
          
      def __call__(self, clip, bbox=[]):
          frame_shape = clip[0].shape
          xmin, xmax, ymin, ymax = self._calculate_center(frame_shape[0], frame_shape[1])
          self.crop_transform._update_bbox(xmin, xmax, ymin, ymax) 
          proc_clip =  self.crop_transform(clip, bbox)
          if isinstance(proc_clip, tuple):
              assert(proc_clip[0].shape[1:3] == (self.crop_h, self.crop_w)), 'Proc frame: {} Crop h,w: {},{}'.format(proc_clip[0].shape,self.crop_h,self.crop_w)
          else:
              assert(proc_clip.shape[1:3] == (self.crop_h, self.crop_w)), 'Proc frame: {} Crop h,w: {},{}'.format(proc_clip.shape,self.crop_h,self.crop_w)
          return proc_clip
  
  
  class RandomFlipClip(PreprocTransform):
      """   
      Specify a flip direction:
      Horizontal, left right flip = 'h' (Default)
      Vertical, top bottom flip = 'v'
      """
      def __init__(self, direction='h', p=0.5, *args, **kwargs):
          super(RandomFlipClip, self).__init__(*args, **kwargs)
          self.direction = direction
          self.p = p
  
      def _update_p(self, p):
          self.p = p
              
      def _random_flip(self):
          flip_prob = np.random.random()
          if flip_prob >= self.p:
              return 0
          else:
              return 1
  
      def _h_flip(self, bbox, frame_size):
          width = frame_size[1]
          bbox_shape = bbox.shape
          output_bbox = np.zeros(bbox_shape)-1
          for bbox_ind, box in enumerate(bbox):
              if np.array_equal(box,-1*np.ones(box.shape)): #only annotated objects
                  continue
  
              if box.shape[-1] == 2: #Operate on point coordinates
                  x = box[:,0]
                  x_new = width - x
  
                  output_bbox[bbox_ind] = np.stack((x_new,box[:,1]),1)
              else: #Operate on bounding box
                  xmin, ymin, xmax, ymax = box
                  xmax_new = width - xmin
                  xmin_new = width - xmax
                  output_bbox[bbox_ind] = xmin_new, ymin, xmax_new, ymax
          return output_bbox
  
      def _v_flip(self, bbox, frame_size):
          height = frame_size[0]
          bbox_shape = bbox.shape
          output_bbox = np.zeros(bbox_shape)-1
          for bbox_ind, box in enumerate(bbox):
              if np.array_equal(box,-1*np.ones(box.shape)): #only annotated objects
                  continue
  
              if box.shape[-1] == 2: #Operate on point coordinates
                  y = box[:,1]
                  y_new = height - y
  
                  output_bbox[bbox_ind] = np.stack((box[:,0],y_new),1)
              else: #Operate on bounding box
                  xmin, ymin, xmax, ymax = box
                  ymax_new = height - ymin
                  ymin_new = height - ymax
                  output_bbox[bbox_ind] = xmin, ymin_new, xmax, ymax_new
          return output_bbox
  
          bbox_shape = bbox.shape
          output_bbox = np.zeros(bbox_shape)-1
          for bbox_ind in range(bbox_shape[0]):
              xmin, ymin, xmax, ymax = bbox[bbox_ind] 
              height = frame_size[0]
              ymax_new = height - ymin 
              ymin_new = height - ymax
              output_bbox[bbox_ind] = xmin, ymin_new, xmax, ymax_new
          return output_bbox 
  
  
      def _flip_data(self, clip, bbox=[]):
          output_bbox = []
          
          if self.direction == 'h':
              output_clip = [cv2.flip(np.array(frame), 1) for frame in clip]
              
              if bbox!=[]:
                  output_bbox = [self._h_flip(frame, output_clip[0].shape) for frame in bbox] 
  
          elif self.direction == 'v':
              output_clip = [cv2.flip(np.array(frame), 0) for frame in clip]
  
              if bbox!=[]:
                  output_bbox = [self._v_flip(frame, output_clip[0].shape) for frame in bbox]
  
          return output_clip, output_bbox 
          
  
      def __call__(self, clip, bbox=[]):
          input_shape = np.array(clip).shape
          flip = self._random_flip()
          out_clip = clip
          out_bbox = bbox
          if flip:
              out_clip, out_bbox = self._flip_data(clip, bbox)
  
          out_clip = np.array(out_clip)
          assert(input_shape == out_clip.shape), "Input shape {}, output shape {}".format(input_shape, out_clip.shape)
  
          if bbox!=[]:
              return out_clip, out_bbox
          else:
              return out_clip
  
  class ToTensorClip(PreprocTransform):
      """
      Convert a list of PIL images or numpy arrays to a 5 dimensional pytorch tensor [batch, frame, channel, height, width]
      """
      def __init__(self, *args, **kwargs):
          super(ToTensorClip, self).__init__(*args, **kwargs)
  
          self.transform = torchvision.transforms.ToTensor()
  
      def __call__(self, clip, bbox=[]):
          
          if isinstance(clip[0], Image.Image):
              # a little round-about but it maintains consistency
              temp_clip = []
              for c in clip:
                  temp_clip.append(np.array(c))
              clip = temp_clip 
  
          output_clip = torch.from_numpy(np.array(clip)).float() #Numpy array to Tensor
  
          if bbox!=[]:
              bbox = torch.from_numpy(np.array(bbox))
              return output_clip, bbox
          else:
              return output_clip
          
  
  class RandomRotateClip(PreprocTransform):
      """
      Randomly rotate a clip from a fixed set of angles.
      The rotation is counterclockwise
      """
      def __init__(self,  angles=[0,90,180,270], *args, **kwargs):
          super(RandomRotateClip, self).__init__(*args, **kwargs)
          self.angles = angles
  
      ######
      # Code from: https://stackoverflow.com/questions/20924085/python-conversion-between-coordinates
      def _cart2pol(self, point):
          x,y = point
          rho = np.sqrt(x**2 + y**2)
          phi = np.arctan2(y, x)
          return(rho, phi)
      
      def _pol2cart(self, point):
          rho, phi = point
          x = rho * np.cos(phi)
          y = rho * np.sin(phi)
          return(x, y)
      #####
  
      def _update_angles(self, angles):
          self.angles=angles
  
  
      def _rotate_bbox(self, bboxes, frame_shape, angle):
          angle = np.deg2rad(angle)
          bboxes_shape = bboxes.shape
          output_bboxes = np.zeros(bboxes_shape)-1
          frame_h, frame_w = frame_shape[0], frame_shape[1] 
          half_h = frame_h/2. 
          half_w = frame_w/2. 
  
          for bbox_ind in range(bboxes_shape[0]):
              xmin, ymin, xmax, ymax = bboxes[bbox_ind]
              tl = (xmin-half_w, ymax-half_h)
              tr = (xmax-half_w, ymax-half_h)
              bl = (xmin-half_w, ymin-half_h)
              br = (xmax-half_w, ymin-half_h)
  
              tl = self._cart2pol(tl) 
              tr = self._cart2pol(tr)    
              bl = self._cart2pol(bl)
              br = self._cart2pol(br)
  
              tl = (tl[0], tl[1] - angle)
              tr = (tr[0], tr[1] - angle)
              bl = (bl[0], bl[1] - angle)
              br = (br[0], br[1] - angle)
  
              tl = self._pol2cart(tl) 
              tr = self._pol2cart(tr)    
              bl = self._pol2cart(bl)
              br = self._pol2cart(br)
  
              tl = (tl[0]+half_w, tl[1]+half_h)
              tr = (tr[0]+half_w, tr[1]+half_h)
              bl = (bl[0]+half_w, bl[1]+half_h)
              br = (br[0]+half_w, br[1]+half_h)
  
              xmin_new = int(np.clip(min(floor(tl[0]), floor(tr[0]), floor(bl[0]), floor(br[0])), 0, frame_w-1))
              xmax_new = int(np.clip(max(ceil(tl[0]), ceil(tr[0]), ceil(bl[0]), ceil(br[0])), 0, frame_w-1))
              ymin_new = int(np.clip(min(floor(tl[1]), floor(tr[1]), floor(bl[1]), floor(br[1])), 0, frame_h-1))
              ymax_new = int(np.clip(max(ceil(tl[1]), ceil(tr[1]), ceil(bl[1]), ceil(br[1])), 0, frame_h-1))
  
              output_bboxes[bbox_ind] = [xmin_new, ymin_new, xmax_new, ymax_new]
  
          return output_bboxes
  
  
      def _rotate_coords(self, bboxes, frame_shape, angle):
          angle = np.deg2rad(angle)
          bboxes_shape = bboxes.shape
          output_bboxes = np.zeros(bboxes_shape)-1
          frame_h, frame_w = frame_shape[0], frame_shape[1] 
          half_h = frame_h/2. 
          half_w = frame_w/2. 
  
          for bbox_ind in range(bboxes_shape[0]):
              x, y = bboxes[bbox_ind].transpose()
  
              pts  = (x-half_w, y-half_h)
  
              pts = self._cart2pol(pts)
  
              pts = (pts[0], pts[1]-angle)
  
              pts = self._pol2cart(pts)
  
              pts  = (pts[0]+half_w, pts[1]+half_h)
  
              output_bboxes[bbox_ind,:,0] = (np.clip(pts[0], 0, frame_w-1))
              output_bboxes[bbox_ind,:,1] = (np.clip(pts[1], 0, frame_h-1))
  
          return output_bboxes
  
      def __call__(self, clip, bbox=[]):
          angle = np.random.choice(self.angles)
          output_clip = []
          clip = self._to_numpy(clip)
          for frame in clip:
              output_clip.append(ndimage.rotate(frame, angle, reshape=False))
  
          if bbox!=[]:
              bbox = np.array(bbox)
              output_bboxes = np.zeros(bbox.shape)-1
              for bbox_ind in range(bbox.shape[0]):
                  if bbox.shape[-1] == 2:
                      output_bboxes[bbox_ind] = self._rotate_coords(bbox[bbox_ind], clip[0].shape, angle)
                  else:
                      output_bboxes[bbox_ind] = self._rotate_bbox(bbox[bbox_ind], clip[0].shape, angle)
  
              return output_clip, output_bboxes 
  
          return output_clip
  
  class RandomTranslateClip(PreprocTransform):
      """
      Random horizontal and/or vertical shift on frames in a clip. All frames receive same shifting 
      Shift will be bounded by object bounding box (if given). Meaning, object will always be in view
      Input numpy array must be of type np.uint8
  
      Args:
          - translate (Tuple)
              - max_x (float): maximum absolute fraction for horizontal shift 
              - max_y (float): maximum absolute fraction for vertical shift 
      """
      def __init__(self, translate, **kwargs):
          super(RandomTranslateClip, self).__init__(**kwargs)
  
          self.max_x, self.max_y = translate
  
          assert(self.max_x >= 0.0 and self.max_y >= 0.0)
          assert(self.max_x < 1.0  and self.max_y < 1.0) #Cannot shift past image bounds
  
      def _shift_frame(self, bbox, frame, tx, ty):
          M       = np.array([[1, 0, tx],[0, 1, ty]], dtype=np.float) # 2 x 3 transformation matrix
          out_frame = cv2.warpAffine(frame, M, (frame.shape[1], frame.shape[0]))
  
          if bbox is not None:
              bbox_h = np.reshape(bbox, (-1,2)) #x-y coords
              bbox_h = np.concatenate((bbox_h, np.ones((bbox_h.shape[0],1))), axis=1).transpose() #homography coords
  
              out_box = M @ bbox_h
  
              if bbox.shape[-1] == 2: #Operate on point coordinates
                  out_box = np.reshape(out_box.transpose(), (bbox.shape[0], bbox.shape[1],2))
              else: #Operate on bounding box
                  out_box = np.reshape(out_box.transpose(), (-1,4))
  
              return out_frame, out_box 
          else:
              return out_frame 
  
      def __call__(self, clip, bbox=[]):
          out_clip = []
          clip = self._to_numpy(clip)
  
          frac_x = np.random.rand()*(2*self.max_x)-self.max_x 
          frac_y = np.random.rand()*(2*self.max_y)-self.max_y  
  
          if bbox != []:
              out_bbox = []
              
              for frame, box in zip(clip,bbox):
                  img_h, img_w = frame.shape[:2] 
                  tx = int(img_w * frac_x)
                  ty = int(img_h * frac_y) 
  
                  #Bound translation amount so all objects remain in scene
                  if box.shape[-1] == 2: #Operate on point coordinates
                      mask = box[:,:,0] != -1
                      tx = np.clip(tx, np.max(-1*box[mask,0]), np.min(img_w-box[mask,0]))
                      ty = np.clip(ty, np.max(-1*box[mask,1]), np.min(img_h-box[mask,1]))
                      out_frame, out_box = self._shift_frame(box, frame, tx, ty)
                      out_box[~mask] = -1*np.ones(2)
  
                  else: #Operate on bounding box 
                      #bbox is bounding box object
                      mask = box[:,0] != -1
                      tx = np.clip(tx, np.max(-1*box[mask,0]), np.min(img_w-box[mask,2]))
                      ty = np.clip(ty, np.max(-1*box[mask,1]), np.min(img_h-box[mask,3]))
                      out_frame, out_box = self._shift_frame(box, frame, tx, ty)
                      out_box[~mask] = -1*np.ones(4)
  
                  out_clip.append(out_frame)
                  out_bbox.append(out_box)
  
              return out_clip, out_bbox 
          else:
              for frame in clip:
                  img_h, img_w = frame.shape[:2] 
                  tx = int(img_w * frac_x)
                  ty = int(img_h * frac_y) 
  
                  out_clip.append(self._shift_frame(None, frame, tx, ty))
  
              return out_clip 
  
  class RandomZoomClip(PreprocTransform):
      """
      Random zoom on all frames in a clip. All frames receive same scaling
      Scale will be bounded by object bounding box (if given). Meaning, object will always be in view
      If zooming out, the borders will be filled with black.
  
      >1: Zoom in
      <1: Zoom out
      =1: Same size
  
      Args:
          - scale (Tuple)
              - min_scale (float): minimum scaling on frame 
              - max_scale (float): maximum scaling on frame  
      """
      def __init__(self, scale, **kwargs):
          super(RandomZoomClip, self).__init__(**kwargs)
  
          self.min_scale, self.max_scale = scale
  
          assert(self.min_scale > 0 and self.min_scale <= self.max_scale)
  
      def _scale_frame(self, bbox, frame, sc):
          M = cv2.getRotationMatrix2D((frame.shape[1]/2, frame.shape[0]/2), 0, sc) # 2 x 3 rotation matrix
          out_frame = cv2.warpAffine(frame, M, (frame.shape[1], frame.shape[0]))
  
          if bbox is not None:
              bbox_h = np.reshape(bbox, (-1,2)) #x-y coords
              bbox_h = np.concatenate((bbox_h, np.ones((bbox_h.shape[0],1))), axis=1).transpose() #homography coords
  
              out_box = M @ bbox_h
  
              if bbox.shape[-1] == 2: #Operate on point coordinates
                  out_box = np.reshape(out_box.transpose(), (bbox.shape[0], bbox.shape[1],2))
              else: #Operate on bounding box
                  out_box = np.reshape(out_box.transpose(), (-1,4))
  
              return out_frame, out_box 
          else:
              return out_frame 
  
      def __call__(self, clip, bbox=[]):
          out_clip = []
          clip = self._to_numpy(clip)
  
          sc = np.random.uniform(self.min_scale, self.max_scale) 
  
          if bbox != []:
              out_bbox = []
              
              for frame, box in zip(clip,bbox):
                  img_h, img_w = frame.shape[:2]
                  cx, cy = (img_w/2, img_h/2)
  
                  #Bound scaling so all objects remain in scene
                  if box.shape[-1] == 2: #Operate on point coordinates
                      mask = box[:,:,0] != -1
  
                      max_x = min(img_w, np.max(cx + sc * (box[mask,0] - cx)))
                      min_x = max(0, np.min(cx + sc * (box[mask,0] - cx)))
                      sx = (max_x - cx) / np.max(box[mask,0] - cx)
                      if min_x == 0:
                          sx = min(sx, (min_x - cx) / np.min(box[mask,0] - cx))
  
                      max_y = min(img_h, np.max(cy + sc * (box[mask,1] - cy)))
                      min_y = max(0, np.min(cy + sc * (box[mask,1] - cy)))
                      sy = (max_y - cy) / np.max(box[mask,1] - cy)
                      if min_y == 0:
                          sy = min(sy, (min_y - cy) / np.min(box[mask,1] - cy))
              
                      sc = min(sx, sy)
                      out_frame, out_box = self._scale_frame(box, frame, sc)
                      out_box[~mask] = -1*np.ones(2)
  
                  else: #Operate on bounding box 
                      mask = box[:,0] != -1
  
                      max_x = min(img_w, np.max(cx + sc * (box[mask,2] - cx)))
                      min_x = max(0, np.min(cx + sc * (box[mask,0] - cx)))
                      sx = (max_x - cx) / np.max(box[mask,2] - cx)
                      if min_x == 0:
                          sx = min(sx, (min_x - cx) / np.min(box[mask,0] - cx))
  
                      max_y = min(img_h, np.max(cy + sc * (box[mask,3] - cy)))
                      min_y = max(0, np.min(cy + sc * (box[mask,1] - cy)))
                      sy = (max_y - cy) / np.max(box[mask,3] - cy)
                      if min_y == 0:
                          sy = min(sy, (min_y - cy) / np.min(box[mask,1] - cy))
              
                      sc = min(sx, sy)
                      out_frame, out_box = self._scale_frame(box, frame, sc)
                      out_box[~mask] = -1*np.ones(4)
  
                  out_clip.append(out_frame)
                  out_bbox.append(out_box)
  
              return out_clip, out_bbox 
          else:
              for frame in clip:
                  img_h, img_w = frame.shape[:2]
                  sx = int(img_w * sc)
                  sy = int(img_h * sc) 
  
                  out_clip.append(self._scale_frame(None, frame, sc))
              return out_clip 
  
  class SubtractMeanClip(PreprocTransform):
      def __init__(self, **kwargs):
          super(SubtractMeanClip, self).__init__(**kwargs)
  #        self.clip_mean = torch.tensor(kwargs['clip_mean']).float()
          self.clip_mean = kwargs['clip_mean']
  #        self.clip_mean      = []
  #
  #        for frame in self.clip_mean_args:
  #            self.clip_mean.append(Image.fromarray(frame))
  
          
      def __call__(self, clip, bbox=[]):
          #clip = clip-self.clip_mean
          for clip_ind in range(len(clip)):
              clip[clip_ind] = clip[clip_ind] - self.clip_mean[clip_ind]
  
          
          if bbox!=[]:
              return clip, bbox
  
          else:
              return clip
  
  class SubtractRGBMean(PreprocTransform):
      def __init__(self, **kwargs):
          super(SubtractRGBMean, self).__init__(**kwargs)
          self.rgb_mean = kwargs['subtract_mean']
      
      def __call__(self, clip, bbox=[]):
  
          clip = self._to_numpy(clip)
          out_clip = []
          out_bbox = []
  
          for frame_ind in range(len(clip)):
              frame = clip[frame_ind]
  
              proc_frame = frame - self.rgb_mean
              out_clip.append(proc_frame)
  
          if bbox != []:
              return out_clip, bbox
          else:
              return out_clip
  
  class ApplyToPIL(PreprocTransform):
      """
      Apply standard pytorch transforms that require PIL images as input to their __call__ function, for example Resize
  
      NOTE: The __call__ function of this class converts the clip to a list of PIL images in the form of integers from 0-255. If the clips are floats (for example afer mean subtraction), then only call this transform before the float transform
  
      Bounding box coordinates are not guaranteed to be transformed properly!
  
      https://pytorch.org/docs/stable/_modules/torchvision/transforms/transforms.html
      """
      def __init__(self, **kwargs):
          """
          class_kwargs is a dictionary that contains the keyword arguments to be passed to the chosen pytorch transform class
          """
          super(ApplyToPIL, self).__init__( **kwargs)
          self.kwargs = kwargs
          self.class_kwargs = kwargs['class_kwargs']
          self.transform = kwargs['transform'](**self.class_kwargs)
  
      def __call__(self, clip, bbox=[]):
          input_pil = True
          output_clip = []
  
          if not isinstance(clip[0], Image.Image):
              clip = self._to_pil(clip)
              clip = [frame.convert('RGB') for frame in clip]
              input_pil = False
  
          if input_pil:
              for frame in clip:
                  transformed_frame = self.transform(frame)
                  if isinstance(transformed_frame, tuple) or isinstance(transformed_frame, list):
                      for tf in transformed_frame:
                          output_clip.append(tf)
                  else:
                      output_clip.append(self.transform(frame)) #Apply transform and convert back to Numpy
  
          else:
              for frame in clip:
                  transformed_frame = self.transform(frame)
                  if isinstance(transformed_frame, tuple) or isinstance(transformed_frame, list):
                      for tf in transformed_frame:
                          output_clip.append(np.array(tf))
                  else:
                      output_clip.append(np.array(self.transform(frame))) #Apply transform and convert back to Numpy
  
          if bbox!=[]:
              return output_clip, bbox
  
          else:
              return output_clip
  
  
  class ApplyToTensor(PreprocTransform):
      """
      Apply standard pytorch transforms that require pytorch Tensors as input to their __call__ function, for example Normalize 
  
      NOTE: The __call__ function of this class converts the clip to a pytorch float tensor. If other transforms require PIL inputs, call them prior tho this one
      Bounding box coordinates are not guaranteed to be transformed properly!
  
      https://pytorch.org/docs/stable/_modules/torchvision/transforms/transforms.html
      """
      def __init__(self, **kwargs):
          super(ApplyToTensor, self).__init__(**kwargs)
          self.kwargs = kwargs
          self.class_kwargs = kwargs['class_kwargs']
          self.transform = kwargs['transform'](**self.class_kwargs)
  
      def __call__(self, clip, bbox=[]):
          if not isinstance(clip, torch.Tensor):
              clip = self._to_tensor(clip)
  
          output_clip = []
          for frame in clip:
              output_clip.append(self.transform(frame))
  
          output_clip = torch.stack(output_clip)
  
          if bbox!=[]:
              return output_clip, bbox
  
          else:
              return output_clip
  
  class ApplyOpenCV(PreprocTransform):
      """
      Apply opencv transforms that require numpy arrays as input to their __call__ function, for example Rotate 
  
      NOTE: The __call__ function of this class converts the clip to a Numpy array. If other transforms require PIL inputs, call them prior tho this one
  
      Bounding box coordinates are not guaranteed to be transformed properly!
      """
      def __init__(self, **kwargs):
          super(ApplyOpenCV, self).__init__(**kwargs)
          self.kwargs = kwargs
          self.class_kwargs = kwargs['class_kwargs']
          self.transform = kwargs['transform']
  
      def __call__(self, clip, bbox=[]):
          if not isinstance(clip, torch.Tensor):
              clip = self._to_numpy(clip)
  
          output_clip = []
          for frame in clip:
              output_clip.append(self.transform(frame, **self.class_kwargs))
  
  
          if bbox!=[]:
              return output_clip, bbox
  
          else:
              return output_clip
  
  
  
  
  class TestPreproc(object):
      def __init__(self):
          self.resize = ResizeClip(resize_shape = [2,2])
          self.crop = CropClip(0,0,0,0, crop_shape=[2,2])
          self.rand_crop = RandomCropClip(crop_shape=[2,2])
          self.cent_crop = CenterCropClip(crop_shape=[2,2])
          self.rand_flip_h = RandomFlipClip(direction='h', p=1.0)
          self.rand_flip_v = RandomFlipClip(direction='v', p=1.0)
          self.rand_rot = RandomRotateClip(angles=[90])
          self.rand_trans = RandomTranslateClip(translate=(0.5,0.5))
          self.rand_zoom  = RandomZoomClip(scale=(1.25,1.25)) 
          self.sub_mean = SubtractMeanClip(clip_mean=np.zeros(1))
          self.applypil = ApplyToPIL(transform=torchvision.transforms.ColorJitter, class_kwargs=dict(brightness=1))
          self.applypil2 = ApplyToPIL(transform=torchvision.transforms.FiveCrop, class_kwargs=dict(size=(64,64)))
          self.applytensor = ApplyToTensor(transform=torchvision.transforms.Normalize, class_kwargs=dict(mean=torch.tensor([0.,0.,0.]), std=torch.tensor([1.,1.,1.])))
          self.applycv = ApplyOpenCV(transform=cv2.threshold, class_kwargs=dict(thresh=100, maxval=100, type=cv2.THRESH_TRUNC))
          self.preproc = PreprocTransform()
  
      def resize_test(self):
          inp = np.array([[[.1,.2,.3,.4],[.1,.2,.3,.4],[.1,.2,.3,.4]]]).astype(float)
          inp2 = np.array([[[.1,.1,.1,.1],[.2,.2,.2,.2],[.3,.3,.3,.3]]]).astype(float)
          expected_out = np.array([[[.15,.35],[.15,.35]]]).astype(float)
          expected_out2 = np.array([[[.125,.125],[.275,.275]]]).astype(float)
          out = self.resize(inp)
          out2 = self.resize(inp2)
          assert (False not in np.isclose(out,expected_out)) and (False not in np.isclose(out2,expected_out2))
          
          bbox = np.array([[[0,0,3,3]]]).astype(float)
          _, bbox_out = self.resize(inp, bbox)
          exp_bbox = np.array([[[0,0,1,2]]])
          assert (False not in np.isclose(bbox_out, exp_bbox))
  
          coord_pts = np.array([[[[1,1], [7,5], [9,6]]]]).astype(float)
          _, bbox_out = self.resize(inp, coord_pts)
          exp_bbox = np.array([[[[0., 0.],
                                 [3., 3.],
                                 [4., 4.]]]])
          assert (False not in np.isclose(bbox_out, exp_bbox))
  
  
      def crop_test(self):
          inp = np.array([[[.1,.2,.3],[.4,.5,.6],[.7,.8,.9]]]).astype(float)
          self.crop._update_bbox(1, 3, 1, 3)
  
          exp_out = np.array([[[.5,.6],[.8,.9]]]).astype(float)
          out = self.crop(inp)
          assert (False not in np.isclose(out,exp_out))
  
      def cent_crop_test(self):
          inp = np.array([[[.1,.2,.3,.4],[.1,.2,.3,.4],[.1,.2,.3,.4],[.1,.2,.3,.4]]]).astype(float)
          exp_out = np.array([[[.2,.3],[.2,.3]]]).astype(float)
          out = self.cent_crop(inp)
          assert (False not in np.isclose(out, exp_out))
  
      def rand_crop_test(self):
          inp = np.array([[[.1,.2,.3,.4],[.3,.4,.5,.6],[.2,.3,.4,.5],[.1,.2,.3,.4]]]).astype(float)
          out = self.rand_crop(inp)
          coords = self.rand_crop.get_random_sample()
          exp_out = np.array(inp[:, coords[2]:coords[3], coords[0]:coords[1]]).astype(float)
          assert (False not in np.isclose(out, exp_out))
  
      def rand_flip_test(self):
          inp = np.array([[[.1,.2,.3],[.4,.5,.6],[.7,.8,.9]]]).astype(float)
          exp_outh = np.array([[[.3,.2,.1],[.6,.5,.4],[.9,.8,.7]]]).astype(float)
          exp_outv = np.array([[[.7,.8,.9],[.4,.5,.6],[.1,.2,.3]]]).astype(float)
          outh = self.rand_flip_h(inp)
          outv = self.rand_flip_v(inp)       
          assert (False not in np.isclose(outh,exp_outh)) and (False not in np.isclose(outv,exp_outv))
  
  
          inp2 = np.arange(36).reshape(6,6)
          bbox = np.array([[[0,0,2,2]]]).astype(float)
          exp_bboxh = np.array([[[4,0,6,2]]]).astype(float)
          exp_bboxv = np.array([[[0,4,2,6]]]).astype(float)
          _, bboxh = self.rand_flip_h([inp2], bbox)
          _, bboxv = self.rand_flip_v([inp2], bbox)       
           
          assert (False not in np.isclose(bboxh, exp_bboxh)) and (False not in np.isclose(bboxv, exp_bboxv))
  
      def rand_flip_vis(self):
          import matplotlib.pyplot as plt
          x = np.arange(112*112).reshape(112,112)
          x[:, 10] = 10000
          x[:, 50] = 5000
          x[10, :] = 5000
          x[50, :] = 10000
  
          plt.subplot(1,3,1); plt.imshow(x); plt.title('Original image')
          h = self.rand_flip_h([x])
          plt.subplot(1,3,2); plt.imshow(h[0]); plt.title('Flip Horizontal')
          v = self.rand_flip_v([x])
          plt.subplot(1,3,3); plt.imshow(v[0]); plt.title('Flip Vertical')
          
          plt.show()
  
      def rand_rot_test(self):
          inp = np.array([[[.1,.2,.3],[.4,.5,.6],[.7,.8,.9]]]).astype(float)
          exp_out = np.array([[[.3,.6,.9],[.2,.5,.8],[.1,.4,.7]]]).astype(float)
  
          out = self.rand_rot(inp)
  
  
          self.rand_rot._update_angles([45])
          inp2 = np.arange(6*6).reshape(6,6)
          bbox = [[2,2,4,4]]
          exp_bbox = [1,1,5,5]
          out_bbox = self.rand_rot([inp2], np.array([bbox]))[1][0].tolist()
          assert (False not in np.isclose(out, exp_out)) and (False not in np.isclose(exp_bbox, out_bbox))
  
  
      def rand_trans_test(self):
          x = np.arange(112*112).reshape(112,112).astype(np.uint8)
          out = self.rand_trans([x])
          out2 = self.rand_trans([x], bbox=[np.array([[32,32,96,96]])])
  
          assert (out2[1][0].min() >= 0) and (out[0].shape==(112,112)) and (out2[0][0].shape==(112,112))
  
      def rand_rot_vis(self):
          import matplotlib.pyplot as plt
          import matplotlib.patches as patches 
          angle = 45
          self.rand_rot._update_angles([angle])
          x = np.arange(112*112).reshape(112,112)
  
          bbox = [30,40,50,100]
          pts = np.array([[30,40],[30,80]])
          fig = plt.figure()
          ax1 = fig.add_subplot(121)
          x[bbox[1]:bbox[3], bbox[0]] = 0
          x[bbox[1]:bbox[3], bbox[2]-1] = 0
          x[bbox[1], bbox[0]:bbox[2]] = 0
          x[bbox[3]-1, bbox[0]:bbox[2]] = 0
          
          ax1.imshow(x); ax1.set_title('Original image')
          rect = patches.Rectangle((bbox[0],bbox[1]), bbox[2]-bbox[0],\
                                    bbox[3]-bbox[1], linewidth=1, edgecolor='k', facecolor='none')
          #ax1.add_patch(rect)
          ax1.scatter(pts[:,0], pts[:,1], c='r')
  
          out2 = self.rand_rot([x], np.array([[bbox]]))
          x_rot = out2[0][0]
          bbox_rot = out2[1][0,0]
  
          out2 = self.rand_rot([x], np.array([[pts]]))
          pts_rot  = out2[1][0,0]
  
          ax2 = fig.add_subplot(122)
          rect = patches.Rectangle((bbox_rot[0],bbox_rot[1]), bbox_rot[2]-bbox_rot[0],\
                                    bbox_rot[3]-bbox_rot[1], linewidth=1, edgecolor='k', facecolor='none')
          ax2.add_patch(rect)
          ax2.imshow(x_rot); ax2.set_title('Rotation: {} degress'.format(angle))
          ax2.scatter(pts_rot[:,0],pts_rot[:,1], c='r')
          plt.show()
  
      def rand_zoom_test(self):
          inp = np.array([[[.1,.2,.3],[.4,.5,.6],[.7,.8,.9]]]).astype(float)
          exp_out = np.array([[0.225   , 0.303125, 0.384375],
                              [0.459375, 0.5375  , 0.61875 ],
                              [0.703125, 0.78125 , 0.8625  ]]).astype(float)
          out = self.rand_zoom(inp)
  
          inp2 = np.arange(6*6, dtype=np.uint8).reshape(6,6)
          bbox = [[2,2,4,4]]
          exp_bbox = [1.75,1.75,4.25,4.25]
          _,out_bbox = self.rand_zoom([inp2], np.array([bbox]))
  
          assert (False not in np.isclose(out, exp_out)) and (False not in np.isclose(exp_bbox, out_bbox))
  
      def rand_zoom_vis(self):
          import matplotlib.pyplot as plt
          import matplotlib.patches as patches 
          x = np.arange(112*112, dtype=np.uint8).reshape(112,112)
  
          bbox = [30,40,50,100]
          pts = np.array([[30,40],[30,80]])
          fig = plt.figure()
          ax1 = fig.add_subplot(121)
  
          x[bbox[1]:bbox[3], bbox[0]] = 0
          x[bbox[1]:bbox[3], bbox[2]-1] = 0
          x[bbox[1], bbox[0]:bbox[2]] = 0
          x[bbox[3]-1, bbox[0]:bbox[2]] = 0
          ax1.imshow(x); ax1.set_title('Original image')
          ax1.scatter(pts[:,0], pts[:,1], c='r')
  
          out = self.rand_zoom([x], np.array([[pts]]))
          pts_zoom = out[1][0][0]
  
          out = self.rand_zoom([x], np.array([[bbox]]))
          x_zoom = out[0][0]
          bbox_zoom = out[1][0][0]
  
          ax2 = fig.add_subplot(122)
          rect = patches.Rectangle((bbox_zoom[0],bbox_zoom[1]), bbox_zoom[2]-bbox_zoom[0],\
                                    bbox_zoom[3]-bbox_zoom[1], linewidth=1, edgecolor='k', facecolor='none')
          ax2.add_patch(rect)
          ax2.imshow(x_zoom); ax2.set_title('Zoomed image')
          ax2.scatter(pts_zoom[:,0],pts_zoom[:,1], c='r')
          
          plt.show()
  
      def applypil_test(self):
          inp = np.arange(112*112).reshape(112,112)
          np_inp = [inp, inp]
          inp = self.applypil._to_pil([inp, inp])
          inp = [inp[0].convert('RGB'), inp[1].convert('RGB')]
          out = self.applypil(inp)
          out2 = self.applypil2(out)
          out3 = self.applypil(np_inp)
          assert (len(out2)==2*5) and (out2[0].size==(64,64)) and (isinstance(out2[0], Image.Image)) and (isinstance(out3[0], np.ndarray))
  
      def applytensor_test(self):
          inp = np.arange(112*112*3).reshape(3,112,112).astype('float32')
          inp = torch.from_numpy(inp)
          out = self.applytensor([inp, inp])
          assert False not in np.array(inp==out)
  
  
      def applycv_test(self):
          inp = np.arange(112*112).reshape(112,112).astype('float32')
          out = self.applycv([inp])
          assert (out[0][1].min()==0.0) and (out[0][1].max()==100.0) 
  
  
      def to_numpy_test(self):
          inp_torch = [torch.zeros((3,112,112))]
          inp_pil = [Image.fromarray(np.zeros((112,112,3)).astype('uint8'))]
          out_torch = self.preproc._to_numpy(inp_torch)
          out_pil = self.preproc._to_numpy(inp_pil)
  
          assert (False not in np.array(out_pil==out_torch))
  
  
      def to_tensor_test(self):
          inp_np_f = np.zeros((1,112,112,3)).astype('float')+201
          inp_np = np.zeros((1,112,112,3)).astype('uint8')+1
          inp_pil = [Image.fromarray(inp_np[0], mode='RGB')]
          out_np_f = self.preproc._to_tensor(inp_np_f)
          out_np = self.preproc._to_tensor(inp_np)
          out_pil = self.preproc._to_tensor(inp_pil)
          assert (False not in np.array(out_np[0]==out_pil[0])) and isinstance(out_np_f[0], torch.DoubleTensor)
  
  
      def to_pil_test(self):
          inp_np = [np.zeros((112,112,3)).astype('int32')]
          inp_torch = [torch.zeros((3,112,112))]
          out_np = self.preproc._to_pil(inp_np)
          out_torch = self.preproc._to_pil(inp_torch)
  
          assert (False not in np.array(np.array(out_np[0])==np.array(out_torch[0])))
  
  
      def run_tests(self):
          self.resize_test()
          self.crop_test()
          self.cent_crop_test()
          self.rand_crop_test()
          self.rand_flip_test()
          self.rand_rot_test()
          self.rand_trans_test()
          self.rand_zoom_test()
          self.applypil_test()
          self.applytensor_test()
          self.applycv_test()
          self.to_tensor_test()
          self.to_pil_test()
          self.to_numpy_test()
          print("Tests passed")
  
          self.rand_flip_vis()
          self.rand_rot_vis()
          self.rand_zoom_vis()
          
  if __name__=='__main__':
      test = TestPreproc()
      test.run_tests()
  
  
  