From b79ab602d53e9e37a15553fdb7e9a7e78796a591 Mon Sep 17 00:00:00 2001
From: Kaiyuan Hu <kaiyuan.hu@zilliz.com>
Date: Mon, 26 Sep 2022 17:33:51 +0800
Subject: [PATCH] update select

Signed-off-by: Kaiyuan Hu <kaiyuan.hu@zilliz.com>
---
 select_video.py | 28 ++++++++++++++++++++--------
 1 file changed, 20 insertions(+), 8 deletions(-)

diff --git a/select_video.py b/select_video.py
index 7ff9b3b..a8413f6 100644
--- a/select_video.py
+++ b/select_video.py
@@ -25,18 +25,13 @@ class SelectVideo:
         self.reduce_function = reduce_function
         self.reverse = reverse
 
-    def __call__(self, video_urls: List[str], scores: List[float]) -> List[str]:
-        """
-        args:
-            video_urls
-            scores
-        return:
-            video_list
-        """
+    def _select(self, video_urls, scores):
         assert len(video_urls) == len(scores), 'len(video_urls) must equal len(scores)'
+
         video2score_list = [{'video_url': video_url, 'score': score} for video_url, score in zip(video_urls, scores)]
         video2score_list.sort(key=lambda d: d['video_url'])
         reduced_dict = {}
+
         for video_url, items in groupby(video2score_list, key=lambda d: d['video_url']):
             video_score_list = [item['score'] for item in items]
             if self.reduce_function == 'sum':
@@ -50,10 +45,27 @@ class SelectVideo:
             else:
                 raise Exception('unknown reduce_function')
             reduced_dict[video_url] = reduced_value
+
         sorted_video_urls = [k for k, v in sorted(reduced_dict.items(), key=lambda item: item[1], reverse=self.reverse)]
         top_k = min(self.top_k, len(sorted_video_urls))
+
         return sorted_video_urls[:top_k]
 
+    def __call__(self, video_urls: List[str], scores: List[float]) -> List[str]:
+        """
+        args:
+            video_urls
+            scores
+        return:
+            video_list
+        """
+        if isinstance(video_urls[0], list):
+            video_urls = sum(video_urls, [])
+            scores = sum(scores, [])
+        if isinstance(video_urls[0], str):
+            return self._select(video_urls, scores)
+
+
 # if __name__ == '__main__':
 #     op = SelectVideo(top_k=2, reverse=True, reduce_function='min')
 #     res = op(['a', 'a', 'c', 'a', 'b', 'b', 'c', 'c'],