updates

Author: Ye Shaokai

action/chatgpt_utils.py

@@ -149,14 +149,21 @@ class GPTInferenceAnnotator(ChatGPT):
     Given the images, this class will annotate the video frames
     """
 
-    def __init__(self, root, prediction_save_folder, clip_length = 4, debug = False):
+    def __init__(self, 
+                 root, 
+                 prediction_save_folder, 
+                 clip_length = 4, 
+                 debug = False,
+                 topk = 10
+                 ):
         super().__init__(clip_length = clip_length)
         self.root = root
         self.prediction_save_folder = prediction_save_folder 
         self.prediction_analysis = PredictionAnalysis(self.prediction_save_folder)
         self.prediction_analysis.load()
         self.data = self.prediction_analysis.data
         self.debug = debug
+        self.topk = topk
 
     def multi_process_run(self):
         prediction_analysis = PredictionAnalysis(self.prediction_save_folder)
@@ -187,6 +194,8 @@ def parse_item(self, item):
 
         gt_name = item['gt_name']
         avion_predictions = item['avion_preds']['predictions']
+        assert self.topk <= len(avion_predictions)
+        avion_predictions = avion_predictions[:self.topk]
         # _avion_predictions = [e.replace(':', ' ', 1) for e in avion_predictions]
         # if gt_name not in _avion_predictions:
         #     print ('gt_name not in avion_predictions')
@@ -458,12 +467,18 @@ def explore_wrong_examples(root, prediction_save_folder):
                                     debug = True)
     annotator.explore_wrong_examples()
 
-def multi_process_inference(root, prediction_save_folder):
+def multi_process_inference(root, 
+                            prediction_save_folder, 
+                            clip_length = 4,
+                            topk = 10, 
+                            debug = False):
 
     annotator = GPTInferenceAnnotator(root, 
     prediction_save_folder, 
-    clip_length = 4,
-    debug = True)
+    clip_length = clip_length,
+    debug = debug,
+    topk = topk)
+
     annotator.multi_process_run()
 
 if __name__ == '__main__':
@@ -473,7 +488,8 @@ def multi_process_inference(root, prediction_save_folder):
     root = '/data/EK100/EK100_320p_15sec_30fps_libx264'    
     pred_folder = '/data/epic_kitchen/llavavideo_avion_mc_top10_5epoch_preds'
 
-    multi_process_annotate(train_file_path, root)
+    #multi_process_annotate(train_file_path, root)
     #explore_wrong_examples(root, pred_folder)
-
-    #multi_process_inference(root, pred_folder)
+    multi_process_inference(root, pred_folder, debug = False,
+                            clip_length = 4,
+                            topk = 5)