This repository contains an implementation of object detection using YOLOv8 specifically designed for detecting weapons in images and videos. The repository includes pre-trained models and sample data for testing.
Language: Python 3.9
Dataset: https://universe.roboflow.com/fend-tech/weapon-detection-dinou
Here I have used a labelled dataset but if are unable to find a pre-labelled dataset you can label it yourself.
Split your dataset into three classes train test valid.
Use labelImg -> https://github.com/heartexlabs/labelImg Make sure to select the txt format for the labels as that is what YOLO requires.
YOLO also requires a data.yaml file. If this file is not generated by labelimg simply create one, here is what my data.yaml contains :
names:
- Grenade
- Gun
- Knife
- Pistol
- handgun
- rifle
nc: 6
test: ../test/images
train: ../train/images
val: ../valid/images
The Directory should be as follows:
object_detection
└── train
├── images
│ ├── image_1.jpg
│ ├── image_2.jpg
│ └── ...
│
└── labels
├── image_1.txt
├── image_2.txt
└── ...
└── test
├── images
│ ├── image_1.jpg
│ ├── image_2.jpg
│ └── ...
│
└── labels
├── image_1.txt
├── image_2.txt
└── ...
└── valid
├── images
│ ├── image_1.jpg
│ ├── image_2.jpg
│ └── ...
│
└── labels
├── image_1.txt
├── image_2.txt
└── ...
├── data.yaml
├──weapon-detection.ipynb
$python -m venv venv
# In cmd.exe
venv\Scripts\activate.bat
# In PowerShell
venv\Scripts\Activate.ps1
source myvenv/bin/activate
$pip install ultralytics
$yolo task=detect mode=predict model=yolov8n.pt conf=0.25 source='https://media.roboflow.com/notebooks/examples/dog.jpeg'
$yolo task=detect mode=train model=yolov8s.pt data={dataset location}/data.yaml epochs=100 imgsz=640
Your model will begin training and run for several minutes, or hours, depending on how big the dataset is and which training options you chose. A folder named runs will be created.
yolo task=detect mode=val model=/runs/detect/train/weights/best.pt data={dataset location}/data.yaml
yolo task=detect mode=predict model={HOME}/runs/detect/train/weights/best.pt conf=0.25 source={path to the image u wish to run inference on}
You can also run prediction on an entire folder containing test images
yolo task=detect mode=predict model={HOME}/runs/detect/train/weights/best.pt conf=0.25 source={dataset.location}/test/images
Note :Roboflow is a quick way to get your hands on some labelled data. https://universe.roboflow.com/roboflow-100