code: 
  
  
  
  
  
  '''simple docstring'''
  
  from typing import SupportsIndex
  
  import numpy as np
  from scipy.ndimage import convolve
  
  
  
  
  
  def 		UpperCamelCase       (							_lowerCamelCase      :      np.ndarray		, _lowerCamelCase      :      np.ndarray		, _lowerCamelCase      :      np.ndarray ):
      A__       =    np.stack((horizontal_flow, vertical_flow)		, 2 )
  
      # Create a grid of all pixel coordinates and subtract the flow to get the
      # target pixels coordinates
      A__       =    np.stack(
          np.meshgrid(np.arange(0		, image.shape[1] )		, np.arange(0		, image.shape[0] ) )		, 2 )
      A__       =    np.round(grid - flow ).astype(np.intaa )
  
      # Find the locations outside of the original image
      A__       =    (grid < 0) | (grid >= np.array([image.shape[1], image.shape[0]] ))
      A__       =    0
  
      A__       =    image[grid[:, :, 1], grid[:, :, 0]]
  
      # Set pixels at invalid locations to 0
      A__       =    0
  
      return warped
  
  
  
  
  
  def 		UpperCamelCase       (							_lowerCamelCase      :      np.ndarray		, _lowerCamelCase      :      np.ndarray		, _lowerCamelCase      :      SupportsIndex		, _lowerCamelCase      :      float | None = None		, ):
      if alpha is None:
          A__       =    0.1
  
      # Initialize flow
      A__       =    np.zeros_like(_lowerCamelCase )
      A__       =    np.zeros_like(_lowerCamelCase )
  
      # Prepare kernels for the calculation of the derivatives and the average velocity
      A__       =    np.array([[-1, 1], [-1, 1]] ) * 0.2_5
      A__       =    np.array([[-1, -1], [1, 1]] ) * 0.2_5
      A__       =    np.array([[1, 1], [1, 1]] ) * 0.2_5
      A__       =    np.array(
          [[1 / 12, 1 / 6, 1 / 12], [1 / 6, 0, 1 / 6], [1 / 12, 1 / 6, 1 / 12]] )
  
      # Iteratively refine the flow
      for _ in range(_lowerCamelCase ):
          A__       =    warp(_lowerCamelCase		, _lowerCamelCase		, _lowerCamelCase )
          A__       =    convolve(_lowerCamelCase		, _lowerCamelCase ) + convolve(_lowerCamelCase		, _lowerCamelCase )
          A__       =    convolve(_lowerCamelCase		, _lowerCamelCase ) + convolve(_lowerCamelCase		, _lowerCamelCase )
          A__       =    convolve(_lowerCamelCase		, _lowerCamelCase ) + convolve(_lowerCamelCase		, -kernel_t )
  
          A__       =    convolve(_lowerCamelCase		, _lowerCamelCase )
          A__       =    convolve(_lowerCamelCase		, _lowerCamelCase )
  
          # This updates the flow as proposed in the paper (Step 12)
          A__       =    (
              derivative_x * avg_horizontal_velocity
              + derivative_y * avg_vertical_velocity
              + derivative_t
          )
          A__       =    update / (alpha**2 + derivative_x**2 + derivative_y**2)
  
          A__       =    avg_horizontal_velocity - derivative_x * update
          A__       =    avg_vertical_velocity - derivative_y * update
  
      return horizontal_flow, vertical_flow
  
  
  if __name__ == "__main__":
    import doctest
  
    doctest.testmod()
  
  
  
  
style_context: 
  
  
  
  
  
  '''simple docstring'''
  
  import bisect
  import copy
  from collections import defaultdict
  
  import numpy as np
  from torch.utils.data import BatchSampler, Sampler
  
  from utils import logger
  
  
  
  
  
  def 		UpperCamelCase       (							_lowerCamelCase      :      str		, _lowerCamelCase      :      Optional[Any] ):
      A__       =    copy.deepcopy(_lowerCamelCase )
      A__       =    sorted(_lowerCamelCase )
      A__       =    [bisect.bisect_right(_lowerCamelCase		, _lowerCamelCase ) for y in x]
      return quantized
  
  
  
  
  
  def 		UpperCamelCase       (							_lowerCamelCase      :      Optional[int]		, _lowerCamelCase      :      Tuple=0 ):
      A__       =    np.arange(start=3		, stop=_lowerCamelCase		, step=4 ).tolist() if k > 0 else [10]
      A__       =    _quantize(_lowerCamelCase		, _lowerCamelCase )
      # count number of elements per group
      A__       =    np.unique(_lowerCamelCase		, return_counts=_lowerCamelCase )[1]
      A__       =    [0] + bins + [np.inf]
      logger.info("Using {} as bins for aspect lengths quantization".format(_lowerCamelCase ) )
      logger.info("Count of instances per bin: {}".format(_lowerCamelCase ) )
      return groups
  
  
  
  
  
  
  class      UpperCAmelCase			( UpperCamelCase__       ):
   def __init__(       self   :Optional[Any]		, lowercase_   :List[Any]		, lowercase_   :Optional[int]		, lowercase_   :str							)->							Dict:
       if not isinstance(lowercase_		, lowercase_							):
           raise ValueError(
               "sampler should be an instance of torch.utils.data.Sampler, but got sampler={}".format(lowercase_							)							)
       A__       =    sampler
       A__       =    group_ids
       A__       =    batch_size
   def __iter__(       self   :int							)->							Tuple:
       A__       =    defaultdict(lowercase_							)
       A__       =    defaultdict(lowercase_							)
  
       A__       =    0
       for idx in self.sampler:
           A__       =    self.group_ids[idx]
           buffer_per_group[group_id].append(lowercase_							)
           samples_per_group[group_id].append(lowercase_							)
           if len(buffer_per_group[group_id]							) == self.batch_size:
               yield buffer_per_group[group_id]  # TODO
               num_batches += 1
               del buffer_per_group[group_id]
           assert len(buffer_per_group[group_id]							) < self.batch_size
  
       # now we have run out of elements that satisfy
       # the group criteria, let's return the remaining
       # elements so that the size of the sampler is
       # deterministic
       A__       =    len(self							)
       A__       =    expected_num_batches - num_batches
       if num_remaining > 0:
           # for the remaining batches, group the batches by similar lengths
           A__       =    []
           for group_id, idxs in sorted(buffer_per_group.items()		, key=lambda lowercase_							: x[0]							):
               batch_idx.extend(lowercase_							)
               if len(lowercase_							) >= self.batch_size:
                   yield batch_idx[: self.batch_size]
                   A__       =    batch_idx[self.batch_size :]
                   num_remaining -= 1
           if len(lowercase_							) > 0:
               yield batch_idx
               num_remaining -= 1
       assert num_remaining == 0
  
  
  
   def __len__(       self   :int							)->							Tuple:
       return (len(self.sampler							) + self.batch_size - 1) // self.batch_size
  
  
  
  