content: 
  import torch
  import numpy as np
  import sys
  import glob
  import json
  
  def neighborhoods(mu, x_range, y_range, sigma, circular_x=True, gaussian=False):
      """ Generate masks centered at mu of the given x and y range with the
          origin in the centre of the output
      Inputs:
          mu: tensor (N, 2)
      Outputs:
          tensor (N, y_range, s_range)
      """
      x_mu = mu[:,0].unsqueeze(1).unsqueeze(1)
      y_mu = mu[:,1].unsqueeze(1).unsqueeze(1)
  
      # Generate bivariate Gaussians centered at position mu
      x = torch.arange(start=0,end=x_range, device=mu.device, dtype=mu.dtype).unsqueeze(0).unsqueeze(0)
      y = torch.arange(start=0,end=y_range, device=mu.device, dtype=mu.dtype).unsqueeze(1).unsqueeze(0)
  
      y_diff = y - y_mu
      x_diff = x - x_mu
      if circular_x:
          x_diff = torch.min(torch.abs(x_diff), torch.abs(x_diff + x_range))
      if gaussian:
          output = torch.exp(-0.5 * ((x_diff/sigma[0])**2 + (y_diff/sigma[1])**2 ))
      else:
          output = torch.logical_and(
              torch.abs(x_diff) <= sigma[0], torch.abs(y_diff) <= sigma[1]
          ).type(mu.dtype)
  
      return output
  
  
  def nms(pred, max_predictions=10, sigma=(1.0,1.0), gaussian=False):
      ''' Input (batch_size, 1, height, width) '''
  
      shape = pred.shape
  
      output = torch.zeros_like(pred)
      flat_pred = pred.reshape((shape[0],-1))  # (BATCH_SIZE, 24*48)
      supp_pred = pred.clone()
      flat_output = output.reshape((shape[0],-1))  # (BATCH_SIZE, 24*48)
  
      for i in range(max_predictions):
          # Find and save max over the entire map
          flat_supp_pred = supp_pred.reshape((shape[0],-1))
          val, ix = torch.max(flat_supp_pred, dim=1)
          indices = torch.arange(0,shape[0])
          flat_output[indices,ix] = flat_pred[indices,ix]
  
          # Suppression
          y = ix / shape[-1]
          x = ix % shape[-1]
          mu = torch.stack([x,y], dim=1).float()
  
          g = neighborhoods(mu, shape[-1], shape[-2], sigma, gaussian=gaussian)
  
          supp_pred *= (1-g.unsqueeze(1))
  
      output[output < 0] = 0
      return output
  
  
  def print_progress(iteration, total, prefix='', suffix='', decimals=1, bar_length=50):
      """
      Call in a loop to create terminal progress bar
      @params:
          iteration   - Required  : current iteration (Int)
          total       - Required  : total iterations (Int)
          prefix      - Optional  : prefix string (Str)
          suffix      - Optional  : suffix string (Str)
          decimals    - Optional  : positive number of decimals in percent complete (Int)
          bar_length  - Optional  : character length of bar (Int)
      """
      str_format = "{0:." + str(decimals) + "f}"
      percents = str_format.format(100 * (iteration / float(total)))
      filled_length = int(round(bar_length * iteration / float(total)))
      bar = '█' * filled_length + '-' * (bar_length - filled_length)
  
      sys.stdout.write('\r%s |%s| %s%s %s' % (prefix, bar, percents, '%', suffix)),
  
      if iteration == total:
          sys.stdout.write('\n')
      sys.stdout.flush()
  
  
  def get_attention_mask(num_imgs=12, neighbor=1):
      assert neighbor <= 5
  
      mask = np.zeros((num_imgs,num_imgs))
      t = np.zeros(num_imgs)
      t[:neighbor+1] = np.ones(neighbor+1)
      if neighbor != 0:
          t[-neighbor:] = np.ones(neighbor)
      for ri in range(num_imgs):
          mask[ri] = t
          t = np.roll(t, 1)
  
      return torch.from_numpy(mask).reshape(1,1,num_imgs,num_imgs).long()