GeekAlexis / FastMOT

High-performance multiple object tracking based on YOLO, Deep SORT, and KLT 🚀
MIT License
1.12k stars 256 forks source link

[Question] How to print total object counts on the video stream? #231

Closed AngelaYZhang closed 2 years ago

AngelaYZhang commented 2 years ago

Heya,

I was wondering how should I edit the code to print more information on the video stream? Currently default is print the number of objects visible for the current frame on the top left like this: replicate

I would like to get the total number of objects identified (ie. max value of unique id) across the entire video stream shown on the top left too, and possibly in future separate into total counts per class.

I've been looking at the tracker.py 'Track._count' variable to start off, am I in the right direction? https://github.com/GeekAlexis/FastMOT/blob/a185b78b771546151d42e02397f6940e4ee47bd8/fastmot/track.py#L223-L225 How do I get the info to print on the video stream rather than the terminal?

Many thanks for your help!

AngelaYZhang commented 2 years ago

If anyone else is interested, I've added to _draw function in mot.py to print total cumulative count. Please correct me if it's incorrect :)

def _draw(self, frame, detections):
        visible_tracks = list(self.visible_tracks())
        self.visualizer.render(frame, visible_tracks, detections, self.tracker.klt_bboxes.values(),
                               self.tracker.flow.prev_bg_keypoints, self.tracker.flow.bg_keypoints)        
        cv2.putText(frame, f'Visible: {len(visible_tracks)}', (30, 30),
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (255,255,255), 2, cv2.LINE_AA)

        #Print the total cumulative count
        track_list = list(track.trk_id for track in self.tracker.tracks.values())
        cv2.putText(frame, f'Total: {max(track_list)}', (30, 60),
                     cv2.FONT_HERSHEY_SIMPLEX, 1, (255,255,255), 2, cv2.LINE_AA)