From 0b763b3dd9ffef55020d5993549263b9abe68d30 Mon Sep 17 00:00:00 2001
From: George Ni <nemonameless@qq.com>
Date: Thu, 10 Jun 2021 16:07:07 +0800
Subject: [PATCH] [MOT] decord video infer (#3251)

* decord video infer

* add video infer write results
---
 ppdet/data/source/mot.py | 25 ++++++++++---------------
 ppdet/engine/tracker.py  |  2 ++
 requirements.txt         |  3 ++-
 3 files changed, 14 insertions(+), 16 deletions(-)

diff --git a/ppdet/data/source/mot.py b/ppdet/data/source/mot.py
index 1e61f4408..9e34241b4 100644
--- a/ppdet/data/source/mot.py
+++ b/ppdet/data/source/mot.py
@@ -13,8 +13,8 @@
 # limitations under the License.
 
 import os
-import cv2
 import numpy as np
+import decord as de
 from collections import OrderedDict
 try:
     from collections.abc import Sequence
@@ -317,26 +317,23 @@ class MOTVideoDataset(DetDataset):
         self.dataset_dir = dataset_dir
         self.keep_ori_im = keep_ori_im
         self.roidbs = None
+        self.frame_rate = 25
 
     def parse_dataset(self, ):
         if not self.roidbs:
             self.roidbs = self._load_video_images()
 
     def _load_video_images(self):
-        self.cap = cv2.VideoCapture(self.video_file)
-        self.vn = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
-        self.frame_rate = int(round(self.cap.get(cv2.CAP_PROP_FPS)))
-        logger.info('Length of the video: {:d} frames'.format(self.vn))
-        res = True
-        ct = 0
+        self.video_frames = de.VideoReader(self.video_file)
+        self.video_length = len(self.video_frames)
+        logger.info('Length of the video: {:d} frames.'.format(
+            self.video_length))
         records = []
-        while res:
-            res, img = self.cap.read()
-            image = np.ascontiguousarray(img, dtype=np.float32)
-            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        for idx in range(self.video_length):
+            image = self.video_frames.get_batch([idx]).asnumpy()[0]
             im_shape = image.shape
             rec = {
-                'im_id': np.array([ct]),
+                'im_id': np.array([idx]),
                 'image': image,
                 'h': im_shape[0],
                 'w': im_shape[1],
@@ -347,10 +344,8 @@ class MOTVideoDataset(DetDataset):
             }
             if self.keep_ori_im:
                 rec.update({'ori_image': image})
-            ct += 1
             records.append(rec)
-        records = records[:-1]
-        assert len(records) > 0, "No image file found"
+        assert len(records) > 0, "No image file found."
         return records
 
     def set_video(self, video_file):
diff --git a/ppdet/engine/tracker.py b/ppdet/engine/tracker.py
index 0d72be36e..9144a63d1 100644
--- a/ppdet/engine/tracker.py
+++ b/ppdet/engine/tracker.py
@@ -372,6 +372,8 @@ class Tracker(object):
         else:
             raise ValueError(model_type)
 
+        self.write_mot_results(result_filename, results, data_type)
+
         if save_videos:
             output_video_path = os.path.join(save_dir, '..',
                                              '{}_vis.mp4'.format(seq))
diff --git a/requirements.txt b/requirements.txt
index 864e1b138..d8864d8da 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -13,4 +13,5 @@ setuptools>=42.0.0
 lap
 sklearn
 motmetrics
-openpyxl
\ No newline at end of file
+openpyxl
+decord
\ No newline at end of file
-- 
GitLab