text: import config
  import matplotlib.pyplot as plt
  import matplotlib.patches as patches
  import numpy as np
  import os
  import random
  import torch
  from nms import nms
  from collections import Counter
  from torch.utils.data import DataLoader
  from tqdm import tqdm
  
  def plot_particles(image, boxes, scores=True, pixels=False):
      """Plots predicted  particles on the image"""
      im = np.array(image)
      height, width, _ = im.shape if not pixels else (1,1,1)
  
      # Create figure and axes
      bboxes = np.asarray(boxes)
      # bboxes = bboxes[bboxes[...,0]>0.5]
      if len(bboxes) > 0:
          if scores:
              plt.scatter(bboxes[..., 1] * width, bboxes[..., 2] * height)
          else:
              plt.scatter(bboxes[..., 0] * width, bboxes[..., 1] * height)
      # Display the image
      plt.imshow(im, cmap="gray")
  
      plt.show()
  
  
  def cells_to_bboxes(predictions, S, is_preds=True):
      """
      Scales the predictions coming from the model to
      be relative to the entire image such that they for example later
      can be plotted or compared to targets.
      :param predictions: tensor of size (N, 3, S, S, num_classes+5)
      :param S: the number of cells the image is divided in on the width (and height)
      :param is_preds: whether the input is predictions or the true bounding boxes
      :returns converted_bboxes: the converted boxes of sizes (N, num_anchors, S, S, 1+5) with
                        object score, bounding box coordinates
      """
      BATCH_SIZE = predictions.shape[0]
      box_predictions = predictions[..., 1:6]
      if is_preds:
          box_predictions[..., 0:2] = torch.sigmoid(box_predictions[..., 0:2])
          box_predictions[..., 3:4] = box_predictions[..., 3:4] * 1e-7  # r predictions
          scores = torch.sigmoid(predictions[..., 0:1])
      else:
          scores = predictions[..., 0:1]
      cell_indices = (
          torch.arange(S)
          .repeat(predictions.shape[0], 3, S, 1)
          .unsqueeze(-1)
          .to(predictions.device)
      )
      x = 1 / S * (box_predictions[..., 0:1] + cell_indices)
      y = 1 / S * (box_predictions[..., 1:2] + cell_indices.permute(0, 1, 3, 2, 4))
      converted_bboxes = torch.cat(
          (scores, x, y, box_predictions[..., 2:5]), dim=-1
      ).reshape(BATCH_SIZE, 3 * S * S, 6)
      return converted_bboxes.tolist()
  
  
  
  
  def get_mean_std(loader):
      # var[X] = E[X**2] - E[X]**2
      channels_sum, channels_sqrd_sum, num_batches = 0, 0, 0
  
      for data, _ in tqdm(loader):
          channels_sum += torch.mean(data, dim=[0, 2, 3])
          channels_sqrd_sum += torch.mean(data ** 2, dim=[0, 2, 3])
          num_batches += 1
  
      mean = channels_sum / num_batches
      std = (channels_sqrd_sum / num_batches - mean ** 2) ** 0.5
  
      return mean, std
  
  
  def save_checkpoint(model, optimizer, filename="my_checkpoint.pth.tar"):
      print("=> Saving checkpoint")
      checkpoint = {
          "state_dict": model.state_dict(),
          "optimizer": optimizer.state_dict(),
      }
      torch.save(checkpoint, filename)
  
  
  def load_checkpoint(checkpoint_file, model, optimizer, lr):
      print("=> Loading checkpoint")
      checkpoint = torch.load(checkpoint_file, map_location=config.DEVICE)
      model.load_state_dict(checkpoint["state_dict"])
      optimizer.load_state_dict(checkpoint["optimizer"])
  
      # If we don't do this then it will just have learning rate of old checkpoint
      # and it will lead to many hours of debugging \:
      for param_group in optimizer.param_groups:
          param_group["lr"] = lr
  
  
  def get_loaders(train_dir, test_dir, batch_size, num_workers, pin_memory):
      from load_simulated_data import YOLODataset
  
      IMAGE_SIZE = config.IMAGE_SIZE
      train_dataset = YOLODataset(
          anchors = config.ANCHORS,
          datadir=train_dir,
          transform=config.transform,
          S=[IMAGE_SIZE // 32, IMAGE_SIZE // 16, IMAGE_SIZE // 8],
      )
      test_dataset = YOLODataset(
          anchors=config.ANCHORS,
          datadir=test_dir,
          transform=config.transform,
          S=[IMAGE_SIZE // 32, IMAGE_SIZE // 16, IMAGE_SIZE // 8],
      )
      train_loader = DataLoader(
          dataset=train_dataset,
          batch_size=batch_size,
          num_workers=num_workers,
          pin_memory=pin_memory,
          shuffle=True,
          drop_last=False,
      )
      test_loader = DataLoader(
          dataset=test_dataset,
          batch_size=batch_size,
          num_workers=num_workers,
          pin_memory=pin_memory,
          shuffle=False,
          drop_last=False,
      )
  
      return train_loader, test_loader
  
  
  def plot_couple_examples(model, loader, threshold, nms_threshold):
  
      model.eval()
      x, y = next(iter(loader))
      x = x.to(config.DEVICE).squeeze(0)
      with torch.cuda.amp.autocast():
          with torch.no_grad():
              out = model(x)
              bboxes = [[] for _ in range(x.shape[0])]
              for i in range(3):
                  batch_size, A, S, _, _ = out[i].shape
                  boxes_scale_i = cells_to_bboxes(out[i], S=S, is_preds=True)
                  for idx, (box) in enumerate(boxes_scale_i):
                      bboxes[idx] += box
  
      model.train()
  
      for i in range(batch_size):
          nms_boxes = nms(np.array(bboxes[i]), conf_threshold=threshold, threshold=nms_threshold).tolist()
          print(len(nms_boxes))
          plot_particles(x[i, 2:3, :, :].permute(1, 2, 0).detach().cpu(), nms_boxes)
  
  
  def seed_everything(seed=42):
      os.environ["PYTHONHASHSEED"] = str(seed)
      random.seed(seed)
      np.random.seed(seed)
      torch.manual_seed(seed)
      torch.cuda.manual_seed(seed)
      torch.cuda.manual_seed_all(seed)
      torch.backends.cudnn.deterministic = True
      torch.backends.cudnn.benchmark = False
  