fix(preprocessor): fix bug

gnes-ai · Aug 12, 2019 · 415456d · 415456d
1 parent c150ad5
commit 415456d
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/gnes/preprocessor/video/ffmpeg.py b/gnes/preprocessor/video/ffmpeg.py
@@ -161,8 +161,8 @@ def apply(self, doc: 'gnes_pb2.Document') -> None:
                 elif self.segment_method == 'cut_by_clustering':
                     if self.segment_num >= 2:
                         hash_v = [phash_descriptor(_).hash for _ in frames]
-                        label_v = KMeans(n_clusters=self.segment_num
-                                         ).fit_predict(np.array(hash_v, dtype=np.int32))
+                        hash_v = np.array(hash_v, dtype=np.int32).reshape([len(hash_v), -1])
+                        label_v = KMeans(n_clusters=self.segment_num).fit_predict(hash_v)
                         sub_videos = [[frames[i] for i, j in enumerate(label_v) if j == _] for _ in range(self.segment_num)]
                     else:
                         sub_videos = [frames]

diff --git a/tests/test_video_preprocessor.py b/tests/test_video_preprocessor.py
@@ -14,7 +14,7 @@ def setUp(self):
         self.yml_path = os.path.join(self.dirname, 'yaml', 'preprocessor-ffmpeg.yml')
         self.yml_path_2 = os.path.join(self.dirname, 'yaml', 'preprocessor-ffmpeg2.yml')
         self.yml_path_3 = os.path.join(self.dirname, 'yaml', 'preprocessor-ffmpeg3.yml')
-        self.yml_path_3 = os.path.join(self.dirname, 'yaml', 'preprocessor-ffmpeg4.yml')
+        self.yml_path_4 = os.path.join(self.dirname, 'yaml', 'preprocessor-ffmpeg4.yml')
         self.video_path = os.path.join(self.dirname, 'videos')
         self.video_bytes = [open(os.path.join(self.video_path, _), 'rb').read()
                             for _ in os.listdir(self.video_path)]