Made it easy to create a project

Author: yeshaokai

amadeusgpt/__init__.py

@@ -10,7 +10,7 @@
 from amadeusgpt.integration_modules import *
 from amadeusgpt.main import AMADEUS
 from amadeusgpt.version import VERSION, __version__
-
+from amadeusgpt.project import create_project
 params = {
     "axes.labelsize": 10,
     "legend.fontsize": 10,

amadeusgpt/main.py

@@ -2,6 +2,7 @@
 
 from amadeusgpt.config import Config
 from amadeusgpt.programs.sandbox import Sandbox
+import yaml
 ##########
 # all these are providing the customized classes for the code execution
 ##########
@@ -20,7 +21,7 @@
 
 
 class AMADEUS:
-    def __init__(self, config: Config, use_vlm=True):
+    def __init__(self, config: Config | dict, use_vlm=True):
         self.config = config
         ### fields that decide the behavior of the application
         self.use_self_debug = True
@@ -42,6 +43,9 @@ def __init__(self, config: Config, use_vlm=True):
         video_suffix = data_info["video_suffix"]
         video_file_paths = glob.glob(str(data_folder / f"*{video_suffix}"))
 
+        if len(video_file_paths) == 0:
+            print (f"No video files found in the data folder {data_folder}. Please check the data folder and the video suffix")
+            return 
         # optionally get the corresponding keypoint files
         keypoint_file_paths = self.get_DLC_keypoint_files(video_file_paths)
 
@@ -114,7 +118,13 @@ def step(self, user_query: str) -> QA_Message:
         return qa_message
 
     def get_video_file_paths(self) -> list[str]:
-        return self.sandbox.video_file_paths
+        data_info = self.config["data_info"]
+        data_folder = data_info['data_folder']
+        video_suffix = data_info['video_suffix']
+        video_file_paths = glob.glob(os.path.join(data_folder, f"*{video_suffix}"))
+
+        return video_file_paths
+
 
     def get_keypoint_file_paths(self) -> list[str]:
         return self.sandbox.keypoint_file_paths
@@ -143,12 +153,11 @@ def register_task_program(self, task_program, creator="human"):
         TaskProgramLibrary.register_task_program(creator=creator)(task_program)
 
     def get_messages(self):
-
         return self.sandbox.message_cache
 
     def get_task_programs(self):
         return TaskProgramLibrary.get_task_programs()
-
+                
 
 if __name__ == "__main__":
     from amadeusgpt.analysis_objects.llm import VisualLLM

amadeusgpt/programs/sandbox.py

@@ -165,7 +165,7 @@ def wrapper(*args, **kwargs):
 class Sandbox(SandboxBase):
     def __init__(
         self,
-        config: Config,
+        config: Config | dict,
         video_file_paths: list[str],
         keypoint_file_paths: list[str],
     ):

amadeusgpt/project.py

@@ -0,0 +1,38 @@
+from amadeusgpt.config import Config
+import os
+import pprint
+import yaml
+
+def create_project(data_folder, 
+                result_folder, 
+                video_suffix=".mp4"):
+    """
+    Create a project config file. Save the config file to the result folder
+    """
+    config = {
+        "data_info": {
+            "data_folder": data_folder,
+            "result_folder": result_folder,
+            "video_suffix": video_suffix,
+        },
+        "llm_info": {
+            "max_tokens": 4096,
+            "temperature": 0.0,
+            "keep_last_n_messages": 2
+        },
+    }
+    # save the dictionary config to yaml
+
+    os.makedirs(result_folder, exist_ok=True)
+
+    file_path = os.path.join(result_folder, "config.yaml")
+
+    with open(file_path, "w") as f:
+        yaml.dump(config, f)
+
+    print (f"Project created at {result_folder}. Results will be saved to {result_folder}")
+    print (f"The project will load video files (*.{video_suffix}) and optionally keypoint files from {data_folder}")        
+    print (f"A copy of the project config file is saved at {file_path}")
+    pprint.pprint(config)
+
+    return config

notebooks/custom_mouse_video.ipynb

@@ -24,7 +24,7 @@
     "from amadeusgpt.config import Config\n",
     "from amadeusgpt.utils import parse_result\n",
     "import amadeusgpt\n",
-    "from pathlib import Path\n",
+    "from amadeusgpt import create_project\n",
     "import matplotlib.pyplot as plt\n",
     "import cv2"
    ]
