Skip to content

Commit c691abe

Browse files
adipandasadipandas
committed
Update IoUTracker
1 parent d46901e commit c691abe

File tree

8 files changed

+133
-44
lines changed

8 files changed

+133
-44
lines changed

‎DOWNLOAD_WEIGHTS.md‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
## Instructions to download pretrained neural-network weights.
2+
3+
##### YOLOv3
4+
```
5+
cd ./examples/pretrained_models/yolo_weights
6+
sudo chmod +x ./get_yolo.sh
7+
./get_yolo.sh
8+
```
9+
10+
##### TensorFlow - MobileNetSSDv2
11+
```
12+
cd ./pretrained_models/tensorflow_weights
13+
sudo chmod +x ./get_ssd_model.sh
14+
./get_ssd_model.sh
15+
```
16+
17+
##### Caffemodel - MobileNetSSD
18+
```
19+
cd ./pretrained_models/caffemodel_weights
20+
sudo chmod +x ./get_caffemodel.sh
21+
./get_caffemodel.sh
22+
```

‎README.md‎

Lines changed: 13 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -2,16 +2,16 @@
22
[cows-tf-ssd-output]: examples/assets/cows.gif "Sample Output with SSD"
33

44
# Multi-object trackers in Python
5-
Object detection using deep learning and multi-object tracking
5+
Various multi-object tracking algorithms.
66

