WeaponDetectionYOLOv5.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/patriciacs99/WeaponDetectionYOLOv5/blob/main/WeaponDetectionYOLOv5.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7Nlu9T9k8E5u"
      },
      "source": [
        "# 🔫 Firearms detection with YOLOV5\n",
        "\n",
        "*Author: Patricia Corral Sanz*\n",
        "\n",
        "In this project the following code snippets have been used to detect firearms and classify them into long guns (shotguns, rifles, machine guns and snipers) and handguns (pistols and revolvers).\n",
        "\n",
        "For the experiments performed, the necessary cells will be executed for each of them depending on the objective of the experiment. In each case changing the necessary values as indicated in this notebook.\n",
        "https://drive.google.com/drive/folders/15O3lpCT-JYyuhEc5WftPzr0vCELjLPSS?usp=drive_link  \n",
        "\n",
        "\n",
        "---"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "X7alvtOhBfx5"
      },
      "outputs": [],
      "source": [
        "#imports\n",
        "import os\n",
        "import shutil\n",
        "import random\n",
        "import cv2\n",
        "import torch\n",
        "import matplotlib.pyplot as plt\n",
        "import torchvision\n",
        "import numpy as np\n",
        "import glob\n",
        "import re\n",
        "import time\n",
        "from pathlib import Path\n",
        "import itertools\n",
        "from IPython.display import Image,display\n",
        "from google.colab.patches import cv2_imshow"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Check GPU\n",
        "gpu_info = !nvidia-smi\n",
        "gpu_info = '\\n'.join(gpu_info)\n",
        "if gpu_info.find('failed') >= 0:\n",
        "  print('Not connected to a GPU')\n",
        "else:\n",
        "  print(gpu_info)"
      ],
      "metadata": {
        "id": "uN0ORk4dXP4Y"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9_O69Z_qfIa8"
      },
      "source": [
        "##Prepare data and Yolov5\n",
        "\n",
        "\n",
        "\n",
        ">In this section we will access Google Drive, where we will have organized the datasets to be used in the experiments.   \n",
        "We will clone the Ultralytics YOLOv5 repository to be able to work with it and install its requirements.  \n",
        "We will also download the YOLOv5 models with weights pre-trained with the COCO dataset we are going to use. In our case, we will only use YOLOv5s and YOLOv5m.\n",
        "\n",
        "⚠ *As a result of the execution of the YOLOv5 files (train.py,detect.py and val.py) a \"runs\" folder will be generated in the YOLOv5 directory where we will find the results of the executions.*\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "q9QN-80oVxP8"
      },
      "outputs": [],
      "source": [
        "#Connecting Google Drive\n",
        "\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive',force_remount=True)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cU799ikTW_fx"
      },
      "outputs": [],
      "source": [
        "# Clone YOLOv5\n",
        "%cd content\n",
        "!git clone https://github.com/ultralytics/yolov5.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YLulDOv8Q3hI"
      },
      "outputs": [],
      "source": [
        "#Install YOLOv5\n",
        "%cd /content/yolov5\n",
        "%pip install -qr requirements.txt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Wb6of7pXXIU0"
      },
      "outputs": [],
      "source": [
        "# Download pre trained models\n",
        "'''\n",
        "  Modelos a probar:\n",
        "    - yolov5s.pt\n",
        "    - yolov5m.pt\n",
        "    - yolov5l.pt\n",
        "    - yolov5x.pt\n",
        "'''\n",
        "# Replace \"yolov5m.pt\" with the preferred model name\n",
        "!wget https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5m.pt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KJZ8-j-pHONy"
      },
      "source": [
        "## 👟 Model training\n",
        "\n",
        "\n",
        "> We will train the model for our particular case: detecting short and long firearms.  \n",
        "In the following cells, they should be replaced by their corresponding routes:\n",
        "*   **conf_file_path**: path to the yaml file of the dataset used.\n",
        "*   **weigths_model_YOLOv5_path**: path to the model with the pre-trained weights previously downloaded.\n",
        "*   **results_path**: path to the directory where the results have been generated.\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ACbed_Re711-"
      },
      "outputs": [],
      "source": [
        "!python train.py --img 416 --batch 32 --epochs 100 --data conf_file_path --weights weigths_model_YOLOv5_path   --name \"resultados_entrenamiento\" --nosave"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CQkYUq41syHy"
      },
      "outputs": [],
      "source": [
        "#Download training results\n",
        "from google.colab import files\n",
        "\n",
        "!zip -r train_results_yolov5.zip results_path/resultados_entrenamiento\n",
        "files.download('train_results_yolov5.zip')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V0L0b39uM2--"
      },
      "source": [
        "### Tensorboard Training Results\n",
        "\n",
        "\n",
        "> We will run Tensorboard to visualize the graphs resulting from the training.  \n",
        "Replace:\n",
        "*   **training_path**: path of the directory where the training results have been generated.\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "89rT52JvM56x"
      },
      "outputs": [],
      "source": [
        "%reload_ext tensorboard\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NI2vQ-xn3j7N"
      },
      "outputs": [],
      "source": [
        "%tensorboard --logdir training_path\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IkO_f-goz8_i"
      },
      "source": [
        "## 🎯 Inference\n",
        "\n",
        "> The weights of the trained model will be used to detect and classify weapons in new test images that the model has not seen before.  In the following cells, they should be replaced by their corresponding paths:\n",
        "*   **test_images_path**:path to the test images of the dataset.\n",
        "*   **weights_path**: path to weights generated in training.\n",
        "*   **detections_path**: path to the directory where the results of the inference have been generated.\n",
        "*   **results_path**: path to the directory where the results have been generated.\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xehIJ1w5z8U5"
      },
      "outputs": [],
      "source": [
        "!python detect.py --source test_images_path  --weights weights_path --img 416 --save-conf  --save-txt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YX343_Ub0M0C"
      },
      "outputs": [],
      "source": [
        "#Show test images detections\n",
        "for imageName in glob.glob(\"detections_path/*.jpg\"):\n",
        "  display(Image(filename=imageName))\n",
        "  print(\"\\n\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "k1gstUP3-JTF"
      },
      "outputs": [],
      "source": [
        "#Download results\n",
        "from google.colab import files\n",
        "\n",
        "!zip -r train_detect_yolov5.zip results_path\n",
        "files.download('train_detect_yolov5.zip')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8ZBV-H0IPmOA"
      },
      "source": [
        "## 📈 Evaluation\n",
        "\n",
        "\n",
        ">   \n",
        "\n",
        "> We will extract the metrics of the model on the new test images in order to evaluate its performance.  \n",
        "In the following cells, they should be replaced by their corresponding paths:\n",
        "*   **conf_file_path**: path to the yaml file of the dataset used.\n",
        "*   **weights_path**:path where the file with the weights resulting from the training of the model is located.\n",
        "*   **results_path**: path to the directory where the results of the execution of the val.py file have been generated.\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7CXnggn7kPrP"
      },
      "outputs": [],
      "source": [
        "!python val.py --weights weights_path --data conf_file_path --workers 0 --img 416 --save-txt --conf-thres 0.5 --task test"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RFsdEvaquFum"
      },
      "outputs": [],
      "source": [
        "#Download results\n",
        "from google.colab import files\n",
        "\n",
        "!zip -r train_metrics_yolov5.zip results_path\n",
        "files.download('train_metrics_yolov5.zip')"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "---\n",
        "\n",
        "## 📑 Additional code\n",
        "\n",
        "\n",
        "> The following code snippets have been useful for the elaboration of the experiments.\n",
        "\n"
      ],
      "metadata": {
        "id": "eLLZhBhyNieQ"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bvrh4RF680Lw"
      },
      "source": [
        "### Intersection Over Union\n",
        "\n",
        "\n",
        ">Although YOLO calculates the IoU itself, here the intersections are calculated for each image in order to draw them on each image.   \n",
        "When executing the following cells, the intersection between the real label and the detection is calculated, and the Ground Truth is drawn on the image containing the detection in order to clearly visualize the difference.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b-pz6NJusKke"
      },
      "outputs": [],
      "source": [
        " #Format: x,y,w,h  ---> Format: x1,y1,x2,y2\n",
        "def rectangulo(box):\n",
        "    box_x1 = box[0] - box[2]/2\n",
        "    box_y1 = box[1] - box[3]/2\n",
        "    box_x2 = box[0] + box[2]/2\n",
        "    box_y2 = box[1] + box[3]/2\n",
        "    return [box_x1,box_y1,box_x2,box_y2]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "p86zjRKj9_pO"
      },
      "outputs": [],
      "source": [
        "def intersection_over_union(target_bbox,predicted_bbox):\n",
        "\n",
        "    #Format: x,y,w,h\n",
        "\n",
        "    #Calculate corners bboxes\n",
        "    tg = rectangulo(target_bbox)\n",
        "    pred = rectangulo(predicted_bbox)\n",
        "\n",
        "    #Calculate corners interseccion\n",
        "    x1 = max(pred[0], tg[0])\n",
        "    y1 = max(pred[1], tg[1])\n",
        "    x2 = min(pred[2], tg[2])\n",
        "    y2 = min(pred[3], tg[3])\n",
        "\n",
        "    #area = w * h\n",
        "    intersection = max(0,(x2 - x1)) * max(0,(y2 - y1))\n",
        "\n",
        "    box1_area = abs((pred[2] - pred[0]) * (pred[3] - pred[1]))\n",
        "    box2_area = abs((tg[2] - tg[0]) * (tg[3] - tg[1]))S\n",
        "\n",
        "    return intersection / (box1_area + box2_area - intersection + 1e-16)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1C05eBa8oDk4"
      },
      "outputs": [],
      "source": [
        "#Draw detections,IoU,groundTruth\n",
        "\n",
        "predicted = \"/content/yolov5/runs/detect/exp/labels/0205_jpg.rf.eb76fc2a71b61cec8cdf25fd3e51c6cb.txt\"\n",
        "real = \"/content/drive/MyDrive/Deteccion_armas_TFG/Datasets/Customized/WeaponsDataset.v8-weapons_strectch_oclussionsomited.yolov5pytorch/test/images/0205_jpg.rf.eb76fc2a71b61cec8cdf25fd3e51c6cb.jpg\"\n",
        "\n",
        "#Read detection image\n",
        "image = cv2.imread(\"/content/drive/MyDrive/Deteccion_armas_TFG/Datasets/weapons.v2-weapons_blackedges_contrast.yolov5pytorch/test/images/0251_jpg.rf.1d75a35dbae43d7e1f13a8266a5d8726.jpg\")\n",
        "\n",
        "with open(real) as archivo: # Open gt coordenates file\n",
        "    for linea in archivo:\n",
        "      coord1 =[] # file coordenates storing list\n",
        "      clase = linea[0] # getting detection class - first character of the read line\n",
        "      coord1.append(linea[1:].split()) # Removing first character and adding remaining coordenates\n",
        "      coord1 =list(itertools.chain(*coord1)) # Flattening list\n",
        "      x=0\n",
        "      for i in coord1: # Casting values\n",
        "            coord1[x]= float(i)\n",
        "            x +=1\n",
        "      coords= [] # Coordinates of all possible detections will have to be saved for comparison with the gt.\n",
        "      if os.path.exists(predicted):\n",
        "        with open(predicted) as archivo:\n",
        "          for linea in archivo:\n",
        "                coord2 =[]\n",
        "                coord2.append(linea[1:].split())\n",
        "                coord2= list(itertools.chain(*coord2))\n",
        "                x=0\n",
        "                for i in coord2:\n",
        "                        coord2[x]= float(i)\n",
        "                        x +=1\n",
        "                coords.append(coord2)\n",
        "\n",
        "        ious = [] # Storing ious\n",
        "        for i in coords:\n",
        "            ious.append(intersection_over_union(i,coord1))\n",
        "        # Selecting larger IoU value detection\n",
        "        mayor = ious.index(max(ious))\n",
        "        deteccion = coords[mayor]\n",
        "\n",
        "        pred = rectangulo(deteccion)\n",
        "        gt = rectangulo(coord1)\n",
        "        iou = max(ious)\n",
        "\n",
        "      else:\n",
        "        pred = rectangulo([0,0,0,0])\n",
        "        gt = rectangulo(coord1)\n",
        "        iou = 0.000\n",
        "\n",
        "      for i in range(0,4):\n",
        "          gt[i] = int(gt[i]*416)\n",
        "\n",
        "      start_point = tuple(gt[:2])\n",
        "      end_point = tuple(gt[2:])\n",
        "\n",
        "      #Print rectangles on detection image\n",
        "      cv2.rectangle(image,start_point,end_point,(0, 255, 0), 4)\n",
        "\n",
        "      cv2.putText(image, \"IoU: {:.3f}\".format(iou), (gt[0]-10, gt[1]-10),\n",
        "            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)\n",
        "      if iou !=0:\n",
        "        cv2.putText(image, \"Class: \"+ clase, (gt[0]+105, gt[1]- 10),\n",
        "              cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255,0, 0), 2)\n",
        "\n",
        "# Show image with IoU value, detection and ground truth\n",
        "cv2_imshow(image)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Size of image instances\n",
        "\n",
        "\n",
        "> Code to divide test images according to their size used in experiment 2.\n",
        "\n"
      ],
      "metadata": {
        "id": "BRwffXFufBeC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Create 6 directories\n",
        "os.makedirs(\"dir_peq/images\")\n",
        "os.makedirs(\"dir_peq/labels\")\n",
        "os.makedirs(\"dir_med/images\")\n",
        "os.makedirs(\"dir_med/labels\")\n",
        "os.makedirs(\"dir_gra/images\")\n",
        "os.makedirs(\"dir_gra/labels\")\n",
        "os.makedirs(\"dir_mez/images\")\n",
        "os.makedirs(\"dir_mez/labels\")"
      ],
      "metadata": {
        "id": "QbKYhCGSuMoH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "For images containing only one weapon:"
      ],
      "metadata": {
        "id": "0A9DJbJGwdc-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "dir_origen = \"/content/drive/MyDrive/Deteccion_armas_TFG/Datasets/weapons.v2-weapons_blackedges_contrast.yolov5pytorch/test/images/\"\n",
        "\n",
        "# images path\n",
        "imagenes = glob.glob(\"/content/drive/MyDrive/Deteccion_armas_TFG/Datasets/weapons.v2-weapons_blackedges_contrast.yolov5pytorch/test/images/*.jpg\")\n",
        "\n",
        "dir_gra = \"/content/dir_gra\"\n",
        "dir_med = \"/content/dir_med\"\n",
        "dir_peq = \"/content/dir_peq\"\n",
        "\n",
        "\n",
        "for img in imagenes:\n",
        "  image = cv2.imread(img)\n",
        "  # Image size (PIXELS)\n",
        "  w = image.shape[0]\n",
        "  h = image.shape[1]\n",
        "  image_size = w*h\n",
        "  real = img.replace(\"/images/\",\"/labels/\",1)\n",
        "  real = real[:-3] + \"txt\"\n",
        "  img_name = re.sub(dir_origen,'',img)\n",
        "  img_label = img_name[:-3] + \"txt\"\n",
        "  print(img_label)\n",
        "  print(img_name)\n",
        "  with open(real, 'r+') as archivo:\n",
        "    lista = archivo.readlines()\n",
        "    num = len(lista)\n",
        "    if num == 1:\n",
        "        linea = lista[0]\n",
        "\n",
        "        coord1 =[]\n",
        "        clase = linea[0]\n",
        "        coord1.append(linea[1:].split())\n",
        "        coord1 =list(itertools.chain(*coord1))\n",
        "        x=0\n",
        "        for i in coord1:\n",
        "              coord1[x]= float(i)\n",
        "              x +=1\n",
        "        # Selecting height and wide values from the coordenates\n",
        "        w1 = coord1[2] * 416\n",
        "        h1 = coord1[3] * 416\n",
        "        # Calculating instance size using those values\n",
        "        instance_size = w1*h1\n",
        "\n",
        "        # Calculating percentage of the image ocuppied\n",
        "        percentage_size = instance_size * 100 / image_size\n",
        "        print(\"La instancia ocupa el \"+ str(percentage_size) +\"% de la imagen\")\n",
        "\n",
        "        if 0 <= percentage_size <= 15:\n",
        "          shutil.copy(img, \"dir_peq/images/\"+ img_name)\n",
        "          shutil.copy(real,\"dir_peq/labels/\"+ img_label)\n",
        "          print(\"Instancia peque\")\n",
        "        elif 16 <= percentage_size <= 35:\n",
        "          shutil.copy(img, \"dir_med/images/\"+  img_name)\n",
        "          shutil.copy(real,\"dir_med/labels/\"+ img_label)\n",
        "          print(\"Instancia mediana\")\n",
        "        else:\n",
        "          shutil.copy(img, \"dir_gra/images/\"+  img_name)\n",
        "          shutil.copy(real,\"dir_gra/labels/\"+ img_label)\n",
        "          print(\"Instancia grande\")\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "xcMVjaSDfGVP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "For images containing one or more than one weapon:"
      ],
      "metadata": {
        "id": "2yHS9YvfwUh0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "dir_origen = \"/content/drive/MyDrive/WeaponDetectionYOLOv5/Experimento_2/dataset/test/images/\"\n",
        "\n",
        "imagenes = glob.glob(\"/content/drive/MyDrive/WeaponDetectionYOLOv5/Experimento_2/dataset/test/images/*.jpg\")\n",
        "\n",
        "dir_gra = \"/content/dir_gra\"\n",
        "dir_med = \"/content/dir_med\"\n",
        "dir_peq = \"/content/dir_peq\"\n",
        "dir_mez = \"/content/dir_mez\"\n",
        "\n",
        "\n",
        "for img in imagenes:\n",
        "  image = cv2.imread(img)\n",
        "\n",
        "  w = image.shape[0]\n",
        "  h = image.shape[1]\n",
        "  image_size = w*h\n",
        "  real = img.replace(\"/images/\",\"/labels/\",1)\n",
        "  real = real[:-3] + \"txt\"\n",
        "  img_name = re.sub(dir_origen,'',img)\n",
        "  img_label = img_name[:-3] + \"txt\"\n",
        "  print(img_name)\n",
        "  with open(real, 'r+') as archivo:\n",
        "    lista = archivo.readlines()\n",
        "    num = len(lista)\n",
        "    contador_p = 0\n",
        "    contador_m = 0\n",
        "    contador_g = 0\n",
        "    for linea in lista:\n",
        "        coord1 =[]\n",
        "        clase = linea[0]\n",
        "        coord1.append(linea[1:].split())\n",
        "        coord1 =list(itertools.chain(*coord1))\n",
        "        x=0\n",
        "        for i in coord1:\n",
        "              coord1[x]= float(i)\n",
        "              x +=1\n",
        "\n",
        "        w1 = coord1[2] * 416\n",
        "        h1 = coord1[3] * 416\n",
        "\n",
        "        instance_size = w1*h1\n",
        "\n",
        "\n",
        "        percentage_size = instance_size * 100 / image_size\n",
        "        print(\"La instancia ocupa el \"+ str(percentage_size) +\"% de la imagen\")\n",
        "\n",
        "        # Adding the number of instances of each size\n",
        "        if 0 <= percentage_size <= 15:\n",
        "          contador_p = contador_p + 1\n",
        "        elif 16 <= percentage_size <= 35:\n",
        "          contador_m = contador_m + 1\n",
        "        else:\n",
        "          contador_g = contador_g + 1\n",
        "\n",
        "\n",
        "    # Assign a specific directory\n",
        "\n",
        "    if contador_p == num :\n",
        "          shutil.copy(img, \"dir_peq/images/\"+ img_name)\n",
        "          shutil.copy(real,\"dir_peq/labels/\"+ img_label)\n",
        "          print(\"PEQUEÑO\")\n",
        "    elif contador_m == num:\n",
        "          shutil.copy(img, \"dir_med/images/\"+  img_name)\n",
        "          shutil.copy(real,\"dir_med/labels/\"+ img_label)\n",
        "          print(\"MEDIANO\")\n",
        "    elif contador_g == num:\n",
        "          shutil.copy(img, \"dir_gra/images/\"+  img_name)\n",
        "          shutil.copy(real,\"dir_gra/labels/\"+ img_label)\n",
        "          print(\"GRANDE\")\n",
        "    else:\n",
        "          shutil.copy(img, \"dir_mez/images/\"+  img_name)\n",
        "          shutil.copy(real,\"dir_mez/labels/\"+ img_label)\n",
        "          print(\"MEZCLA\")\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "4Un7MQQpvyRy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Analyze images with size mixing in instances\n",
        "dir_origen = \"/content/dir_mez/images\"\n",
        "imagenes = glob.glob(\"/content/dir_mez/images/*.jpg\")\n",
        "\n",
        "for img in imagenes:\n",
        "  image = cv2.imread(img)\n",
        "\n",
        "  w = image.shape[0]\n",
        "  h = image.shape[1]\n",
        "  image_size = w*h\n",
        "  real = img.replace(\"/images/\",\"/labels/\",1)\n",
        "  real = real[:-3] + \"txt\"\n",
        "  img_name = re.sub(dir_origen,'',img)\n",
        "  img_label = img_name[:-3] + \"txt\"\n",
        "  print(img_name)\n",
        "  with open(real, 'r+') as archivo:\n",
        "    lista = archivo.readlines()\n",
        "    num = len(lista)\n",
        "    contador_p = 0\n",
        "    contador_m = 0\n",
        "    contador_g = 0\n",
        "    for linea in lista:\n",
        "        coord1 =[]\n",
        "        clase = linea[0]\n",
        "        coord1.append(linea[1:].split())\n",
        "        coord1 =list(itertools.chain(*coord1))\n",
        "        x=0\n",
        "        for i in coord1:\n",
        "              coord1[x]= float(i)\n",
        "              x +=1\n",
        "\n",
        "        w1 = coord1[2] * 416\n",
        "        h1 = coord1[3] * 416\n",
        "\n",
        "        instance_size = w1*h1\n",
        "\n",
        "\n",
        "        percentage_size = instance_size * 100 / image_size\n",
        "        print(\"La instancia ocupa el \"+ str(percentage_size) +\"% de la imagen\")\n",
        "\n",
        "\n",
        "        if 0 <= percentage_size <= 15:\n",
        "          print(\"Instancia peque\")\n",
        "          print(linea)\n",
        "        elif 16 <= percentage_size <= 35:\n",
        "          print(\"Instancia mediana\")\n",
        "          print(linea)\n",
        "        else:\n",
        "          print(\"Instancia grande\")\n",
        "          print(linea)"
      ],
      "metadata": {
        "id": "UlAskABByZMa"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Download images by size\n",
        "from google.colab import files\n",
        "\n",
        "!zip -r imagenes_medianas.zip /content/dir_med\n",
        "!zip -r imagenes_pequeñas.zip /content/dir_peq\n",
        "!zip -r imagenes_grandes.zip /content/dir_gra\n",
        "!zip -r imagenes_mezcla.zip /content/dir_mez\n",
        "files.download('imagenes_medianas.zip')\n",
        "files.download('imagenes_pequeñas.zip')\n",
        "files.download('imagenes_grandes.zip')\n",
        "files.download('imagenes_mezcla.zip')"
      ],
      "metadata": {
        "id": "bPZvpWieG3rl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Modify labels value\n",
        "\n",
        ">If there is a problem because the dataset classes are changed, this code fragment modifies the class value in the labels.  \n",
        "Replace  **labels_path** with the path of the directory containing the labels.\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "Lws4kFaVRmsg"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import glob\n",
        "dir = glob.glob('labels_path/*.txt')\n",
        "\n",
        "for elem in dir:\n",
        " fich1 = open(elem,'r')\n",
        " for line in fich1:\n",
        "      if(line[0] == \"0\"):\n",
        "        data = line.replace(\"0\",\"1\",1)\n",
        "        fich2 = open(elem,'w')\n",
        "        fich2.write(data)\n",
        "        fich2.close()\n",
        "      else:\n",
        "        data = line.replace(\"1\",\"0\",1)\n",
        "        fich2 = open(elem,'w')\n",
        "        fich2.write(data)\n",
        "        fich2.close()\n",
        " fich1.close()\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "buSfhnQZ4yQL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Download directory\n",
        "from google.colab import files\n",
        "\n",
        "!zip -r labels.zip labels_path\n",
        "files.download('labels.zip')"
      ],
      "metadata": {
        "id": "r497aWQz8R7H"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": [
        "KJZ8-j-pHONy",
        "V0L0b39uM2--",
        "IkO_f-goz8_i",
        "8ZBV-H0IPmOA",
        "eLLZhBhyNieQ",
        "bvrh4RF680Lw",
        "BRwffXFufBeC",
        "Lws4kFaVRmsg"
      ],
      "gpuType": "T4",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}