{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Introduction \nThis notebook uses the DeepMAC model from the \"[The surprising impact of mask-head architecture on novel class segmentation](https://arxiv.org/abs/2104.00613)\" paper.\nIn this colab, we use bounding boxes from [MegaDetector](https://github.com/microsoft/CameraTraps/blob/master/megadetector.md) model and use the\nDeepMAC model to generate instance masks. \n\nSee also: [DeepMAC colab](https://github.com/tensorflow/models/blob/master/research/object_detection/colab_tutorials/deepmac_colab.ipynb) in the Object Detection API","metadata":{}},{"cell_type":"markdown","source":"# Imports and Definitions","metadata":{}},{"cell_type":"code","source":"import glob\nimport io\nimport json\nimport logging\nimport os\nimport random\nimport warnings\n\nimport imageio\nfrom IPython.display import display, Javascript\nfrom IPython.display import Image as IPyImage\nimport matplotlib\nfrom matplotlib import patches\nimport matplotlib.pyplot as plt\nimport numpy as np\nfrom PIL import Image, ImageDraw, ImageFont\nimport scipy.misc\nfrom six import BytesIO\nfrom skimage import color\nfrom skimage import transform\nfrom skimage import util\nfrom skimage.color import rgb_colors\nimport tensorflow as tf\n\nCOLORS = ([rgb_colors.cyan, rgb_colors.orange, rgb_colors.pink,\n           rgb_colors.purple, rgb_colors.limegreen , rgb_colors.crimson] +\n          [(color) for (name, color) in color.color_dict.items()])\nrandom.shuffle(COLORS)\n\nlogging.disable(logging.WARNING)\n\n\ndef read_image(path):\n  \"\"\"Read an image and optionally resize it for better plotting.\"\"\"\n  with open(path, 'rb') as f:\n    img = Image.open(f)\n    return np.array(img, dtype=np.uint8)\n\ndef read_json(path):\n  with open(path) as f:\n    return json.load(f)\n\ndef create_detection_map(annotations):\n  \"\"\"Creates a dict mapping IDs to detections.\"\"\"\n\n  ann_map = {}\n  for image in annotations['images']:\n    ann_map[image['id']] = image['detections']\n  return ann_map\n\ndef get_mask_prediction_function(model):\n  \"\"\"Get single image mask prediction function using a model.\"\"\"\n\n  @tf.function\n  def predict_masks(image, boxes):\n    height, width, _ = image.shape.as_list()\n    batch = image[tf.newaxis]\n    boxes = boxes[tf.newaxis]\n\n    detections = model(batch, boxes)\n    masks = detections['detection_masks']\n\n    return reframe_box_masks_to_image_masks(masks[0], boxes[0],\n                                             height, width)\n\n  return predict_masks\n\ndef convert_boxes(boxes):\n  xmin, ymin, width, height = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]\n  ymax = ymin + height\n  xmax = xmin + width\n\n  return np.stack([ymin, xmin, ymax, xmax], axis=1).astype(np.float32)\n\n\n# Copied from tensorflow/models\ndef reframe_box_masks_to_image_masks(box_masks, boxes, image_height,\n                                     image_width, resize_method='bilinear'):\n  \"\"\"Transforms the box masks back to full image masks.\n  Embeds masks in bounding boxes of larger masks whose shapes correspond to\n  image shape.\n  Args:\n    box_masks: A tensor of size [num_masks, mask_height, mask_width].\n    boxes: A tf.float32 tensor of size [num_masks, 4] containing the box\n           corners. Row i contains [ymin, xmin, ymax, xmax] of the box\n           corresponding to mask i. Note that the box corners are in\n           normalized coordinates.\n    image_height: Image height. The output mask will have the same height as\n                  the image height.\n    image_width: Image width. The output mask will have the same width as the\n                 image width.\n    resize_method: The resize method, either 'bilinear' or 'nearest'. Note that\n      'bilinear' is only respected if box_masks is a float.\n  Returns:\n    A tensor of size [num_masks, image_height, image_width] with the same dtype\n    as `box_masks`.\n  \"\"\"\n  resize_method = 'nearest' if box_masks.dtype == tf.uint8 else resize_method\n  # TODO(rathodv): Make this a public function.\n  def reframe_box_masks_to_image_masks_default():\n    \"\"\"The default function when there are more than 0 box masks.\"\"\"\n    def transform_boxes_relative_to_boxes(boxes, reference_boxes):\n      boxes = tf.reshape(boxes, [-1, 2, 2])\n      min_corner = tf.expand_dims(reference_boxes[:, 0:2], 1)\n      max_corner = tf.expand_dims(reference_boxes[:, 2:4], 1)\n      denom = max_corner - min_corner\n      # Prevent a divide by zero.\n      denom = tf.math.maximum(denom, 1e-4)\n      transformed_boxes = (boxes - min_corner) / denom\n      return tf.reshape(transformed_boxes, [-1, 4])\n\n    box_masks_expanded = tf.expand_dims(box_masks, axis=3)\n    num_boxes = tf.shape(box_masks_expanded)[0]\n    unit_boxes = tf.concat(\n        [tf.zeros([num_boxes, 2]), tf.ones([num_boxes, 2])], axis=1)\n    reverse_boxes = transform_boxes_relative_to_boxes(unit_boxes, boxes)\n\n    # TODO(vighneshb) Use matmul_crop_and_resize so that the output shape\n    # is static. This will help us run and test on TPUs.\n    resized_crops = tf.image.crop_and_resize(\n        image=box_masks_expanded,\n        boxes=reverse_boxes,\n        box_indices=tf.range(num_boxes),\n        crop_size=[image_height, image_width],\n        method=resize_method,\n        extrapolation_value=0)\n    return tf.cast(resized_crops, box_masks.dtype)\n\n  image_masks = tf.cond(\n      tf.shape(box_masks)[0] > 0,\n      reframe_box_masks_to_image_masks_default,\n      lambda: tf.zeros([0, image_height, image_width, 1], box_masks.dtype))\n  return tf.squeeze(image_masks, axis=3)\n\ndef plot_image_annotations(image, boxes, masks, darken_image=0.5):\n  fig, ax = plt.subplots(figsize=(16, 12))\n  ax.set_axis_off()\n  image = (image * darken_image).astype(np.uint8)\n  ax.imshow(image)\n\n  height, width, _ = image.shape\n\n  num_colors = len(COLORS)\n  color_index = 0\n\n  for box, mask in zip(boxes, masks):\n    ymin, xmin, ymax, xmax = box\n    ymin *= height\n    ymax *= height\n    xmin *= width\n    xmax *= width\n\n    color = COLORS[color_index]\n    color = np.array(color)\n    rect = patches.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,\n                             linewidth=2.5, edgecolor=color, facecolor='none')\n    ax.add_patch(rect)\n    mask = (mask > 0.5).astype(np.float32)\n    color_image = np.ones_like(image) * color[np.newaxis, np.newaxis, :]\n    color_and_mask = np.concatenate(\n        [color_image, mask[:, :, np.newaxis]], axis=2)\n\n    ax.imshow(color_and_mask, alpha=0.5)\n\n    color_index = (color_index + 1) % num_colors\n\n  return ax","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Downloading the checkpoint\nMake sure internet access is enabled from the right panel.","metadata":{}},{"cell_type":"code","source":"!curl -o /kaggle/working/deepmac_1024x1024_coco17.tar.gz http://download.tensorflow.org/models/object_detection/tf2/20210329/deepmac_1024x1024_coco17.tar.gz\n!tar -xzf /kaggle/working/deepmac_1024x1024_coco17.tar.gz\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Load the model into memory\nThis can take a minute or so","metadata":{}},{"cell_type":"code","source":"model = tf.keras.models.load_model('/kaggle/working/deepmac_1024x1024_coco17/saved_model')\nprediction_function = get_mask_prediction_function(model)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Load metadata information","metadata":{}},{"cell_type":"code","source":"BOX_ANNOTATION_FILE = '/kaggle/input/iwildcam2021-fgvc8/metadata/iwildcam2021_megadetector_results.json'\ndetection_map = create_detection_map(read_json(BOX_ANNOTATION_FILE))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Running inference","metadata":{}},{"cell_type":"code","source":"\nimage_path = '/kaggle/input/iwildcam2021-fgvc8/train/905e980e-21bc-11ea-a13a-137349068a90.jpg'\nimage_id = os.path.basename(image_path).rstrip('.jpg')\n\nif image_id not in detection_map:\n  print(f'Image {image_path} is missing detection data.')\nelif len(detection_map[image_id]) == 0:\n  print(f'There are no detected objects in the image {image_path}.')\nelse:\n  detections = detection_map[image_id]\n  image = read_image(image_path)\n  bboxes = np.array([det['bbox'] for det in detections])\n  bboxes = convert_boxes(bboxes)\n  masks = prediction_function(tf.convert_to_tensor(image),\n                              tf.convert_to_tensor(bboxes, dtype=tf.float32))\n  plot_image_annotations(image, bboxes, masks.numpy(), darken_image=0.75)","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}