@@ -41,33 +41,48 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 3,
    "id": "be76dc87-fbe8-452f-b85c-2af3e95a03bf",
    "metadata": {},
    "outputs": [
     {
-     "name": "stderr",
+     "name": "stdout",
      "output_type": "stream",
      "text": [
-      "INFO:httpx:HTTP Request: POST https://api.openai.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n"
+      "Project created at temp_result_folder. Results will be saved to temp_result_folder\n",
+      "The project will load video files (*..mp4) and optionally keypoint files from temp_data_folder\n",
+      "A copy of the project config file is saved at temp_result_folder/config.yaml\n",
+      "{'data_info': {'data_folder': 'temp_data_folder',\n",
+      "               'result_folder': 'temp_result_folder',\n",
+      "               'video_suffix': '.mp4'},\n",
+      " 'llm_info': {'max_tokens': 4096, 'temperature': 1.0}}\n",
+      "No video files found in the data folder temp_data_folder. Please check the data folder and the video suffix\n"
      ]
     },
     {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "current total cost 0.004 $\n",
-      "current input tokens 25667\n",
-      "current accumulated tokens 25843\n",
-      "['/Users/shaokaiye/OFT/OFT_43_short.mp4']\n"
+     "ename": "AttributeError",
+     "evalue": "'AMADEUS' object has no attribute 'sandbox'",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[3], line 13\u001b[0m\n\u001b[1;32m     10\u001b[0m config[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mscene_frame_number\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m scene_frame_number\n\u001b[1;32m     12\u001b[0m amadeus \u001b[38;5;241m=\u001b[39m AMADEUS(config)\n\u001b[0;32m---> 13\u001b[0m video_file_paths \u001b[38;5;241m=\u001b[39m \u001b[43mamadeus\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_video_file_paths\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     14\u001b[0m \u001b[38;5;28mprint\u001b[39m (video_file_paths)\n",
+      "File \u001b[0;32m~/AmadeusGPT-dev/amadeusgpt/main.py:121\u001b[0m, in \u001b[0;36mAMADEUS.get_video_file_paths\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    120\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mget_video_file_paths\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28mlist\u001b[39m[\u001b[38;5;28mstr\u001b[39m]:\n\u001b[0;32m--> 121\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msandbox\u001b[49m\u001b[38;5;241m.\u001b[39mvideo_file_paths\n",
+      "\u001b[0;31mAttributeError\u001b[0m: 'AMADEUS' object has no attribute 'sandbox'"
      ]
     }
    ],
    "source": [
     "scene_frame_number = 400\n",
-    "amadeus_root = Path(amadeusgpt.__file__).parent.parent\n",
-    "config = Config(amadeus_root / \"amadeusgpt/configs/template.yaml\")\n",
-    "config['data_info']['data_folder'] = \"/Users/shaokaiye/OFT\"\n",
+    "\n",
+    "# where you store you video and (optionally) keypoint files\n",
+    "data_folder = \"temp_data_folder\"\n",
+    "result_folder = \"temp_result_folder\"\n",
+    "video_suffix = \".mp4\"\n",
+    "\n",
+    "config = create_project(data_folder, result_folder, video_suffix = video_suffix)\n",
+    "\n",
+    "config[\"scene_frame_number\"] = scene_frame_number\n",
     "\n",
     "amadeus = AMADEUS(config)\n",
     "video_file_paths = amadeus.get_video_file_paths()\n",

tests/test_project_creation.py

@@ -0,0 +1,15 @@
+import os
+if 'OPENAI_API_KEY' not in os.environ:  
+     os.environ['OPENAI_API_KEY'] = 'your key'
+
+from amadeusgpt import create_project
+from amadeusgpt import AMADEUS
+# Create a project
+
+data_folder = "temp_data_folder"
+result_folder = "temp_result_folder"
+
+config = create_project(data_folder, result_folder)
+
+# Create an AMADEUS instance
+amadeus = AMADEUS(config)