Explorar o código

support batch and webm for videoclas

liuhongen1234567 hai 10 meses
pai
achega
0dfe515326

+ 1 - 1
paddlex/inference/common/batch_sampler/video_batch_sampler.py

@@ -25,7 +25,7 @@ from .base_batch_sampler import BaseBatchSampler
 
 
 class VideoBatchSampler(BaseBatchSampler):
 class VideoBatchSampler(BaseBatchSampler):
 
 
-    SUFFIX = ["mp4", "avi", "mkv"]
+    SUFFIX = ["mp4", "avi", "mkv", "webm"]
 
 
     # XXX: auto download for url
     # XXX: auto download for url
     def _download_from_url(self, in_path):
     def _download_from_url(self, in_path):

+ 10 - 4
paddlex/inference/models_new/video_classification/predictor.py

@@ -18,11 +18,17 @@ from ....modules.video_classification.model_list import MODELS
 from ...common.batch_sampler import VideoBatchSampler
 from ...common.batch_sampler import VideoBatchSampler
 from ...common.reader import ReadVideo
 from ...common.reader import ReadVideo
 from ..common import (
 from ..common import (
-    ToBatch,
     StaticInfer,
     StaticInfer,
 )
 )
 from ..base import BasicPredictor
 from ..base import BasicPredictor
-from .processors import Scale, CenterCrop, Image2Array, NormalizeVideo, VideoClasTopk
+from .processors import (
+    Scale,
+    CenterCrop,
+    Image2Array,
+    NormalizeVideo,
+    VideoClasTopk,
+    ToBatch,
+)
 from .result import TopkVideoResult
 from .result import TopkVideoResult
 
 
 
 
@@ -76,8 +82,8 @@ class VideoClasPredictor(BasicPredictor):
         batch_videos = self.pre_tfs["Scale"](videos=batch_raw_videos)
         batch_videos = self.pre_tfs["Scale"](videos=batch_raw_videos)
         batch_videos = self.pre_tfs["CenterCrop"](videos=batch_videos)
         batch_videos = self.pre_tfs["CenterCrop"](videos=batch_videos)
         batch_videos = self.pre_tfs["Image2Array"](videos=batch_videos)
         batch_videos = self.pre_tfs["Image2Array"](videos=batch_videos)
-        x = self.pre_tfs["NormalizeVideo"](videos=batch_videos)
-
+        batch_videos = self.pre_tfs["NormalizeVideo"](videos=batch_videos)
+        x = self.pre_tfs["ToBatch"](videos=batch_videos)
         batch_preds = self.infer(x=x)
         batch_preds = self.infer(x=x)
 
 
         batch_class_ids, batch_scores, batch_label_names = self.post_op["Topk"](
         batch_class_ids, batch_scores, batch_label_names = self.post_op["Topk"](

+ 18 - 4
paddlex/inference/models_new/video_classification/processors.py

@@ -115,10 +115,9 @@ class Scale:
                         if self.do_round
                         if self.do_round
                         else int(w * self.short_size / h)
                         else int(w * self.short_size / h)
                     )
                     )
-            if self.keep_ratio is not None:
-                resized_imgs.append(
-                    cv2.resize(img, (ow, oh), interpolation=cv2.INTER_LINEAR)
-                )
+            resized_imgs.append(
+                cv2.resize(img, (ow, oh), interpolation=cv2.INTER_LINEAR)
+            )
         imgs = resized_imgs
         imgs = resized_imgs
         return imgs
         return imgs
 
 
@@ -392,3 +391,18 @@ class VideoClasTopk:
         ]
         ]
         label_names = [[self.class_id_map[i] for i in index] for index in indexes]
         label_names = [[self.class_id_map[i] for i in index] for index in indexes]
         return indexes, scores, label_names
         return indexes, scores, label_names
+
+
+class ToBatch:
+    """A class for batching videos."""
+
+    def __call__(self, videos: List[np.ndarray]) -> List[np.ndarray]:
+        """Call method to stack videos into a batch.
+
+        Args:
+            videos (list of np.ndarrays): List of videos to process.
+
+        Returns:
+            list of np.ndarrays: List containing a stacked tensor of the videos.
+        """
+        return [np.concatenate(videos, axis=0).astype(dtype=np.float32, copy=False)]