updates

Author: yeshaokai

amadeusgpt/managers/visual_manager.py

@@ -550,21 +550,26 @@ def write_video(self, out_folder, video_file_path, out_name, events):
             return
 
         fourcc = cv2.VideoWriter_fourcc(*"avc1")  # Adjust the codec as needed
-        out = cv2.VideoWriter(
-            os.path.join(out_folder, f"{out_name}"),
-            fourcc,
-            30.0,
-            (int(cap.get(3)), int(cap.get(4))),
-        )
 
-        for triple in data:
+        out_videos = []        
+
+        for idx, triple in enumerate(data):
+
+            out = cv2.VideoWriter(
+                os.path.join(out_folder, f"{out_name}_{idx}"),
+                fourcc,
+                30.0,
+                (int(cap.get(3)), int(cap.get(4))),
+            )
+            out_videos.append(os.path.join(out_folder, f"{out_name}_{idx}"))
+            
             time_slice = triple["time_slice"]
             sender_animal_name = triple["sender_animal_name"]
             sender_keypoints = triple["sender_keypoints"]
             receiver_keypoints = triple["receiver_keypoints"]
             cap.set(cv2.CAP_PROP_POS_FRAMES, time_slice[0])
             offset = 0
-
+            
             while cap.isOpened():
                 current_frame = time_slice[0] + offset
                 ret, frame = cap.read()
@@ -669,6 +674,8 @@ def write_video(self, out_folder, video_file_path, out_name, events):
         cap.release()
         out.release()
         cv2.destroyAllWindows()
+        
+        return out_videos
 
     def generate_video_clips_from_events(
         self, out_folder, video_file, events: List[BaseEvent], behavior_name
@@ -683,5 +690,5 @@ def generate_video_clips_from_events(
         videoname = video_file.split("/")[-1].replace(".mp4", "").replace(".avi", "")
         video_name = f"{videoname}_{behavior_name}_video.mp4"
 
-        self.write_video(out_folder, video_file, video_name, events)
-        return os.path.join(out_folder, video_name)
+        out_videos = self.write_video(out_folder, video_file, video_name, events)
+        return out_videos

amadeusgpt/programs/sandbox.py

@@ -31,7 +31,7 @@ def create_message(query, sandbox):
         "error_message": None,
         "function_rets": None,
         "sandbox": sandbox,
-        "out_videos": [],
+        "out_videos": None,
         "pose_video": None,
         "meta_info": None
     }
@@ -424,9 +424,8 @@ def render_qa_message(self, qa_message):
                                 bodypart_names=bodypart_names, events=e
                             )
                         )
-                        qa_message["out_videos"].append(
-                            self.events_to_videos(e, self.get_function_name_from_string(qa_message["code"]))
-                        )
+                        qa_message["out_videos"] = self.events_to_videos(e, self.get_function_name_from_string(qa_message["code"]))
+                        
 
         elif (
             isinstance(function_rets, list)
@@ -442,9 +441,8 @@ def render_qa_message(self, qa_message):
             plots.append(
                 visual_manager.get_ethogram_visualization(events=function_rets)
             )
-            qa_message["out_videos"].append(
-                self.events_to_videos(function_rets, self.get_function_name_from_string(qa_message["code"]))
-            )
+            qa_message["out_videos"] = self.events_to_videos(function_rets, self.get_function_name_from_string(qa_message["code"]))
+            
         else:
             pass
         qa_message["plots"].extend(plots)