MaskRCNN Inference (tinygrad#884)

* MaskRCNN weights loading * backbone maybe works * backbone works, but resnet body atol 1e-3 * RPN Call, but veryy wrong output * fixed topk * RPN maybe works, not sure about nms * Fix cursed modules * add back editorconfig * Full call, wrong output * Full call works * fix mask * use NMS from retinanet * Removing extra funcs * refactor * readable * Add example to run model * remove filter * Fix split, batched inference is worse * Fix image sizes * Matching reference * merge master * add filter on top detections * cuda backend fixed * add model eval and spec * convert images to rgb * fix eval * simplify examples code * remove extra code * meshgrid using tinygrad * removing numpy * roi align, floor, ceil * remove numpy from level_mapper * remove numpy from pooler * Revert "Merge branch 'master' of github.com:kunwar31/tinygrad into mrcnn-inference" This reverts commit 4b95a3c, reversing changes made to 98f2b1f. * roi align gather * fix master merge * revert to old floor, ceil as ints present in domain * use log2 op * fix indexes * weird bug with ints and gpu * weird bug with ints and gpu * refactors, add env var for gather * floor with contiguous, where * refactor topk, sort * remove staticmethod * refactor stride * remove log2 mlop * realize -> contiguous * refactor forward * remove num_classes, stride_in_1x1 from state * refactor forward * refactoring * flake8 * removing numpy in anchor gen, use numpy for gather, nonzero, optimize topk * keep using tinygrad for smaller gathers * fix empty tensors * comms * move from tensor.py * resnet test passing * add coco dataset back * fix spaces * add test for log2 * no need to create Tensors * no need to create Tensors --------- Co-authored-by: Kunwar Raj Singh <[email protected]>
Crocozx · Jun 25, 2023 · 5d3310c · 5d3310c
1 parent 0f281e7
commit 5d3310c
Show file tree

Hide file tree

Showing 8 changed files with 1,854 additions and 23 deletions.
diff --git a/datasets/coco.py b/datasets/coco.py
@@ -0,0 +1,200 @@
+import json
+import pathlib
+import zipfile
+import numpy as np
+from extra.utils import download_file
+import pycocotools._mask as _mask
+from examples.mask_rcnn import Masker
+from pycocotools.coco import COCO
+from pycocotools.cocoeval import COCOeval
+
+iou         = _mask.iou
+merge       = _mask.merge
+frPyObjects = _mask.frPyObjects
+
+BASEDIR = pathlib.Path(__file__).parent.parent / "datasets/COCO"
+
+def create_dict(key_row, val_row, rows): return {row[key_row]:row[val_row] for row in rows}
+
+
+if not pathlib.Path(BASEDIR/'val2017').is_dir():
+  fn = BASEDIR/'val2017.zip'
+  download_file('http://images.cocodataset.org/zips/val2017.zip',fn)
+  with zipfile.ZipFile(fn, 'r') as zip_ref:
+    zip_ref.extractall(BASEDIR)
+  fn.unlink()
+
+
+if not pathlib.Path(BASEDIR/'annotations').is_dir():
+  fn = BASEDIR/'annotations_trainval2017.zip'
+  download_file('http://images.cocodataset.org/annotations/annotations_trainval2017.zip',fn)
+  with zipfile.ZipFile(fn, 'r') as zip_ref:
+    zip_ref.extractall(BASEDIR)
+  fn.unlink()
+
+with open(BASEDIR/'annotations/instances_val2017.json', 'r') as f:
+  annotations_raw = json.loads(f.read())
+images = annotations_raw['images']
+categories = annotations_raw['categories']
+annotations = annotations_raw['annotations']
+file_name_to_id = create_dict('file_name', 'id', images)
+id_to_width = create_dict('id', 'width', images)
+id_to_height = create_dict('id', 'height', images)
+json_category_id_to_contiguous_id = {v['id']: i + 1 for i, v in enumerate(categories)}
+contiguous_category_id_to_json_id = {v:k for k,v in json_category_id_to_contiguous_id.items()}
+
+
+def encode(bimask):
+  if len(bimask.shape) == 3:
+    return _mask.encode(bimask)
+  elif len(bimask.shape) == 2:
+    h, w = bimask.shape
+    return _mask.encode(bimask.reshape((h, w, 1), order='F'))[0]
+
+def decode(rleObjs):
+  if type(rleObjs) == list:
+    return _mask.decode(rleObjs)
+  else:
+    return _mask.decode([rleObjs])[:,:,0]
+
+def area(rleObjs):
+  if type(rleObjs) == list:
+    return _mask.area(rleObjs)
+  else:
+    return _mask.area([rleObjs])[0]
+
+def toBbox(rleObjs):
+  if type(rleObjs) == list:
+    return _mask.toBbox(rleObjs)
+  else:
+    return _mask.toBbox([rleObjs])[0]
+
+
+def convert_prediction_to_coco_bbox(file_name, prediction):
+  coco_results = []
+  try:
+    original_id = file_name_to_id[file_name]
+    if len(prediction) == 0:
+      return coco_results
+
+    image_width = id_to_width[original_id]
+    image_height = id_to_height[original_id]
+    prediction = prediction.resize((image_width, image_height))
+    prediction = prediction.convert("xywh")
+
+    boxes = prediction.bbox.numpy().tolist()
+    scores = prediction.get_field("scores").numpy().tolist()
+    labels = prediction.get_field("labels").numpy().tolist()
+
+    mapped_labels = [contiguous_category_id_to_json_id[int(i)] for i in labels]
+
+    coco_results.extend(
+      [
+        {
+          "image_id": original_id,
+          "category_id": mapped_labels[k],
+          "bbox": box,
+          "score": scores[k],
+        }
+          for k, box in enumerate(boxes)
+      ]
+    )
+  except Exception as e:
+    print(file_name, e)
+  return coco_results
+
+masker = Masker(threshold=0.5, padding=1)
+
+def convert_prediction_to_coco_mask(file_name, prediction):
+  coco_results = []
+  try:
+    original_id = file_name_to_id[file_name]
+    if len(prediction) == 0:
+      return coco_results
+
+    image_width = id_to_width[original_id]
+    image_height = id_to_height[original_id]
+    prediction = prediction.resize((image_width, image_height))
+    masks = prediction.get_field("mask")
+
+    scores = prediction.get_field("scores").numpy().tolist()
+    labels = prediction.get_field("labels").numpy().tolist()
+
+    masks = masker([masks], [prediction])[0].numpy()
+
+    rles = [
+      encode(np.array(mask[0, :, :, np.newaxis], order="F"))[0]
+      for mask in masks
+    ]
+    for rle in rles:
+      rle["counts"] = rle["counts"].decode("utf-8")
+
+    mapped_labels = [contiguous_category_id_to_json_id[int(i)] for i in labels]
+
+    coco_results.extend(
+      [
+        {
+          "image_id": original_id,
+          "category_id": mapped_labels[k],
+          "segmentation": rle,
+          "score": scores[k],
+        }
+          for k, rle in enumerate(rles)
+      ]
+    )
+  except Exception as e:
+    print(file_name, e)
+  return coco_results
+
+
+
+def accumulate_predictions_for_coco(coco_results, json_result_file, rm=False):
+  path = pathlib.Path(json_result_file)
+  if rm and path.exists(): path.unlink()
+  with open(path, "a") as f:
+    for s in coco_results:
+      f.write(json.dumps(s))
+      f.write('\n')
+
+def remove_dup(l):
+  seen = set()
+  seen_add = seen.add
+  return [x for x in l if not (x in seen or seen_add(x))]
+
+class NpEncoder(json.JSONEncoder):
+  def default(self, obj):
+    if isinstance(obj, np.integer):
+      return int(obj)
+    if isinstance(obj, np.floating):
+      return float(obj)
+    if isinstance(obj, np.ndarray):
+      return obj.tolist()
+    return super(NpEncoder, self).default(obj)
+
+
+def evaluate_predictions_on_coco(json_result_file, iou_type="bbox"):
+  coco_results = []
+  with open(json_result_file, "r") as f:
+    for line in f:
+      coco_results.append(json.loads(line))
+
+  coco_gt = COCO(str(BASEDIR/'annotations/instances_val2017.json'))
+  set_of_json = remove_dup([json.dumps(d, cls=NpEncoder) for d in coco_results])
+  unique_list = [json.loads(s) for s in set_of_json]
+
+  with open(f'{json_result_file}.flattend', "w") as f:
+    json.dump(unique_list, f)
+
+  coco_dt = coco_gt.loadRes(str(f'{json_result_file}.flattend')) 
+  coco_eval = COCOeval(coco_gt, coco_dt, iou_type)
+  coco_eval.evaluate()
+  coco_eval.accumulate()
+  coco_eval.summarize()
+  return coco_eval
+
+def iterate(files, bs=1):
+  batch = []
+  for file in files:
+    batch.append(file)
+    if len(batch) >= bs: yield batch; batch = []
+  if len(batch) > 0: yield batch; batch = []