77
[![DOI](https://zenodo.org/badge/148338463.svg)](https://zenodo.org/badge/latestdoi/148338463)
88

99
### Available Multi Object Trackers
1010

1111
```
1212
CentroidTracker
13-
CentroidKF_Tracker
1413
IOUTracker
14+
CentroidKF_Tracker
1515
SORT
1616
```
1717

@@ -50,43 +50,26 @@ pip install -e .
5050

5151
For using the opencv `dnn`-based object detection modules provided in this repository with GPU, you may have to compile a CUDA enabled version of OpenCV from source.
5252

53-
For building opencv from source, you can refer the following: [[link-1](https://docs.opencv.org/master/df/d65/tutorial_table_of_content_introduction.html)], [[link-2](https://www.pyimagesearch.com/2020/02/03/how-to-use-opencvs-dnn-module-with-nvidia-gpus-cuda-and-cudnn/)]
53+
For building opencv from source, you can refer the following:
54+
[[link-1](https://docs.opencv.org/master/df/d65/tutorial_table_of_content_introduction.html)],
55+
[[link-2](https://www.pyimagesearch.com/2020/02/03/how-to-use-opencvs-dnn-module-with-nvidia-gpus-cuda-and-cudnn/)]
5456

55-
### How to use?: Examples
57+
#### How to use?: Examples
5658

5759
Examples for how to use `motrackers` are provided [examples](./examples/) folder of this repository.
5860
You can clone and run the examples as shown in the [readme](examples/readme.md) inside the [examples](./examples/) folder.
5961

60-
### Pretrained Object Detection Models
61-
62-
You will have to download the pretrained weights for the model. The shell scripts for downloading are provided in [examples](examples/) folder.
63-
64-
##### YOLOv3
65-
```
66-
cd ./examples/pretrained_models/yolo_weights
67-
sudo chmod +x ./get_yolo.sh
68-
./get_yolo.sh
69-
```
62+
#### Pretrained object detection models
7063

71-
##### TensorFlow MobileNetSSDv2
72-
```
73-
cd ./pretrained_models/tensorflow_weights
74-
sudo chmod +x ./get_ssd_model.sh
75-
./get_ssd_model.sh
76-
```
77-
78-
##### Caffemodel
79-
```
80-
cd ./pretrained_models/caffemodel_weights
81-
sudo chmod +x ./get_caffemodel.sh
82-
./get_caffemodel.sh
83-
```
64+
You will have to download the pretrained weights for the neural-network models.
65+
The shell scripts for downloading these are provided in [examples](examples/) folder.
66+
Please refer [DOWNLOAD_WEIGHTS.md](DOWNLOAD_WEIGHTS.md) for more details.
8467

85-
### References and Credits
68+
#### References and Credits
8669

87-
Please see [references.md](REFERENCES.md)
70+
Please see [REFERENCES.md](REFERENCES.md)
8871

89-
### Citation
72+
#### Citation
9073

9174
If you use this repository in your work, please consider citing it with:
9275
```

‎examples/example_scripts/mot_Caffe_SSDMobileNet.py‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
import cv2 as cv
2+
from motrackers.detectors import Caffe_SSDMobileNet
3+
from motrackers import CentroidTracker, CentroidKF_Tracker, SORT, IOUTracker
4+
from motrackers.utils import draw_tracks
5+
6+
7+
def main(video_path, weights_path, config_path, use_gpu, tracker):
8+
9+
model = Caffe_SSDMobileNet(
10+
weights_path=weights_path,
11+
configfile_path=config_path,
12+
confidence_threshold=0.5,
13+
nms_threshold=0.2,
14+
draw_bboxes=True,
15+
use_gpu=use_gpu)
16+
17+
cap = cv.VideoCapture(video_path)
18+
while True:
19+
ok, image = cap.read()
20+
21+
if not ok:
22+
print("Cannot read the video feed.")
23+
break
24+
25+
image = cv.resize(image, (700, 500))
26+
27+
bboxes, confidences, class_ids = model.detect(image)
28+
tracks = tracker.update(bboxes, confidences, class_ids)
29+
updated_image = model.draw_bboxes(image.copy(), bboxes, confidences, class_ids)
30+
31+
updated_image = draw_tracks(updated_image, tracks)
32+
33+
cv.imshow("image", updated_image)
34+
if cv.waitKey(1) & 0xFF == ord('q'):
35+
break
36+
37+
cap.release()
38+
cv.destroyAllWindows()
39+
40+
41+
if __name__ == '__main__':
42+
import argparse
43+
44+
parser = argparse.ArgumentParser(
45+
description='Object detections in input video using Caffemodel of MobileNetSSD.')
46+
47+
parser.add_argument(
48+
'--video', '-v', type=str, default="./../video_data/cars.mp4", help='Input video path.')
49+
50+
parser.add_argument(
51+
'--weights', '-w', type=str, default="./../pretrained_models/caffemodel_weights/MobileNetSSD_deploy.caffemodel",
52+
help='path to weights file of Caffe-MobileNetSSD, i.e., `.caffemodel` file.')
53+
54+
parser.add_argument(
55+
'--config', '-c', type=str, default="./../pretrained_models/caffemodel_weights/MobileNetSSD_deploy.prototxt",
56+
help='path to config file of Caffe-MobileNetSSD, i.e., `.prototxt` file.')
57+
58+
parser.add_argument(
59+
'--gpu', type=bool, default=False, help='Flag to use gpu to run the deep learning model. Default is `False`')
60+
61+
parser.add_argument(
62+
'--tracker', type=str, default='IOUTracker', help="Tracker used to track objects."
63+
" Options include ['CentroidTracker', 'CentroidKF_Tracker', 'SORT', IOUTracker]")
64+
65+
args = parser.parse_args()
66+
67+
if args.tracker == 'CentroidTracker':
68+
tracker = CentroidTracker(max_lost=0, tracker_output_format='mot_challenge')
69+
elif args.tracker == 'CentroidKF_Tracker':
70+
tracker = CentroidKF_Tracker(max_lost=0, tracker_output_format='mot_challenge')
71+
elif args.tracker == 'SORT':
72+
tracker = SORT(max_lost=3, tracker_output_format='mot_challenge', iou_threshold=0.3)
73+
elif args.tracker == 'IOUTracker':
74+
tracker = IOUTracker(max_lost=2, iou_threshold=0.5, min_detection_confidence=0.4, max_detection_confidence=0.7,
75+
tracker_output_format='mot_challenge')
76+
else:
77+
raise NotImplementedError
78+
79+
main(args.video, args.weights, args.config, args.gpu, tracker)

‎examples/example_scripts/mot_TF_SSDMobileNet.py‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import cv2 as cv
22
from motrackers.detectors import TF_SSDMobileNetV2
3-
from motrackers import CentroidTracker, CentroidKF_Tracker, SORT
3+
from motrackers import CentroidTracker, CentroidKF_Tracker, SORT, IOUTracker
44
from motrackers.utils import draw_tracks
55

66

@@ -67,8 +67,9 @@ def main(video_path, weights_path, config_path, use_gpu, tracker):
6767
)
6868

6969
parser.add_argument(
70-
'--tracker', type=str, default='SORT',
71-
help="Tracker used to track objects. Options include ['CentroidTracker', 'CentroidKF_Tracker', 'SORT']")
70+
'--tracker', type=str, default='IOUTracker',
71+
help="Tracker used to track objects. Options include ['CentroidTracker', 'CentroidKF_Tracker', "
72+
"'SORT', IOUTracker]")
7273

7374
args = parser.parse_args()
7475

@@ -78,6 +79,9 @@ def main(video_path, weights_path, config_path, use_gpu, tracker):
7879
tracker = CentroidKF_Tracker(max_lost=0, tracker_output_format='mot_challenge')
7980
elif args.tracker == 'SORT':
8081
tracker = SORT(max_lost=3, tracker_output_format='mot_challenge', iou_threshold=0.3)
82+
elif args.tracker == 'IOUTracker':
83+
tracker = IOUTracker(max_lost=2, iou_threshold=0.5, min_detection_confidence=0.4, max_detection_confidence=0.7,
84+
tracker_output_format='mot_challenge')
8185
else:
8286
raise NotImplementedError
8387

‎examples/example_scripts/mot_YOLOv3.py‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import cv2 as cv
22
from motrackers.detectors import YOLOv3
3-
from motrackers import CentroidTracker, CentroidKF_Tracker, SORT
3+
from motrackers import CentroidTracker, CentroidKF_Tracker, SORT, IOUTracker
44
from motrackers.utils import draw_tracks
55

66

@@ -73,17 +73,20 @@ def main(video_path, weights_path, config_path, labels_path, use_gpu, tracker):
7373
)
7474

7575
parser.add_argument(
76-
'--tracker', type=str, default='CentroidTracker',
76+
'--tracker', type=str, default='IOUTracker',
7777
help="Tracker used to track objects. Options include ['CentroidTracker', 'CentroidKF_Tracker', 'SORT']")
7878

7979
args = parser.parse_args()
8080

8181
if args.tracker == 'CentroidTracker':
82-
tracker = CentroidTracker(max_lost=3, tracker_output_format='mot_challenge')
82+
tracker = CentroidTracker(max_lost=0, tracker_output_format='mot_challenge')
8383
elif args.tracker == 'CentroidKF_Tracker':
84-
tracker = CentroidKF_Tracker(max_lost=3, tracker_output_format='mot_challenge')
84+
tracker = CentroidKF_Tracker(max_lost=0, tracker_output_format='mot_challenge')
8585
elif args.tracker == 'SORT':
86-
tracker = SORT(max_lost=3, tracker_output_format='mot_challenge', iou_threshold=0.3, time_step=1)
86+
tracker = SORT(max_lost=3, tracker_output_format='mot_challenge', iou_threshold=0.3)
87+
elif args.tracker == 'IOUTracker':
88+
tracker = IOUTracker(max_lost=2, iou_threshold=0.5, min_detection_confidence=0.4, max_detection_confidence=0.7,
89+
tracker_output_format='mot_challenge')
8790
else:
8891
raise NotImplementedError
8992

‎motrackers/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,4 +10,4 @@
1010
from motrackers.tracker import Tracker as CentroidTracker
1111
from motrackers.centroid_kf_tracker import CentroidKF_Tracker
1212
from motrackers.sort_tracker import SORT
13-
# from motrackers.iou_tracker import IOUTracker
13+
from motrackers.iou_tracker import IOUTracker

‎motrackers/detectors/tf.py‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,7 @@
33

44

55
class TF_SSDMobileNetV2(Detector):
6-
def __init__(
7-
self, weights_path, configfile_path,
8-
confidence_threshold=0.5, nms_threshold=0.4, draw_bboxes=True, use_gpu=False):
6+
def __init__(self, weights_path, configfile_path, confidence_threshold=0.5, nms_threshold=0.4, draw_bboxes=True, use_gpu=False):
97
self.image_size = (300, 300)
108
self.net = cv.dnn.readNetFromTensorflow(weights_path, configfile_path)
119

‎motrackers/iou_tracker.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from motrackers.utils import iou
1+
from motrackers.utils.misc import iou_xywh as iou
22
from motrackers.tracker import Tracker
33

44

@@ -26,7 +26,7 @@ class IOUTracker(Tracker):
2626

2727
def __init__(
2828
self,
29-
max_lost=0,
29+
max_lost=2,
3030
iou_threshold=0.5,
3131
min_detection_confidence=0.4,
3232
max_detection_confidence=0.7,

0 commit comments

Comments
 (0)