Added model stats for yolo

jinshengye-git · Mar 12, 2021 · 717e6c7 · 717e6c7
1 parent 6df2918
commit 717e6c7
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 4 deletions.
diff --git a/README.md b/README.md
@@ -11,11 +11,15 @@ We currently support the models:
 The SSD model is taken from [here](https://github.com/qfgaohao/pytorch-ssd).
 The YOLO model is taken from [here](https://github.com/ultralytics/yolov5).
 
+
+![Example](media/detr-example.jpg)
+
+
 ## Dependencies and setup
-* Python >= 3.79
+* Python ≥ 3.7.10
 * OpenCV2
-* PyTorch >= 1.6.0
-* Torchvision >= 0.7.0
+* PyTorch ≥ 1.7.1
+* Torchvision ≥ 0.8.2
 
 To use the video function place `.mp4` videos into the folder `media/DrivingClips`.
 

diff --git a/media/detr-example.jpg b/media/detr-example.jpg
diff --git a/visualizer/stats_model.py b/visualizer/stats_model.py
@@ -38,7 +38,7 @@ def showStats(image, model_type, labels, conf):
     model_stats = dict.fromkeys(keys)
 
     # COCO dataset
-    if (model_type == "-detr" or model_type == "-fasterrcnn"):
+    if (model_type == "-detr" or model_type == "-fasterrcnn" or model_type == "-yolov5s"):
         model_stats["numCars"] = np.count_nonzero(labels == 3)
         model_stats["numTrucksBuses"] = np.count_nonzero(labels == 6) + np.count_nonzero(labels == 8)
         model_stats["numMotorCycles"] = np.count_nonzero(labels == 4)

diff --git a/yolov5s.pt b/yolov5s.pt