text: # AUTOGENERATED! DO NOT EDIT! File to edit: 10_subcoco_utils.ipynb (unless otherwise specified).
  
  __all__ = ['fetch_data', 'fetch_subcoco', 'CocoDatasetStats', 'empty_list', 'load_stats', 'box_within_bounds',
             'is_notebook', 'overlay_img_bbox', 'bbox_to_rect', 'label_for_bbox', 'listify', 'tensorify',
             'SubCocoWrapper', 'iou_calc', 'match_true_false_neg', 'calc_wavg_F1', 'clamp_fn', 'digest_pred']
  
  # Cell
  import albumentations as A
  import cv2
  import glob
  import json
  import matplotlib.pyplot as plt
  import matplotlib.colors as mcolors
  import numpy as np
  import os
  import pickle
  import PIL
  import re
  import requests
  import sys
  import tarfile
  import torch
  import torchvision
  
  from albumentations.pytorch import ToTensorV2
  from collections import defaultdict
  from functools import reduce
  from IPython.utils import io
  from pathlib import Path
  from PIL import Image, ImageStat
  from pycocotools.coco import COCO
  from pycocotools.cocoeval import COCOeval
  from shutil import copyfile, rmtree
  from torch.utils.data import DataLoader, random_split
  from torchvision import transforms
  from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
  from tqdm import tqdm
  from typing import Hashable, List, Tuple, Union
  
  # Cell
  def fetch_data(url:str, datadir: Path, tgt_fname:str, chunk_size:int=8*1024, quiet=False):
      dest = datadir/tgt_fname
      if not quiet: print(f"Downloading from {url} to {dest}...")
      with requests.get(url, stream=True, timeout=10) as response:
          content_len = int(response.headers['content-length'])
          with open(dest, 'wb') as f:
              with tqdm(total=content_len) as pbar:
                  nbytes = 0
                  num_chunks = 0
                  for chunk in response.iter_content(chunk_size=chunk_size):
                      chunk_len = len(chunk)
                      nbytes += chunk_len
                      num_chunks += 1
                      f.write(chunk)
                      pbar.update(chunk_len)
  
      with tarfile.open(dest, 'r') as tar:
          extracted = []
          for item in tar:
              tar.extract(item, datadir)
              extracted.append(item.name)
  
      if not quiet: print(f"Downloaded {nbytes} from {url} to {dest}, extracted in {datadir}: {extracted[:3]},...,{extracted[-3:]}")
  
  # Cell
  def fetch_subcoco(
      datadir:str="workspace",
      url:str="https://s3.amazonaws.com/fast-ai-coco/coco_sample.tgz",
      img_subdir:str="train_sample",
  ):
      fname = url.split('/')[-1]
      froot = (fname.split('.'))[0]
      if not os.path.isdir(Path(datadir)/froot):
          fetch_data(url, Path(datadir), fname, chunk_size=1024*1024)
      json_fname = glob.glob(f"{datadir}/**/{img_subdir}.json", recursive=True)[0]
      with open(json_fname, 'r') as json_f:
          train_json = json.load(json_f)
  
      return train_json
  
  # Cell
  class CocoDatasetStats():
      # num_cats
      # num_imgs
      # num_bboxs
      # cat2name
      # class_map
      # lbl2cat
      # cat2lbl
      # img2fname
      # imgs
      # img2l2bs
      # img2lbs
      # l2ibs
      # avg_ncats_per_img
      # avg_nboxs_per_img
      # avg_nboxs_per_cat
      # img2sz
      # chn_means
      # chn_stds
      # avg_width
      # avg_height
      def __init__(self, ann:dict, img_dir:str):
  
          self.img_dir = Path(img_dir)
          self.num_cats = len(ann['categories'])
          self.num_imgs = len(ann['images'])
          self.num_bboxs = len(ann['annotations'])
  
          # build cat id to name, assign FRCNN
          self.cat2name = { c['id']: c['name'] for c in ann['categories'] }
  
          # need to translate coco subset category id to indexable label id
          # expected labels w 0 = background
          self.lbl2cat = { i: cid for i, cid in enumerate(self.cat2name.keys(),1) }
          self.cat2lbl = { cid: l for l, cid in self.lbl2cat.items() }
          self.lbl2name = { l:self.cat2name[cid] for l, cid in self.lbl2cat.items() }
          self.lbl2cat[0] = 0 # background
          self.cat2lbl[0] = 0 # background
  
          # img_id to file map
          self.img2fname = { img['id']:img['file_name'] for img in ann['images'] }
  
          # compute Images per channel means and std deviation using PIL.ImageStat.Stat()
  
          self.img2sz = {}
          n = 0
          mean = np.zeros((3,))
          stddev = np.zeros((3,))
          avgw = 0
          avgh = 0
          for img_id, img_fname in tqdm(self.img2fname.items()):
              img_fpath = self.img_dir/img_fname
              if not os.path.isfile(img_fpath): continue
              n = n + 1
              img = Image.open(img_fpath)
              istat = ImageStat.Stat(img)
              width, height = img.size
              avgw = (width + (n-1)*avgw)/n
              avgh = (height + (n-1)*avgh)/n
              mean = (istat.mean + (n-1)*mean)/n
              stddev = (istat.stddev + (n-1)*stddev)/n
              self.img2sz[img_id] = (width, height)
  
          # cleanup stats due to missing images
          self.num_imgs = len(self.img2sz)
          self.img2fname = { img_id: fname for img_id, fname in self.img2fname.items() if img_id in self.img2sz }
  
          self.chn_means = mean
          self.chn_stds = stddev
          self.avg_width = avgw
          self.avg_height = avgh
  
          # build up some maps for later analysis
          self.img2l2bs = {}
          self.img2lbs = defaultdict(empty_list)
          self.l2ibs = defaultdict(empty_list)
          anno_id = 0
          for a in ann['annotations']:
              img_id = a['image_id']
              if self.img2sz.get(img_id, None) == None: continue
              cat_id = a['category_id']
              lbl_id = self.cat2lbl[cat_id]
              l2bs_for_img = self.img2l2bs.get(img_id, { l:[] for l in range(1+len(self.cat2name))})
              (x, y, w, h) = a['bbox']
              b = (x, y, w, h)
              ib = (img_id, *b)
              lb = (lbl_id, *b)
              l2bs_for_img[lbl_id].append(b)
              self.l2ibs[lbl_id].append(ib)
              self.img2lbs[img_id].append(lb)
              self.img2l2bs[img_id] = l2bs_for_img
  
          acc_ncats_per_img = 0.0
          acc_nboxs_per_img = 0.0
          for img_id, l2bs in self.img2l2bs.items():
              acc_ncats_per_img += len(l2bs)
              for lbl_id, bs in l2bs.items():
                  acc_nboxs_per_img += len(bs)
  
          self.avg_ncats_per_img = acc_ncats_per_img/self.num_imgs
          self.avg_nboxs_per_img = acc_nboxs_per_img/self.num_imgs
  
          acc_nboxs_per_cat = 0.0
          for lbl_id, ibs in self.l2ibs.items():
              acc_nboxs_per_cat += len(ibs)
  
          self.avg_nboxs_per_cat = acc_nboxs_per_cat/self.num_cats
  
  def empty_list()->list: return [] # cannot use lambda as pickling will fail when saving models
  
  # Cell
  def load_stats(ann:dict, img_dir:str, force_reload:bool=False)->CocoDatasetStats:
      stats_fpath = Path(img_dir).parent/'stats.pkl'
      stats = None
      if os.path.isfile(stats_fpath) and not force_reload:
          try:
              stats = pickle.load( open(stats_fpath, "rb" ) )
          except Exception as e:
              print(f"Failed to read precomputed stats: {e}")
  
      if stats == None:
          stats = CocoDatasetStats(ann, img_dir)
          pickle.dump(stats, open(stats_fpath, "wb" ) )
  
      return stats
  
  # Cell
  def box_within_bounds(bx, by, bw, bh, img_width, img_height, min_margin_ratio, min_width_height_ratio):
      min_width = min_width_height_ratio*img_width
      min_height = min_width_height_ratio*img_height
      if bw < min_width or bh < min_height:
          return False
      top_margin = min_margin_ratio*img_height
      bottom_margin = img_height - top_margin
      left_margin = min_margin_ratio*img_width
      right_margin = img_width - left_margin
      if bx < left_margin or bx > right_margin:
          return False
      if by < top_margin or by > bottom_margin:
          return False
      return True
  
  # Cell
  def is_notebook():
      try:
          shell = get_ipython().__class__.__name__
          if shell == 'ZMQInteractiveShell':
              return True   # Jupyter notebook or qtconsole
          elif shell == 'TerminalInteractiveShell':
              return False  # Terminal running IPython
          else:
              return False  # Other type (?)
      except NameError:
          return False      # Probably standard Python interpreter
  
  # Cell
  def overlay_img_bbox(img:Image, l2bs: dict, l2name: dict):
      l2color = { l: colname for (l, colname) in zip(l2bs.keys(), mcolors.TABLEAU_COLORS.keys()) }
      fig = plt.figure(figsize=(16,10))
      fig = plt.imshow(img)
      for l, bs in l2bs.items():
          for b in bs:
              label_for_bbox(b, l2name[l])
              fig.axes.add_patch(bbox_to_rect(b, l2color[l]))
  
  def bbox_to_rect(bbox:Tuple[int, int, int, int], color:str):
      return plt.Rectangle(
          xy=(bbox[0], bbox[1]), width=bbox[2], height=bbox[3],
          fill=False, edgecolor=color, linewidth=2)
  
  def label_for_bbox(bbox:Tuple[int, int, int, int], label:str):
      return plt.text(bbox[0], bbox[1], f"{label}", color='#ffffff', fontsize=12)
  
  # Cell
  def listify(tensorOrIterable):
      return tensorOrIterable.tolist() if type(tensorOrIterable) == torch.Tensor else list(tensorOrIterable)
  
  def tensorify(tensorOrIterable):
      return tensorOrIterable if type(tensorOrIterable) == torch.Tensor else torch.Tensor(tensorOrIterable)
  
  class SubCocoWrapper():
      def __init__(self, prediction, target, width, height):
          # turn tgt: { "boxes": [...], "labels": [...], "image_id": "xxx", "area": [...], "iscrowd": 0 }
          # into COCO with dataset dict of this form:
          # { images: [{'id':int, 'file_name':str, 'width':int, 'height':int}], categories: [int,...],
          #   annotations: [{'id':int, 'image_id': int, 'category_id': int, 'bbox': (x,y,width,height)}, 'area':float, 'iscrowd':0] }
          # see https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocotools/coco.py
  
          def toCOCO(rec:dict)->COCO:
              coco = COCO()
              img_id = int(rec.get('image_id', 0)) # could be tensor, cast to int
              coco.dataset['images'] = [ {'id': img_id, 'file_name': f'{img_id:012d}.jpg', 'width': width, 'height': height} ]
              coco.dataset['categories'] = []
              coco.dataset['annotations'] = []
              scores = rec.get('scores',[])
              anno_ids =rec.get('ids',[])
              for bi, b in enumerate(listify(rec['boxes'])):
                  x, y, w, h = b
                  cat_id = int(rec['labels'][bi])
                  anno_id = int(anno_ids[bi]) if bi < len(anno_ids) else 0
                  s = float(scores[bi]) if bi < len(scores) else 0.
                  coco.dataset["annotations"].append({'id': anno_id, 'image_id': img_id, 'category_id': cat_id, 'bbox': b, 'iscrowd':0, 'area':w*h, 'score':s})
                  coco.dataset['categories'].append({'id':cat_id})
  
              coco.createIndex()
              return coco
  
          with io.capture_output() as captured:
              self.target = toCOCO(target)
              self.prediction = toCOCO(prediction)
  
      def metrics(self)->float:
          with io.capture_output() as captured:
              cocoeval = COCOeval(self.target, self.prediction, "bbox")
              cocoeval.evaluate()
              cocoeval.accumulate()
              cocoeval.summarize()
          return cocoeval.stats
  
  # Cell
  def iou_calc(x1,y1,w1,h1, x2,y2,w2,h2):
      r1 = x1+w1 # right of box1
      b1 = y1+h1 # bottom of box1
      r2 = x2+w2 # right of box2
      b2 = y2+h2 # bottom of box2
      a1 = 1.0*w1*h1
      a2 = 1.0*w2*h2
      ia = 0.0 # intercept
  
      if x1 <= x2 <= r1:
          if y1 <= y2 <= b1 or y1 <= b2 <= b1:
              ia = (min(r1,r2)-max(x1,x2))*(min(b1,b2)-max(y1,y2))
      elif x1 <= r2 <= r1:
          if y1 <= y2 <= b1 or y1 <= b2 <= b1:
              ia = (min(r1,r2)-max(x1,x2))*(min(b1,b2)-max(y1,y2))
  
      #print(a1, a2, ia)
      iou = ia/(a1+a2-ia)
      return iou
  
  # Cell
  def match_true_false_neg(pred, tgt, scut=0.5, ithr=0.5):
  
      #Init map of labels to 3 counters: True positive, False positive, false Negative
      l2tfn = defaultdict(lambda: (0,0,0))
  
      #Find all prediction above confidence score cutoff
      pscores = listify(pred['scores'])
      # original (pscores > scut).nonzero(as_tuple=True) but can't handle non tensor
      pidxs = reduce(lambda l, idx_sc: l+[idx_sc[0]] if idx_sc[1] > scut else l, enumerate(pscores), [])
      #print(f"pidxs = {pidxs}")
      pboxs = listify(tensorify(pred['boxes'])[pidxs])
  
  
      tboxs = listify(tgt['boxes'])
      tls = listify(tgt['labels'])
      pls = listify(pred['labels'])
  
      for tl,tb in zip(tls,tboxs):
          #init maxIoU, maxIndex, maxTrueIoU, maxTrueIndex
          maxIoU, maxIndex, maxTrueIoU, maxTrueIndex = -1, -1, -1, -1
  
          x1,y1,w1,h1 = float(tb[0]), float(tb[1]), float(tb[2]), float(tb[3])
          for pi,(pl,pb) in enumerate(zip(pls,pboxs)):
              x2,y2,w2,h2 = float(pb[0]), float(pb[1]), float(pb[2]), float(pb[3])
              iou = iou_calc(x1,y1,w1,h1, x2,y2,w2,h2)
              #ensure IoU above threshold
              if iou < ithr: continue
  
              if iou > maxIoU: #Found new Max IoU
                  maxIoU = iou
                  maxIndex = pi
  
              if int(pl)==int(tl) and iou > maxTrueIoU: #Found new Max True IoU
                  maxTrueIoU = iou
                  maxTrueIndex = pi
  
          if maxTrueIndex >= 0:
              #update true positive counter under l2tfn
              l2tfn[tl] = (l2tfn[tl][0]+1, l2tfn[tl][1], l2tfn[tl][2])
              #remove pb@maxTrueIndex from pboxs
              pls.pop(pi)
              pboxs.pop(pi)
          elif maxIndex >= 0:
              #update false positive counter under l2tfn
              l2tfn[tl] = (l2tfn[tl][0], l2tfn[tl][1]+1, l2tfn[tl][2])
              #remove pb@maxIndex from pboxs
              pls.pop(pi)
              pboxs.pop(pi)
          else:
              #update false negative counter under l2tfn
              l2tfn[tl] = (l2tfn[tl][0], l2tfn[tl][1], l2tfn[tl][2]+1)
  
      #Count remaining unmatched predictions as False positive of Background.
      l2tfn[0] = (0, len(pboxs), 0)
  
      return l2tfn
  
  # Cell
  def calc_wavg_F1(pred, tgt, scut=0.5, ithr=0.5):
      l2tfn = match_true_false_neg(pred, tgt, scut=scut, ithr=ithr)
      lset = l2tfn.keys()
      bsum = 0
      l2num = defaultdict(lambda:0)
      for l, (t,f,n) in l2tfn.items():
          boxes = t+f+n
          l2num[l]= boxes
          bsum += boxes
  
      l2f1 = {}
      for l in lset:
          tp, fp, fn = l2tfn[l]
          precision = 0 if tp == 0 else tp/(tp+fp)
          recall = 0 if tp == 0 else tp/(tp+fn)
          f1 = 0 if precision*recall == 0 else 2/(1/precision + 1/recall)
          l2f1[l] = f1
  
      #print(l2tfn, l2f1)
      acc = 0.
      for l, f1 in l2f1.items():
          lnum = l2num[l]
          f1 = l2f1[l]
          acc += f1*(lnum/bsum)
  
      return acc
  
  # Cell
  def clamp_fn(lo, hi):
      return lambda v: min(hi,max(lo,v))
  
  def digest_pred(l2name, pred, cutoff=0.5, img_sz=128):
      scores = pred['scores']
      pass_idxs = [i for i in range(len(scores)) if scores[i] > cutoff]
      lbls = pred['labels'][pass_idxs]
      bboxs = pred['boxes'][pass_idxs]
      if len(pass_idxs) > 1:
          lbls = lbls.squeeze()
          bboxs = bboxs.squeeze()
      l2bs = defaultdict(lambda: [])
      for l, b in zip(lbls.tolist(), bboxs.tolist()):
          l = int(l)
          x,y,w,h = map(clamp_fn(0,img_sz), b)
          bs = l2bs[l]
          bs.append((x,y,w,h))
      return l2bs