Change result cache to defaultdict(dict) and added a new notebook for task program library

Author: yeshaokai

amadeusgpt/analysis_objects/analysis_factory.py

@@ -1,5 +1,6 @@
 from amadeusgpt.implementation import AnimalBehaviorAnalysis
 
+
 analysis_fac = {}
 
 def create_analysis(config):

amadeusgpt/configs/MABe_template.yaml

@@ -12,5 +12,5 @@ llm_info:
 object_info:
   load_objects_from_disk: false
 video_info:
-  scene_frame_number: 100
+  scene_frame_number: 1400
   video_file_path: "examples/MABe/EGS8X2MN4SSUGFWAV976.mp4"

amadeusgpt/main.py

@@ -17,7 +17,7 @@
 from amadeusgpt.analysis_objects.llm import (CodeGenerationLLM, DiagnosisLLM,
                                              SelfDebugLLM, VisualLLM)
 from amadeusgpt.integration_module_hub import IntegrationModuleHub
-
+from collections import defaultdict
 import pickle 
 
 class AMADEUS:
@@ -84,11 +84,16 @@ def get_analysis(self):
         analysis = self.sandbox.exec_namespace["behavior_analysis"]
         return analysis
 
-    def run_task_program(self, task_program_name: str):
+    def run_task_program(self, config: Config, task_program_name: str):
         """
         Execute the task program on the currently holding sandbox
+        Parameters
+        -----------
+        config: a config specifies the movie file and the keypoint file to run task program
+        task_program_name: the name of the task program to run
+
         """
-        return self.sandbox.run_task_program(task_program_name)
+        return self.sandbox.run_task_program(config, task_program_name)
 
     def save_results(self, out_folder: str| None = None):
         """
@@ -102,9 +107,10 @@ def save_results(self, out_folder: str| None = None):
         os.makedirs(result_folder, exist_ok=True)
         results = self.sandbox.result_cache
 
-        ret = {}
+        ret = defaultdict(dict)
         for query in results:
-            ret[query] = results[query].get_serializable()
+            for video_file_path in results[query]:
+                ret[query][video_file_path] = results[query][video_file_path].get_serializable()
 
         # save results to a pickle file
         with open (os.path.join(result_folder, "results.pickle"), "wb") as f:

amadeusgpt/managers/visual_manager.py

@@ -673,8 +673,6 @@ def write_video(self, out_folder, video_file_path, out_name, events):
         # Release everything when job is finished
         cap.release()
         cv2.destroyAllWindows()
-        print("out videos" * 10)
-        print(out_videos)
         return out_videos
 
     def generate_video_clips_from_events(

amadeusgpt/programs/sandbox.py

@@ -17,7 +17,7 @@
 from amadeusgpt.programs.task_program_registry import (TaskProgram,
                                                        TaskProgramLibrary)
 from pathlib import Path
-
+from collections import defaultdict
 
 class QA_Message(dict):
     def __init__(self, *args, **kwargs):
@@ -233,8 +233,17 @@ def __init__(self, config):
         # just easier to pass this around
         self.query = None
         self.matched_modules = []
-        # result cache keeps the qa_message using the query as the key:
-        self.result_cache = {}
+        # example result_cahe
+        """
+        {'query' :  
+            {
+                'file1.mp4':  QA_Message(),
+                'file2.mp4':  QA_Message(),
+            }
+        }         
+        """
+
+        self.result_cache = defaultdict(dict)
         # configure how to save the results to a result folder
         self.result_folder = Path(self.config["result_info"].get("result_folder", "./results"))
 
@@ -497,24 +506,30 @@ def llm_step(self, user_query):
         self.messages.append(qa_message)
         # there might be better way to set this
         self.query = user_query
-        self.llms["code_generator"].speak(self)
-        self.result_cache[user_query] = qa_message
+        self.llms["code_generator"].speak(self)        
+        self.result_cache[user_query][self.config['video_info']['video_file_path']] = qa_message
         return qa_message
 
-    def run_task_program(self, task_program_name):
+    def run_task_program(self, config: Config, task_program_name: str):
         """
         1) sandbox is also responsible for running task program
         2) self.task_program_library references to a singleton so a different sandbox still has reference to the task program
         """
+        # update the config 
+        self.config = config
+
         task_program = self.task_program_library[task_program_name]
         # there might be better way to set this
         self.query = task_program_name
         qa_message = create_message(self.query, self)
         qa_message["code"] = task_program["source_code"]
         self.messages.append(qa_message)
+
+        # code execution will use the latest config, if updated
         self.code_execution(qa_message)
+
         qa_message = self.render_qa_message(qa_message)
-        self.result_cache[task_program_name] = qa_message
+        self.result_cache[task_program_name][config['video_info']['video_file_path']] = qa_message
         return qa_message
 
     def step(self, user_query, number_of_debugs=1):
@@ -539,7 +554,7 @@ def step(self, user_query, number_of_debugs=1):
                 qa_message = self.code_execution(qa_message)
 
         qa_message = self.render_qa_message(qa_message)
-        self.result_cache[user_query] = qa_message
+        self.result_cache[user_query][self.config['video_info']['video_file_path']] = qa_message
         return qa_message