Notebook_PICA_40_B_883.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "<a target=\"_blank\" href=\"https://colab.research.google.com/github/marianbasti/Proyecto-PICA-40-B-883/blob/main/Notebook_PICA_40_B_883.ipynb\">\n",
        "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
        "</a>\n",
        "\n",
        "# Proyecto PICA 40-B-883\n",
        "En este notebook veremos en funcionamiento el entrenamiento del modelo de detección así como la inferencia y extracción de información"
      ],
      "metadata": {
        "id": "IhmJzBkASxVG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# @markdown Antes de comenzar, clonamos el repositorio e instalamos las librerías necesarias\n",
        "!git clone https://github.com/marianbasti/Proyecto-PICA-40-B-883\n",
        "%cd Proyecto-PICA-40-B-883\n",
        "!pip install -r requirements.txt"
      ],
      "metadata": {
        "cellView": "form",
        "id": "Z95BJ0GfS89-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @markdown Y luego importamos las librerías que usaremos\n",
        "import os, json, argparse\n",
        "import cv2\n",
        "import numpy as np\n",
        "from bisect import bisect_left\n",
        "from collections import defaultdict\n",
        "from datetime import datetime, timedelta\n",
        "from statistics import mean\n",
        "from ultralytics import YOLO"
      ],
      "metadata": {
        "cellView": "form",
        "id": "e-YT80FCTExe"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Entrenamiento\n",
        "Adaptaremos el modelo base YOLOv8 para nuestro caso, la detección de avispas y su casta"
      ],
      "metadata": {
        "id": "_TgUCFxcRVq_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# @markdown Descargamos el dataset\n",
        "!gdown 1skVPS8g-JSSWca0zt_500vZNW5f3vPn8\n",
        "!unzip dataset"
      ],
      "metadata": {
        "cellView": "form",
        "id": "XHuo9hd1SU-3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @markdown Y luego realizamos el entrenamiento\n",
        "\n",
        "# Definimos el directorio del dataset con formato YOLOv8\n",
        "dataset = 'dataset'\n",
        "\n",
        "epochs=100 # @param\n",
        "learning_rate=0.015 # @param\n",
        "\n",
        "# Cargamos modelo base\n",
        "model = YOLO('yolov8n.pt')\n",
        "\n",
        "# Entrenamos el modelo con el dataset\n",
        "results = model.train(data=f'/content/Proyecto-PICA-40-B-883/{dataset}/data.yaml', epochs=epochs, imgsz=640, batch=-1, lr0=learning_rate)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "lIDaU0XaTesd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Inferencia"
      ],
      "metadata": {
        "id": "1NMCK_bEUn3t"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# @markdown Descargamos videos+data asociada para probar detección\n",
        "%mkdir videos\n",
        "%cd videos\n",
        "import gdown\n",
        "gdown.download_folder(\"https://drive.google.com/drive/folders/1-ROZBlggbJYK4N6yYCl_nBbJkIsXUWO4?usp=sharing\", quiet=True)\n",
        "%cd .."
      ],
      "metadata": {
        "cellView": "form",
        "id": "qwmFxK7HUv3M"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "upwsTwj2PLUT"
      },
      "outputs": [],
      "source": [
        "# @markdown ## Variables\n",
        "# @markdown Directorio donde está nuestro modelo entrenado\n",
        "model_path = 'yolov8_avispas.pt' # @param\n",
        "# @markdown Directorio de los videos (video/ + sensordata)\n",
        "input_path = './videos' # @param\n",
        "# @markdown Exportar el video para visualización\n",
        "save_vid = False # @param {'type':'boolean'}\n",
        "# @markdown Nombre de salida para CSV\n",
        "txt_path = 'tracks' # @param\n",
        "# @markdown Parámetros de detección\n",
        "iou=0.5 # @param\n",
        "conf=0.5 # @param\n",
        "\n",
        "# Valor del threshold\n",
        "threshold_value = 90\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# @markdown Funciones que usaremos para la detección y procesamientos\n",
        "def load_sensor_data(sensor_file):\n",
        "    sensor_data = {}\n",
        "    with open(sensor_file, 'r') as f:\n",
        "        for line in f:\n",
        "            parts = line.strip().split()\n",
        "            time_str = parts[1]\n",
        "            temp = float(parts[2].split('=')[1][:-2])\n",
        "            humidity = float(parts[3].split('=')[1][:-1])\n",
        "            sensor_data[time_str] = {'temp': temp, 'humidity': humidity}\n",
        "    return sensor_data\n",
        "\n",
        "def interpolate(a, b, alpha):\n",
        "    return a + alpha * (b - a)\n",
        "\n",
        "def parse_time(t):\n",
        "    try:\n",
        "        return datetime.strptime(t, '%H:%M:%S.%f')\n",
        "    except ValueError:\n",
        "        return datetime.strptime(t, '%H:%M:%S')\n",
        "\n",
        "def find_closest_data(time_str, sensor_data):\n",
        "    times = sorted(sensor_data.keys())\n",
        "    pos = bisect_left(times, time_str)\n",
        "\n",
        "    if pos == 0 or pos == len(times):\n",
        "        return sensor_data[times[0]] if pos == 0 else sensor_data[times[-1]]\n",
        "\n",
        "    before_time = parse_time(times[pos - 1])\n",
        "    after_time = parse_time(times[pos])\n",
        "    current_time = parse_time(time_str)\n",
        "\n",
        "    alpha = (current_time - before_time).total_seconds() / (after_time - before_time).total_seconds()\n",
        "\n",
        "    before = sensor_data[times[pos - 1]]\n",
        "    after = sensor_data[times[pos]]\n",
        "\n",
        "    if not (after['temp'] and after['humidity']):\n",
        "        return before\n",
        "\n",
        "    return {\n",
        "        'temp': interpolate(before['temp'], after['temp'], alpha),\n",
        "        'humidity': interpolate(before['humidity'], after['humidity'], alpha)\n",
        "    }\n",
        "\n",
        "def refine_single_bbox(im, bbox_xyxy):\n",
        "    bbox_np = np.array(bbox_xyxy).flatten()\n",
        "    x_center, y_center, w, h = int(bbox_np[0]), int(bbox_np[1]), int(bbox_np[2]), int(bbox_np[3])\n",
        "\n",
        "    x1, y1 = x_center - (w // 2), y_center - (h // 2)\n",
        "    x2, y2 = x1 + w, y1 + h\n",
        "    roi = im[y1:y1+h, x1:x1+w]\n",
        "    # cv2.imwrite('./avispas/roi.jpg', roi)\n",
        "    # Thresholding\n",
        "    try:\n",
        "        roi_gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)\n",
        "    except:\n",
        "        print(f\"Error with cropping image :\\n{bbox_np}\")\n",
        "\n",
        "    _, thresh = cv2.threshold(roi_gray, threshold_value, 255, cv2.THRESH_BINARY_INV)\n",
        "\n",
        "    # Find contours\n",
        "    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n",
        "    largest_contour = max(contours, key=cv2.contourArea)\n",
        "    translated_contour = largest_contour + [x1, y1]\n",
        "    cv2.polylines(annotated_frame, translated_contour, isClosed=True, color=(0, 255, 0), thickness=2)\n",
        "\n",
        "    if len(contours) == 0:\n",
        "        # box - 4 corners of the rectangle\n",
        "        box = np.array([[x1, y1], [x2, y1], [x2, y2], [x1, y2]])\n",
        "        # min_rect - center, size (width, height), angle\n",
        "        center = ((x1 + x2) / 2, (y1 + y2) / 2)\n",
        "        size = (x2 - x1, y2 - y1)\n",
        "        angle = 0\n",
        "        min_rect = (center, size, angle)\n",
        "        return box, min_rect\n",
        "\n",
        "    # Find the smallest enclosing rotated rectangle for closed contours\n",
        "    min_rect = cv2.minAreaRect(np.vstack(largest_contour))\n",
        "\n",
        "    # Find the smallest enclosing rotated rectangle\n",
        "    min_rect = cv2.minAreaRect(np.vstack(largest_contour))\n",
        "    box = cv2.boxPoints(min_rect).astype(int)\n",
        "\n",
        "    # Translate to original image coordinates\n",
        "    box[:, 0] += x1\n",
        "    box[:, 1] += y1\n",
        "\n",
        "    w, h = min_rect[1]\n",
        "\n",
        "    if h > w:\n",
        "        w, h = h, w\n",
        "        angle = min_rect[2]\n",
        "        if angle < 0:\n",
        "            min_rect = (min_rect[0], (w, h), angle + 90)\n",
        "        else:\n",
        "            min_rect = (min_rect[0], (w, h), angle - 90)\n",
        "\n",
        "    cv2.polylines(im, [box], isClosed=True, color=(0, 255, 0), thickness=2)\n",
        "\n",
        "    return box, min_rect"
      ],
      "metadata": {
        "cellView": "form",
        "id": "RT78idDKWKfs"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @markdown Detección\n",
        "# Cargamos el modelo YOLOv8\n",
        "model = YOLO(model_path)\n",
        "\n",
        "# Inicializamos el archivo de salida\n",
        "if not os.path.exists(txt_path + '.csv'):\n",
        "        with open(txt_path + '.csv', 'a') as f:\n",
        "            f.write('id,timestamp,temp,humidity,largo,ancho,movement,time,filename\\n')\n",
        "\n",
        "# Por cada capeta que representa un día de filmaciones\n",
        "for source in os.listdir(input_path):\n",
        "\n",
        "    # Extraemos la fecha\n",
        "    date_str = source.split('\\\\')[-1]\n",
        "\n",
        "    # Cargamos el archivo con datos de los sensores\n",
        "    sensor_file = os.path.join(input_path,source, f\"temperatura {date_str}.txt\")\n",
        "    sensor_data = load_sensor_data(sensor_file)\n",
        "\n",
        "    # Por cada video en la carpeta \"/video\"\n",
        "    for video_path in os.listdir(input_path+'/'+source+'/video'):\n",
        "\n",
        "        # Inicializamos el historial de trackeos\n",
        "        track_history = defaultdict(lambda: {'points': [], 'timestamp': [], 'width': [], 'height': []})\n",
        "\n",
        "        # Abrimos el video\n",
        "        cap = cv2.VideoCapture(os.path.join(input_path,source)+'/video/'+video_path)\n",
        "\n",
        "        # Obtenemos datos del video original\n",
        "        frame_rate = int(cap.get(cv2.CAP_PROP_FPS))\n",
        "        fps = cap.get(cv2.CAP_PROP_FPS)\n",
        "        img_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
        "        img_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
        "\n",
        "        if save_vid:  # Si queremos guardar el video para visualizar el trackeo\n",
        "            vid_writer = cv2.VideoWriter(f\"{video_path.split('.mp4')[0]}.avi\", cv2.VideoWriter_fourcc(*'mp4v'), fps, (img_w, img_h))\n",
        "\n",
        "        # Limite para considerar salida/entrada\n",
        "        significant_move = img_w/4\n",
        "\n",
        "        # Limite para ignorar altos y anchos (en los bordes confunde)\n",
        "        border_threshold = img_w/5\n",
        "\n",
        "        # Extraemos la hora\n",
        "        time_str = video_path.split(' ')[-1].replace('_', ':').split('.')[0]\n",
        "        initial_time = datetime.strptime(f\"{date_str} {time_str}\", '%Y-%m-%d %H:%M:%S')\n",
        "\n",
        "        # Loop a través de cada cuadro del video\n",
        "        while cap.isOpened():\n",
        "\n",
        "            # Leemos el cuadro\n",
        "            success, frame = cap.read()\n",
        "            if success:\n",
        "                # Corremos YOLOv8\n",
        "                results = model.track(frame, persist=True, conf=conf, iou=iou, tracker='avispa_bytetrack.yaml')\n",
        "\n",
        "                # Obtenemos cajas e IDs detectadas\n",
        "                boxes = results[0].boxes.xywh.cpu()\n",
        "                annotated_frame = results[0].plot()\n",
        "\n",
        "                # Si hay detección\n",
        "                if results[0].boxes:\n",
        "                    if results[0].boxes.id != None:\n",
        "                        # Obtenemos datos temporales de la detección\n",
        "                        current_frame = int(cap.get(cv2.CAP_PROP_POS_FRAMES))\n",
        "                        elapsed_time = timedelta(seconds=current_frame / frame_rate)\n",
        "                        timestamp = (initial_time + elapsed_time).strftime('%H:%M:%S.%f')\n",
        "                        track_ids = results[0].boxes.id.int().cpu().tolist()\n",
        "\n",
        "                        # Por cada detección\n",
        "                        for box, track_id in zip(boxes, track_ids):\n",
        "\n",
        "                            # Guardamos los puntos, timestamp y tamaños de la detección\n",
        "                            x, y, w, h = box\n",
        "                            track = track_history[track_id]\n",
        "                            track['points'].append((float(x), float(y)))\n",
        "                            track['timestamp'].append(timestamp)\n",
        "\n",
        "                            # Refinamos el tamaño para que sea el mas compacto posible\n",
        "                            refined_xyxy,min_rect = refine_single_bbox(frame,box)\n",
        "                            r_w, r_h = min_rect[1]\n",
        "                            track['width'].append(r_w)\n",
        "                            track['height'].append(r_h)\n",
        "\n",
        "                            # Por último dibujamos las líneas del trackeo\n",
        "                            points = np.hstack(track['points']).astype(np.int32).reshape((-1, 1, 2))\n",
        "                            if x>border_threshold and x<img_w-border_threshold:\n",
        "                                cv2.polylines(annotated_frame, [points], isClosed=False, color=(20, 255, 0), thickness=4)\n",
        "                            else:\n",
        "                                cv2.polylines(annotated_frame, [points], isClosed=False, color=(0, 20, 250), thickness=4)\n",
        "\n",
        "                if save_vid:  # Si guardamos el video procesado\n",
        "                    vid_writer.write(annotated_frame)\n",
        "\n",
        "            else:\n",
        "                # Salir del loop si terminó el video\n",
        "                break\n",
        "\n",
        "        # Por cada track detectado\n",
        "        for id in track_history:\n",
        "            # Definimos si entró o salió (entrada a la izquierda, salida a la derecha)\n",
        "            initial_x = track_history[id]['points'][0][0]\n",
        "            final_x = track_history[id]['points'][-1][0]\n",
        "            movement = abs(final_x - initial_x)\n",
        "\n",
        "            # Si el movimiento es significativo\n",
        "            if movement > significant_move:\n",
        "                # Definimos si entró o salió\n",
        "                if final_x > initial_x:\n",
        "                    print(f\"Object {id} moved out.\")\n",
        "                    str_movement='out'\n",
        "                elif final_x < initial_x:\n",
        "                    print(f\"Object {id} moved in.\")\n",
        "                    str_movement='in'\n",
        "            else:\n",
        "                str_movement='undetermined'\n",
        "\n",
        "            # Formateamos data\n",
        "            timestamp = track_history[id]['timestamp'][0]\n",
        "            data = find_closest_data(timestamp,sensor_data)\n",
        "            temp = data.get('temp', 'N/A')\n",
        "            humidity = data.get('humidity', 'N/A')\n",
        "\n",
        "            filtered_widths = [w for w, (x, _) in zip(track_history[id]['width'], track_history[id]['points']) if x > border_threshold and x < (img_w - border_threshold)]\n",
        "            filtered_heights = [h for h, (x, _) in zip(track_history[id]['height'], track_history[id]['points']) if x > border_threshold and x < (img_w - border_threshold)]\n",
        "\n",
        "            if len(filtered_widths)>0:\n",
        "                avg_w = round(mean(filtered_widths),2)\n",
        "                avg_h = round(mean(filtered_heights),2)\n",
        "            else:\n",
        "                avg_w = \"N/A\"\n",
        "                avg_h = \"N/A\"\n",
        "\n",
        "            # Calculamos cuanto tiempo tardó en entrar/salir\n",
        "            time_taken = datetime.strptime(track_history[id]['timestamp'][-1],'%H:%M:%S.%f') - datetime.strptime(track_history[id]['timestamp'][0], '%H:%M:%S.%f')\n",
        "\n",
        "            # Armamos el log\n",
        "            log_data=f\"\"\"{id},{timestamp},{round(temp,2)},{round(humidity,2)},{avg_w},{avg_h},{str_movement},{str(time_taken)[2:-4]},{video_path}\\n\"\"\"\n",
        "\n",
        "            # Agregamos la data del video en el csv\n",
        "            with open(f\"{txt_path}.csv\", \"a\") as f:\n",
        "                f.write(log_data)\n",
        "\n",
        "# Finalizar\n",
        "if save_vid:\n",
        "    vid_writer.release()\n",
        "\n",
        "cap.release()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "Zyp9ugzJWX5s"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}