Mask_RCNN_TensorFlow_Object_Detection_API.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "Mask RCNN - TensorFlow Object Detection API.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [
        "Sqzm5Om4rEvq",
        "JL-VU8HwabXj",
        "q0pTF60YSNxZ",
        "PALYIrj9sSrN",
        "pziK2r8-MOQr",
        "Mf_orYdYQrIW"
      ],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.4"
    },
    "latex_envs": {
      "LaTeX_envs_menu_present": true,
      "autoclose": false,
      "autocomplete": true,
      "bibliofile": "biblio.bib",
      "cite_by": "apalike",
      "current_citInitial": 1,
      "eqLabelWithNumbers": true,
      "eqNumInitial": 1,
      "hotkeys": {
        "equation": "Ctrl-E",
        "itemize": "Ctrl-I"
      },
      "labels_anchors": false,
      "latex_user_defs": false,
      "report_style_numbering": false,
      "user_envs_cfg": false
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/anshupandey/Computer-Vision/blob/master/Mask_RCNN_TensorFlow_Object_Detection_API.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Sqzm5Om4rEvq"
      },
      "source": [
        "#### Download a trained model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zdkJQ2MuYwMh"
      },
      "source": [
        "We will download a trained model from [TensorFlow detection model zoo](https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf1_detection_zoo.md). If you already have a trained model then you can use the same here."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IhwY3bSIYhbm"
      },
      "source": [
        "#Get trained model\n",
        "!wget -q https://download.tensorflow.org/models/object_detection/mask_rcnn_inception_resnet_v2_atrous_coco_2018_01_28.tar.gz"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EHsX5ZkvZ6pJ"
      },
      "source": [
        "!ls -l"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ujeoCooOZ4Fo"
      },
      "source": [
        "#Unzip the file\n",
        "!tar -xf mask_rcnn_inception_resnet_v2_atrous_coco_2018_01_28.tar.gz"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KrdlbRHTaBiz"
      },
      "source": [
        "#check the unzipped files in the folder\n",
        "!ls -l mask_rcnn_inception_resnet_v2_atrous_coco_2018_01_28"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ML-49aao-ZKP"
      },
      "source": [
        "!ls -l mask_rcnn_inception_resnet_v2_atrous_coco_2018_01_28/saved_model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y2L_YpSmaKh-"
      },
      "source": [
        "When we export the model (after training), we will get same set of files."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JL-VU8HwabXj"
      },
      "source": [
        "#### Load Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8NdJs8LqSAJb"
      },
      "source": [
        "#This code will work with tf 2.x\n",
        "import tensorflow as tf"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PxNGLv_xWhaO"
      },
      "source": [
        "#Check the tf version\n",
        "tf.__version__"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "by5hnBF5R-ev"
      },
      "source": [
        "#In tf2, we will use saved model rather than frozen_inference_graph.pb\n",
        "model = tf.saved_model.load('mask_rcnn_inception_resnet_v2_atrous_coco_2018_01_28/saved_model')\n",
        "model = model.signatures['serving_default']"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ep4Xz-sxhCmC"
      },
      "source": [
        "#Check model's input\n",
        "model.inputs"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Pti2aYT-hNqn"
      },
      "source": [
        "Here the model input tensor's name is 'image_tensor' and it has a 4D shape (first dimension is for batch size i.e how many images we will feed)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rj0QqON5hqwO"
      },
      "source": [
        "#Check model's output tensors\n",
        "model.outputs"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wHjO7bcahnYi"
      },
      "source": [
        "How to understand 4 outputs here:\n",
        "\n",
        "\n",
        "\n",
        "1.   **num_detections** : Number of prediction boxes we are getting from our model. We limit number of predictions in model configuration file. In this model, output will have top 100 predictions (out of 1000s of anchor boxes). Please note that we get this output after Non-Maximum supression (NMS) step has been completed.\n",
        "2.   **detection_classes** : Index of the class with highest probability for each predicted box. These index values should be matched with index created using Label Encoder during training time. The values will between 1 to number of classes.\n",
        "3.   **detection_scores** : Probability value for highest probability class for each box. The value will be between 0 to 1. This indicates how confident model of a real object in the box.\n",
        "4.   **detection_boxes** : Boundary box co-ordinates for each predicted box. For each predicted box, we get 4 outputs i.e ymin, xmin, ymax, xmax. Please note that these are normalized values.\n",
        "\n",
        "5. **detection_masks** : Masks for each detected box\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9X68OdfVdXUk"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q0pTF60YSNxZ"
      },
      "source": [
        "#### Load Class labels"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f11_p-MBSQtp"
      },
      "source": [
        "Label dictionary (class index to class name mapping) should be taken from training module. Here is the dictionary which was used for this model."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LszBHGAISPka"
      },
      "source": [
        "all_classes = {1 : 'person' , 2: 'bicycle' , 3: 'car' , 4: 'motorcycle' , 5: 'airplane' , 6: 'bus' , 7: 'train' , \n",
        "               8: 'truck' , 9: 'boat' , 10: 'traffic light' , 11: 'fire hydrant' , 13: 'stop sign' , 14: 'parking meter' , \n",
        "               15: 'bench' , 16: 'bird' , 17: 'cat' , 18: 'dog' , 19: 'horse' , 20: 'sheep' , 21: 'cow' , 22: 'elephant' , \n",
        "               23: 'bear' , 24: 'zebra' , 25: 'giraffe' , 27: 'backpack' , 28: 'umbrella' , 31: 'handbag' , 32: 'tie' , \n",
        "               33: 'suitcase' , 34: 'frisbee' , 35: 'skis' , 36: 'snowboard' , 37: 'sports ball' , 38: 'kite' , \n",
        "               39: 'baseball bat' , 40: 'baseball glove' , 41: 'skateboard' , 42: 'surfboard' , 43: 'tennis racket' , \n",
        "               44: 'bottle' , 46: 'wine glass' , 47: 'cup' , 48: 'fork' , 49: 'knife' , 50: 'spoon' , 51: 'bowl' , \n",
        "               52: 'banana' , 53: 'apple' , 54: 'sandwich' , 55: 'orange' , 56: 'broccoli' , 57: 'carrot' , 58: 'hot dog' , \n",
        "               59: 'pizza' , 60: 'donut' , 61: 'cake' , 62: 'chair' , 63: 'couch' , 64: 'potted plant' , 65: 'bed' , \n",
        "               67: 'dining table' , 70: 'toilet' , 72: 'tv' , 73: 'laptop' , 74: 'mouse' , 75: 'remote' , 76: 'keyboard' , \n",
        "               77: 'cell phone' , 78: 'microwave' , 79: 'oven' , 80: 'toaster' , 81: 'sink' , 82: 'refrigerator' , \n",
        "               84: 'book' , 85: 'clock' , 86: 'vase' , 87: 'scissors' , 88: 'teddy bear' , 89: 'hair drier' , 90: 'toothbrush'}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XZP_JtLiVXEf"
      },
      "source": [
        "all_classes[15]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PALYIrj9sSrN"
      },
      "source": [
        "#### Model Prediction"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uQsu9BPC6RFK"
      },
      "source": [
        "def reframe_box_masks_to_image_masks(box_masks, boxes, image_height, image_width):\n",
        "  \"\"\"\n",
        "  Code taken from TensorFlow Object Detection API --> utils/visualization_utils.py \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 tf.float32 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",
        "  Returns:\n",
        "    A tf.float32 tensor of size [num_masks, image_height, image_width].\n",
        "  \"\"\"\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",
        "      transformed_boxes = (boxes - min_corner) / (max_corner - min_corner)\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",
        "    return 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",
        "        extrapolation_value=0.0)\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], dtype=tf.float32))\n",
        "  return tf.squeeze(image_masks, axis=3)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xnPdSqAjoayD"
      },
      "source": [
        "#Function to get predictions from a Detection model\n",
        "def detector_prediction(image_file, confidence_threshold=0.5):\n",
        "\n",
        "    \"\"\"\n",
        "    image_file: File path of the image for which prediction needs to be done\n",
        "    confidence_threshold: Minimum confidence/probability for prediction to be considered\n",
        "    \"\"\"\n",
        "    #Load image\n",
        "    img = tf.keras.preprocessing.image.load_img(image_file)\n",
        "    w, h = img.size\n",
        "    \n",
        "    #Convert to numpy array\n",
        "    img_array = tf.keras.preprocessing.image.img_to_array(img).astype('uint8')\n",
        "    #Make it a batch of one example\n",
        "    img_array = tf.expand_dims(img_array, axis=0)\n",
        "\n",
        "    #Prediction\n",
        "    output = model(img_array) #get list of tensors discussed above as output\n",
        "    \n",
        "    detection_scores = output['detection_scores'].numpy()[0] #get detection scores\n",
        "    detection_classes = output['detection_classes'].numpy()[0]\n",
        "    detection_boxes = output['detection_boxes'].numpy()[0]\n",
        "    detection_masks = output['detection_masks'].numpy()[0]\n",
        "    \n",
        "    #Select predictions for which probability is higher than confidence_threshold\n",
        "    selected_predictions = detection_scores >= confidence_threshold\n",
        "\n",
        "    selected_prediction_scores = detection_scores[selected_predictions]\n",
        "    selected_prediction_classes = detection_classes[selected_predictions]\n",
        "    selected_prediction_boxes = detection_boxes[selected_predictions]\n",
        "    selected_prediction_masks = detection_masks[selected_predictions]\n",
        "\n",
        "    #Resize masks to image size\n",
        "    detection_masks_reframed = reframe_box_masks_to_image_masks(selected_prediction_masks, \n",
        "                                                                selected_prediction_boxes,\n",
        "                                                                h, w)\n",
        "    #Make it a binary array (1 - pixel belongs to object)\n",
        "    detection_masks_reframed = tf.cast(detection_masks_reframed > 0.5,\n",
        "                                       tf.uint8)\n",
        "\n",
        "    #De-normalize box co-ordinates (multiply x-coordinates by image width and y-coords by image height)\n",
        "    img_w, img_h = img.size\n",
        "\n",
        "    for i in range(selected_prediction_boxes.shape[0]):\n",
        "        \n",
        "        selected_prediction_boxes[i,0] *= img_h #ymin * img_w\n",
        "        selected_prediction_boxes[i,1] *= img_w #xmin * img_h\n",
        "        selected_prediction_boxes[i,2] *= img_h #ymax * img_w\n",
        "        selected_prediction_boxes[i,3] *= img_w #xmax * img_h\n",
        "\n",
        "    #Make all co-ordinates as integer\n",
        "    selected_prediction_boxes= selected_prediction_boxes.astype(int)\n",
        "\n",
        "    #Convert class indexes to actual class labels\n",
        "    predicted_classes = []\n",
        "    for i in range(selected_prediction_classes.shape[0]):\n",
        "        predicted_classes.append(all_classes[int(selected_prediction_classes[i])])\n",
        "\n",
        "    #Number of predictions\n",
        "    selected_num_predictions = selected_prediction_boxes.shape[0]\n",
        "\n",
        "    return {'Total Predictions': selected_num_predictions,\n",
        "            'Scores': selected_prediction_scores, \n",
        "            'Classes': predicted_classes, \n",
        "            'Box coordinates': selected_prediction_boxes, \n",
        "            'Masks': detection_masks_reframed}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zBJxH1CUKzag"
      },
      "source": [
        "Let's download couple of images for which we will do predictions."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "I9-Rax5xLCcZ"
      },
      "source": [
        "!wget https://github.com/tensorflow/models/raw/master/research/object_detection/test_images/image1.jpg --quiet\n",
        "!wget https://github.com/tensorflow/models/raw/master/research/object_detection/test_images/image2.jpg --quiet"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "G3ebquawu-Rc"
      },
      "source": [
        "!ls -l"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gcrTicomAhP9"
      },
      "source": [
        "tf.keras.preprocessing.image.load_img('image1.jpg')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MU4klwe-tF15"
      },
      "source": [
        "tf.keras.preprocessing.image.load_img('image2.jpg')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sPbIDpAGxeBC"
      },
      "source": [
        "#Model predictions for image1.jpg\n",
        "detector_prediction('image1.jpg')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pziK2r8-MOQr"
      },
      "source": [
        "#### Visualizing Model Prediction"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jjaKo63LMf7B"
      },
      "source": [
        "import cv2\n",
        "from matplotlib import pyplot as plt\n",
        "import numpy as np"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HLfLuQFgMQgJ"
      },
      "source": [
        "def visualize_output(image_file, confidence_threshold=0.5):\n",
        "\n",
        "    #Call model prediction function above\n",
        "    output = detector_prediction(image_file, confidence_threshold=confidence_threshold)\n",
        "\n",
        "    #Read image\n",
        "    img = cv2.imread(image_file)\n",
        "    \n",
        "    #Draw rectangle for predicted boxes, also add predicted classes\n",
        "    for i in range(output['Box coordinates'].shape[0]):\n",
        "\n",
        "        #Mask\n",
        "        mask = output['Masks'][i]\n",
        "\n",
        "        #Draw mask\n",
        "        colored_mask = np.expand_dims(mask, axis=2) * np.reshape([255, 255, 0], (1,1,3))\n",
        "        img = cv2.addWeighted(colored_mask.astype('uint8'), 0.5, img, 1.0, 0.0)\n",
        "\n",
        "        #Bounding box\n",
        "        box = output['Box coordinates'][i]\n",
        "        \n",
        "        #Draw rectangle \n",
        "        img = cv2.rectangle(img, (box[1], box[0]), (box[3], box[2]), (0,255,0), 2)\n",
        "        \n",
        "        #Add Label - Class name and confidence level\n",
        "        label = output['Classes'][i] + ': ' + str(round(output['Scores'][i],2))\n",
        "        img = cv2.putText(img, label, (box[1], box[0]-10), cv2.FONT_HERSHEY_SIMPLEX, 1, (255,255,255), 2, cv2.LINE_AA)\n",
        "    \n",
        "    #Conver BGR image to RGB to use with Matplotlib\n",
        "    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
        "\n",
        "    #Display image\n",
        "    plt.figure(figsize=(12,8))\n",
        "    plt.imshow(img)\n",
        "    plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6c0umKw2O74F"
      },
      "source": [
        "#Visualize first image\n",
        "visualize_output('image1.jpg', confidence_threshold=0.9)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Sv06rWVDsVZH"
      },
      "source": [
        "#Visualize second image\n",
        "visualize_output('image2.jpg')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mf_orYdYQrIW"
      },
      "source": [
        "#### Extract Object(s) from an Image"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c5e12TCOR_d2"
      },
      "source": [
        "Function to extract object(s) based on predicted boxes"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0RcILkJDIVZs"
      },
      "source": [
        "def extract_object_from_box(image_file, object_name, confidence_threshold=0.5):\n",
        "\n",
        "    #Call model prediction function above\n",
        "    output = detector_prediction(image_file, confidence_threshold=confidence_threshold)\n",
        "\n",
        "    #Read image\n",
        "    img = cv2.imread(image_file)\n",
        "\n",
        "    cropped_images = []\n",
        "\n",
        "    #Draw rectangle for predicted boxes, also add predicted classes\n",
        "    for i in range(output['Box coordinates'].shape[0]):\n",
        "\n",
        "        #Class\n",
        "        label = output['Classes'][i]\n",
        "\n",
        "        if label.lower() == object_name.lower():\n",
        "\n",
        "            #Bounding box\n",
        "            box = output['Box coordinates'][i]\n",
        "\n",
        "            #Crop image\n",
        "            cropped_img = img[box[0]:box[2], box[1]:box[3]]\n",
        "            cropped_images.append(cropped_img)\n",
        "        \n",
        "    \n",
        "    if len(cropped_images) == 0:\n",
        "        print('No ' + object_name + ' found')\n",
        "        return\n",
        "    \n",
        "    print(\"Number of '\" + object_name + \" found:\", len(cropped_images))\n",
        "    plt.figure(figsize=(15,10))\n",
        "\n",
        "    for i in range(len(cropped_images)):\n",
        "        \n",
        "        plt.subplot(1, len(cropped_images), i+1)\n",
        "        plt.imshow(cv2.cvtColor(cropped_images[i], cv2.COLOR_BGR2RGB))\n",
        "    \n",
        "    plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "D2tK3CiTKRQM"
      },
      "source": [
        "extract_object_from_box('image2.jpg', 'person')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QzlA5j_LSFEw"
      },
      "source": [
        "Function to extract object(s) based on the mask"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0Lxo97GtLXi7"
      },
      "source": [
        "def extract_object_from_mask(image_file, object_name, confidence_threshold=0.5):\n",
        "\n",
        "    #Call model prediction function above\n",
        "    output = detector_prediction(image_file, confidence_threshold=confidence_threshold)\n",
        "\n",
        "    #Read image\n",
        "    img = cv2.imread(image_file)\n",
        "\n",
        "    cropped_images = []\n",
        "\n",
        "    #Draw rectangle for predicted boxes, also add predicted classes\n",
        "    for i in range(output['Box coordinates'].shape[0]):\n",
        "\n",
        "        #Class\n",
        "        label = output['Classes'][i]\n",
        "\n",
        "        if label.lower() == object_name.lower():\n",
        "\n",
        "            #Mask\n",
        "            mask = output['Masks'][i].numpy()\n",
        "            mask = np.repeat(np.expand_dims(mask, axis=2), 3, axis=-1)*255\n",
        "\n",
        "            \n",
        "            #Crop image\n",
        "            cropped_img = cv2.bitwise_and(img, mask)\n",
        "            cropped_images.append(cropped_img)\n",
        "        \n",
        "    \n",
        "    if len(cropped_images) == 0:\n",
        "        print('No ' + object_name + ' found')\n",
        "        return\n",
        "    \n",
        "    print(\"Number of '\" + object_name + \" found:\", len(cropped_images))\n",
        "    plt.figure(figsize=(15,10))\n",
        "\n",
        "    for i in range(len(cropped_images)):\n",
        "        \n",
        "        plt.subplot(1, len(cropped_images), i+1)\n",
        "        plt.imshow(cv2.cvtColor(cropped_images[i], cv2.COLOR_BGR2RGB))\n",
        "    \n",
        "    plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LRRRIg5_CShK"
      },
      "source": [
        "extract_object_from_mask('image1.jpg', 'dog')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VmZRNcgCRLjz"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}