Entrenaremos una red CNN Yolov5 para detectar armas
Usando conda en el entorno base, clonamos el repositorio de labelImg
$ git clone https://github.com/tzutalin/labelImg
$ conda install pyqt=5
$ conda install -c anaconda lxml
$ cd labelImg
$ pyrcc5 -o libs/resources.py resources.qrc
$ python labelImg.py
Etiquetar el dataset y guardar usando la siguiente estructura:
--train_data
-- images
--train
--val
-- labels
--train
--val
Seguiremos la guía del sitio oficial YoloV5
Sobre el codgio del enlace anterior haremos las siguientes modificaciones
!unzip -q ../train_data.zip -d ../
Creamos un archivo customdata.yaml con la siguiente informacion, y guardar en la ruta yolov5/data
path: ../train_data # dataset root dir
train: ../train_data/images/train/ # train images (relative to 'path') 128 images
val: ../train_data/images/val/ # val images (relative to 'path') 128 images
test: # test images (optional)
# Classes
nc: 1 # number of classes
names: ['gun'] # class names
Finalmente ejecutamos la linea de entrenamiento
$ !python train.py --img 640 --batch 4 --epochs 100 --data customdata.yaml --weights yolov5s.pt --cache
$ !python detect.py --weights /content/yolov5/runs/train/exp/weights/best.pt --img 640 --conf 0.25 --source "../guntest.jpg"
$ display.Image(filename='../guntest.jpg', width=600)
Vamos a crear un entorno usando conda, y en ese entorno vamos a instalar los requerimientos adecuados. Luego probaremos la inferencia usando la cámara web.
$ conda create -n YoloV5Test
$ conda activate YoloV5Test
$ conda install python=3.7
$ pip install jupyter
$ conda install ipykernel
$ python -m ipykernel install --user --name YoloV5Test --display-name "YoloV5Test"
$ git clone https://github.com/ultralytics/yolov5
$ cd yolov5
$ pip install -qr requirements.txt
$ jupyter notebook
Luego vamos a abrir el archivo yolov5video.ipynb y ejecutamos la inferencia en video.
Comprobar compatibilidad con GPU
Verificar compatibilidad de librerias y toolkit de cuda
$ pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 torchaudio==0.11.0 --extra-index-url https://download.pytorch.org/whl/cu113
$ import torch
$ import torchvision
$ print(torch.cuda.is_available())
$ print(torch.__version__)
$ print(torchvision.__version__)
$ python detect.py --weights best.pt --source 0