Abhishek_Object_detection.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HpS-cWaoLMwT"
      },
      "outputs": [],
      "source": [
        "import cv2\n",
        "import argparse\n",
        "import numpy as np"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "ap = argparse.ArgumentParser()\n",
        "ap.add_argument('-i', '--image', required=True,\n",
        "                help = 'path to input image')\n",
        "ap.add_argument('-c', '--config', required=True,\n",
        "                help = 'path to yolo config file')\n",
        "ap.add_argument('-w', '--weights', required=True,\n",
        "                help = 'path to yolo pre-trained weights')\n",
        "ap.add_argument('-cl', '--classes', required=True,\n",
        "                help = 'path to text file containing class names')\n",
        "args = ap.parse_args()\n"
      ],
      "metadata": {
        "id": "78wz2wSnLNdV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def get_output_layers(net):\n",
        "\n",
        "    layer_names = net.getLayerNames()\n",
        "    try:\n",
        "        output_layers = [layer_names[i - 1] for i in net.getUnconnectedOutLayers()]\n",
        "    except:\n",
        "        output_layers = [layer_names[i[0] - 1] for i in net.getUnconnectedOutLayers()]\n",
        "\n",
        "    return output_layers"
      ],
      "metadata": {
        "id": "2t0-9E0lLNf8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def draw_prediction(img, class_id, confidence, x, y, x_plus_w, y_plus_h):\n",
        "\n",
        "    label = str(classes[class_id])\n",
        "\n",
        "    color = COLORS[class_id]\n",
        "\n",
        "    cv2.rectangle(img, (x,y), (x_plus_w,y_plus_h), color, 2)\n",
        "\n",
        "    cv2.putText(img, label, (x-10,y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)\n",
        "\n",
        ""
      ],
      "metadata": {
        "id": "6I0uoPTwLNkK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "image = cv2.imread(args.image)\n",
        "\n",
        "Width = image.shape[1]\n",
        "Height = image.shape[0]\n",
        "scale = 0.00392\n",
        "\n",
        "classes = None\n",
        "\n",
        "with open(args.classes, 'r') as f:\n",
        "    classes = [line.strip() for line in f.readlines()]\n",
        "\n",
        "COLORS = np.random.uniform(0, 255, size=(len(classes), 3))\n",
        "\n",
        "net = cv2.dnn.readNet(args.weights, args.config)\n",
        "\n",
        "blob = cv2.dnn.blobFromImage(image, scale, (416,416), (0,0,0), True, crop=False)\n",
        "\n",
        "net.setInput(blob)\n",
        "\n",
        "outs = net.forward(get_output_layers(net))"
      ],
      "metadata": {
        "id": "ebQps0e9LNnG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class_ids = []\n",
        "confidences = []\n",
        "boxes = []\n",
        "conf_threshold = 0.5\n",
        "nms_threshold = 0.4\n",
        "\n",
        "\n",
        "for out in outs:\n",
        "    for detection in out:\n",
        "        scores = detection[5:]\n",
        "        class_id = np.argmax(scores)\n",
        "        confidence = scores[class_id]\n",
        "        if confidence > 0.5:\n",
        "            center_x = int(detection[0] * Width)\n",
        "            center_y = int(detection[1] * Height)\n",
        "            w = int(detection[2] * Width)\n",
        "            h = int(detection[3] * Height)\n",
        "            x = center_x - w / 2\n",
        "            y = center_y - h / 2\n",
        "            class_ids.append(class_id)\n",
        "            confidences.append(float(confidence))\n",
        "            boxes.append([x, y, w, h])\n"
      ],
      "metadata": {
        "id": "JwP-F1_9LNqC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "indices = cv2.dnn.NMSBoxes(boxes, confidences, conf_threshold, nms_threshold)\n"
      ],
      "metadata": {
        "id": "QVTAkm2BLNs1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for i in indices:\n",
        "    try:\n",
        "        box = boxes[i]\n",
        "    except:\n",
        "        i = i[0]\n",
        "        box = boxes[i]\n",
        "\n",
        "    x = box[0]\n",
        "    y = box[1]\n",
        "    w = box[2]\n",
        "    h = box[3]\n",
        "    draw_prediction(image, class_ids[i], confidences[i], round(x), round(y), round(x+w), round(y+h))\n",
        "\n",
        "cv2.imshow(\"object detection\", image)\n",
        "cv2.waitKey()\n",
        "\n",
        "cv2.imwrite(\"object-detection.jpg\", image)\n",
        "cv2.destroyAllWindows()"
      ],
      "metadata": {
        "id": "0fcRU7whLNvd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "rqyhSMacLN81"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}