prediction_into_xml.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9725fda4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2 \n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from PIL import Image\n",
    "import os\n",
    "import xml.etree.ElementTree as ET\n",
    "\n",
    "CONFIDENCE_THRESHOLD = 0.4\n",
    "NMS_THRESHOLD = 0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1a99993a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def indent(elem, level=0):\n",
    "    i = \"\\n\" + level*\"  \"\n",
    "    j = \"\\n\" + (level-1)*\"  \"\n",
    "    if len(elem):\n",
    "        if not elem.text or not elem.text.strip():\n",
    "            elem.text = i + \"  \"\n",
    "        if not elem.tail or not elem.tail.strip():\n",
    "            elem.tail = i\n",
    "        for subelem in elem:\n",
    "            indent(subelem, level+1)\n",
    "        if not elem.tail or not elem.tail.strip():\n",
    "            elem.tail = j\n",
    "    else:\n",
    "        if level and (not elem.tail or not elem.tail.strip()):\n",
    "            elem.tail = j\n",
    "    return elem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3e980caa",
   "metadata": {},
   "outputs": [],
   "source": [
    "To_detect = []\n",
    "with open(\"configuration/Annotations.names\", \"r\") as f:\n",
    "    To_detect = [cname.strip() for cname in f.readlines()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "7d1ee844",
   "metadata": {},
   "outputs": [],
   "source": [
    "class_names = []\n",
    "with open(\"configuration/coco.names\", \"r\") as f:\n",
    "    class_names = [cname.strip() for cname in f.readlines()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0e0b97cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "net = cv2.dnn.readNet(\"configuration/yolov4-csp.weights\", \"configuration/yolov4-csp.cfg\")\n",
    "net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)\n",
    "net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA_FP16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4072b201",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = cv2.dnn_DetectionModel(net)\n",
    "model.setInputParams(size=(640, 640), scale=1/255, swapRB=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "6c669119",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'C:\\\\Users\\\\Gaurav\\\\Desktop\\\\Automatic Annotation tool\\\\images'"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "path = os.path.abspath('images')\n",
    "path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "042abd76",
   "metadata": {},
   "outputs": [],
   "source": [
    "for file in os.listdir(path):\n",
    "    #root element\n",
    "    root = ET.Element('annotation')\n",
    "    colors = np.random.uniform(0,255,size=(len(class_names),3))\n",
    "\n",
    "\n",
    "    #book sub-element\n",
    "    folder = ET.SubElement(root, 'folder')\n",
    "    folder.text='Images'\n",
    "\n",
    "    filename = ET.SubElement(root, 'filename')\n",
    "    filename.text=file\n",
    "\n",
    "    size = ET.SubElement(root, 'size')\n",
    "\n",
    "    width = ET.SubElement(size, 'width')\n",
    "    width.text = '416'\n",
    "    height = ET.SubElement(size, 'height')\n",
    "    height.text = '416'\n",
    "    depth = ET.SubElement(size, 'depth')\n",
    "    depth.text = '3'\n",
    "\n",
    "    segmented = ET.SubElement(root, 'segmented')\n",
    "    segmented.text='0'\n",
    "    print(file)\n",
    "    img = cv2.imread(os.path.join(path, file) )\n",
    "    classes, scores, boxes = model.detect(img, CONFIDENCE_THRESHOLD, NMS_THRESHOLD)\n",
    "    for (classid, score, box) in zip(classes, scores, boxes):\n",
    "        if class_names[classid[0]] in detect_classes:\n",
    "            color=colors[classid[0]]\n",
    "            label = \"%s : %f\" % (class_names[classid[0]], score)\n",
    "            cv2.rectangle(img, box,color, 2)\n",
    "            cv2.putText(img, label, (box[0], box[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5,(255,255,255), 2)\n",
    "            object = ET.SubElement(root, 'object')\n",
    "\n",
    "            name = ET.SubElement(object, 'name')\n",
    "            name.text=class_names[classid[0]]\n",
    "\n",
    "            pose = ET.SubElement(object, 'pose')\n",
    "            pose.text='Unspecified'\n",
    "\n",
    "            turncated = ET.SubElement(object, 'turncated')\n",
    "            turncated.text='0'\n",
    "\n",
    "            occluded = ET.SubElement(object, 'occluded')\n",
    "            occluded.text='0'\n",
    "\n",
    "            difficult = ET.SubElement(object, 'difficult')\n",
    "            difficult.text='0'\n",
    "\n",
    "            bndbox = ET.SubElement(object, 'bndbox')\n",
    "\n",
    "            xmin = ET.SubElement(bndbox, 'xmin')\n",
    "            xmin.text = str(box[0])\n",
    "\n",
    "            ymin = ET.SubElement(bndbox, 'ymin')\n",
    "            ymin.text = str(box[1])\n",
    "\n",
    "            xmax = ET.SubElement(bndbox, 'xmax')\n",
    "            xmax.text = str(box[0]+box[2])\n",
    "\n",
    "            ymax = ET.SubElement(bndbox, 'ymax')\n",
    "            ymax.text = str(box[1]+box[3])\n",
    "\n",
    "            tree = ET.ElementTree(indent(root))\n",
    "            tree.write('annotations/'+file[0:-3:]+'xml', xml_declaration=True, encoding='utf-8')\n",
    "\n",
    "        \n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5b563aa",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "yolo2",
   "language": "python",
   "name": "yolo2"
  },
  "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.6.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}