SnapSort is a trash sorting assistant with the YOLO V4 model. It will take in a video feed from the camera and return visual and audio instructions to users. The project is designed to help people sort their waste properly before waste collection. We placed the device, Jetson Nano or GTX 1080 with a screen, on top of trash cans to guide people on sorting waste.
The project is built for the University of Washington GIX TECHIN 514 Hardware Software Lab 1 course, Microsoft Imagine Cup, and Alaska Environment Innovation Challenges. The code and dataset are collected and built by Yun Liu, Joey Wang, and me.
The dataset reaches 71% MAP for 12 different categories on an 80/20 train/test split trained on the YOLO v4 object detection model.
I also convert the Snapsort Yolo model to Onnx and TensorRT format for better performance on Jetson Nano.
- Introduction
- What's new
- Dataset
- Running Snapsort on Windows/Linux
- Runing Snapsort on Jetson Nano
- Others
- Add trained weights with Yolo V4, V4 Tiny model
- Optimized performance for Jetson Nano with Onnx and Tensor RT core
The dataset contains 4600 original images including images from Coco Dataset, Google open images v4, and images we collected by ourselves. Images are labeled into 12 classes manually following Seattle Government's recycling rule. After data augmentation, 80% of images are used for training and 20% of images are used for testing. The best training result we got is 71% MAP @ 50% confidence threshold.
To run Snapsort on a Windows device, please install all the dependencies listed below and Darknet Yolo V4 by AlexyAB(see Prerequisites below).
- Python 3 and above
- OpenCV 3.4.0 for Python
- Numpy
- Winsound
The Yolo V4 model needs to be installed before using our trained weight. The most Windows-friendly version I found is from AlexAB.
Download and unzip all the files into the Yolo build file.
My file location C:\darknet-master\build\darknet\x64.
cd C:\darknet-master\build\darknet\x64
git clone https://github.com/Kuchunan/SnapSort-Trash-Classification-with-YOLO-v4-Darknet-Onnx-TensorRT-for-Jetson-Nano.git
Change the file location according to the weights you want to use at line 159~162 in darknet_video-Ku.py.
Option 1. 12 classes detection
configPath = "./backup/0311_files/yolov3_5l-GIX_0304_test.cfg"
weightPath = "./backup/0311_files/yolov3_5l-GIX_0326_best.weights"
metaPath = "./backup/0311_files/Aug.data"
Option 2. 3 classes detection(Recycle, Landfill, Compost)
configPath = "./backup/0311_files/yolov3_5l-GIX_0308_RLC_test.cfg"
weightPath = "./backup/0311_files/yolov3_5l-GIX_0308_RLC_best.weights"
metaPath = "./backup/0311_files/Aug_RLC.data"
Option 3. 12 classes with Yolo V2
configPath = "./backup/yolov2/yolov2_GIX_test.cfg"
weightPath = "./backup/yolov2/yolov2_GIX_best.weights"
metaPath = "./backup/0311_files/Aug.data"
Option 4. 12 classes with Yolo V4
configPath = "./backup/Yolo_V4&V4_Tiny/yolov4_GIX_best-416.cfg"
weightPath = "./backup/Yolo_V4&V4_Tiny/yolov4_GIX_best-416.weights"
metaPath = "./backup/0311_files/Aug.data"
Option 5. 12 classes with Yolo V4 Tiny (416*416)
configPath = "./backup/Yolo_V4&V4_Tiny/yolov4-tiny_GIX-416.cfg"
weightPath = "./backup/Yolo_V4&V4_Tiny/yolov4-tiny_GIX-416.weights"
metaPath = "./backup/0311_files/Aug.data"
Option 6. 12 classes with Yolo V4 Tiny (640*640)
configPath = "./backup/Yolo_V4&V4_Tiny/yolov4-tiny_GIX-640.cfg"
weightPath = "./backup/Yolo_V4&V4_Tiny/yolov4-tiny_GIX-640.weights"
metaPath = "./backup/0311_files/Aug.data"
Audio instruction will be given every 45 frames(default). It will collect all the detection results in these frames, merge the result for IoU higher than 50%, and give the audio feedback of the highest possibility. You can change the interval on line 14.
AvergeOverTime = 45
You can change the IOU on line 71.
def nms(detections,threshold=.5):
To change video quality:
VWidth = 1280
VHight = 720
To change input camera (0 means built-in camera)
cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)
To change thredhold:
detections = darknet.detect_image(netMain, metaMain, darknet_image, thresh=0.2,hier_thresh=0.3, nms=0.4)
Run the code and the detection will start. Hit ESC to exit.
- Jetson Nano * 1
- 32G micro SD card
- Webcam * 1 (Recommend: Logitech C920 HD, higher resolution and low noise will improve performance)
- 5V DC barrel jack power supply
- Screen (with power supply and HDMI cable)
- (Optional) Mouse & Keyboard
- Format a 32GB SD card(32GB or above) and install Jetpack 4.5.1 Tutorial on Nvidia.com
- Make sure connected to the internet (using an ethernet cable or wifi dongle)
- Open Command Line terminal
sudo apt-get install update
sudo apt-get install upgrade
sudo apt install python3-pip
pip3 install protobuf==3.8.0
git clone https://github.com/jkjung-avt/tensorrt_demos.git
cd tensorrt_demos/ssd
export PATH=${PATH}:/usr/local/cuda/bin
export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/local/cuda/lib64
./install_pycuda.sh
sudo apt-get install python-pip protobuf-compiler libprotoc-dev
pip install Cython
sudo pip3 install onnx==1.4.1
cd ..
cd plugins
Make
- Download all the files in the "Onnx & TensorRT' folder (Click here) to folder “~/tensorrt_demos/yolo”
-
Download yolo_classes.py and visualization.py to ~/tensorrt_demos/utils
Changes in yolo_classes.py from jkjung-avt's file:
- Line 6-21: Replace COCO_CLASSES_LIST with GIX_CLASSES_LIST & GIX_3CLASSES_LIST
- Line 36-39: Add parameter support: category_num=12 & category_num=3
Changes in visualization.py from jkjung-avt's file:
- Line 79: Add color_rule list to map with the color of trash bins
- Line 99-102: Add rules to apply to 3 classes and 12 classes
**Mapping rule**
| 12 classes | 3 classes |
|---|---|
| Plastic_soft | Landfill |
| Food_soiled_coated_paper | Landfill |
| Plastic_utensils | Landfill |
| Foam | Landfill |
| Miscellaneous | Landfill |
| Carboard | Recycle |
| Paper | Recycle |
| Plastic_rigid | Recycle |
| Glass | Recycle |
| Metal | Recycle |
| Food_scrapes | Compost |
| Paper_towels | Compost |
- Open Command Line terminal
- Plug-in a USB webcam
Option 1. 12 classes detection with Yolo V4 (416x416)
cd tensorrt_demos
python3 trt_yolo.py --usb 0 -m yolov4_GIX_best-416 --category_num=12
Option 2. 12 classes detection with Yolo V4 Tiny (416x416)
cd tensorrt_demos
python3 trt_yolo.py --usb 0 -m yolov4-tiny_GIX-416 --category_num=12
Option 3. 12 classes detection with Yolo V4 Tiny (640x640)
cd tensorrt_demos
python3 trt_yolo.py --usb 0 -m yolov4-tiny_GIX-640 --category_num=12
Option 4. 3 classes detection with Yolo V4 (416x416)
cd tensorrt_demos
python3 trt_yolo.py --usb 0 -m yolov4_GIX_best-416 --category_num=3
Option 5. 3 classes detection with Yolo V4 Tiny (416x416)
cd tensorrt_demos
python3 trt_yolo.py --usb 0 -m yolov4-tiny_GIX-416 --category_num=3
Option 6. 3 classes detection with Yolo V4 Tiny (640x640)
cd tensorrt_demos
python3 trt_yolo.py --usb 0 -m yolov4-tiny_GIX-640 --category_num=3
- "F" for fullscreen
| Model | Resolution | Fps on Jetson nano |
|---|---|---|
| Yolo v4 Tiny | 416x416 | ~19 |
| Yolo v4 | 416x416 | ~4 |
| Yolo V4 | 416x416 | ~3.x |
- Double click /tensorrt_demos/trt_yolo.py to edit
- Add this line to Line #69 (Leave 6 blank spaces for indention)
set_display(WINDOW_NAME, True)
- Save the file