Merge branch 'shaokai/add_notebooks'

Author: yeshaokai

amadeusgpt/analysis_objects/llm.py

@@ -8,6 +8,8 @@
 import openai
 from openai import OpenAI
 import base64
+import cv2
+import io
 
 class LLM(AnalysisObject):
     total_tokens = 0
@@ -76,11 +78,8 @@ def connect_gpt_oai_1(self, messages, **kwargs):
                     "model": self.gpt_model,
                     "messages": messages,
                     "max_tokens": self.max_tokens,
-                    "stop": None,
-                    "top_p": 1,
                     "temperature": 0.0,
                 }
-
                 response = client.chat.completions.create(**json_data)
 
                 LLM.total_tokens =  LLM.total_tokens + response.usage.prompt_tokens + response.usage.completion_tokens
@@ -121,36 +120,32 @@ def update_history(self, role, content, encoded_image = None, replace=False):
                 self.context_window.append({"role": role, "content": content})
         else:
 
+            if encoded_image is None:
+                self.history.append({"role": role, "content": content})
+                num_AI_messages = (len(self.context_window) - 1) // 2
+                if num_AI_messages == self.keep_last_n_messages:
+                    print ("doing active forgetting")
+                    # we forget the oldest AI message and corresponding answer
+                    self.context_window.pop(1)
+                    self.context_window.pop(1)
+                new_message = {"role": role, "content": content}
+            else:
+                new_message = {"role": "user", "content": [
+                    {"type": "text", "text": ""},
+                    {"type": "image_url", "image_url": {
+                        "url": f"data:image/jpeg;base64,{encoded_image}"}
+                    }
+                ]}
+
+            self.history.append(new_message)
+            
             if replace == True:
-                if len(self.history) == 2:
-                    self.history[1]["content"] = content
-                    self.context_window[1]["content"] = content
+                if len(self.context_window) == 2:
+                    self.context_window[1] = new_message
                 else:
-                    self.history.append({"role": role, "content": content})
-                    self.context_window.append({"role": role, "content": content})
+                    self.context_window.append(new_message)
+                
 
-            else:
-                if encoded_image is None:
-                    self.history.append({"role": role, "content": content})
-                    num_AI_messages = (len(self.context_window) - 1) // 2
-                    if num_AI_messages == self.keep_last_n_messages:
-                        print ("doing active forgetting")
-                        # we forget the oldest AI message and corresponding answer
-                        self.context_window.pop(1)
-                        self.context_window.pop(1)
-                    self.context_window.append({"role": role, "content": content})
-                else:
-                    message = {
-                        "role": "user", "content": [
-                        {"type": "text", "text": content},
-                        {"type": "image_url", "image_url": {
-                        "url": f"data:image/png;base64,{encoded_image}"}
-                        }]
-                    }                                            
-                    self.context_window.append(message) 
-    
-                    
-               
 
     def clean_context_window(self):
         while len(self.context_window) > 1:
@@ -194,13 +189,26 @@ def speak(self, sandbox):
 
         """
 
-        from amadeusgpt.system_prompts.visual import _get_system_prompt
+        from amadeusgpt.system_prompts.visual_llm import _get_system_prompt
         self.system_prompt = _get_system_prompt()
         analysis = sandbox.exec_namespace["behavior_analysis"]
         scene_image = analysis.visual_manager.get_scene_image()
-        encoded_image = self.encode_image(scene_image)
-        self.update_history("user", encoded_image)
 
+        result, buffer = cv2.imencode('.jpeg', scene_image)     
+        image_bytes = io.BytesIO(buffer)
+        base64_image = base64.b64encode(image_bytes.getvalue()).decode('utf-8')       
+        self.update_history("system", self.system_prompt)
+        self.update_history("user", "here is the image", encoded_image = base64_image, replace = True)
+        response = self.connect_gpt(self.context_window, max_tokens=2000)        
+        text = response.choices[0].message.content.strip()
+        print (text)
+        pattern = r"```json(.*?)```"
+        if len(re.findall(pattern, text, re.DOTALL)) == 0:
+            raise ValueError("can't parse the json string correctly", text)
+        else:
+            json_string = re.findall(pattern, text, re.DOTALL)[0]
+            json_obj = json.loads(json_string)
+            return json_obj
 
 class CodeGenerationLLM(LLM):
     """
@@ -394,3 +402,14 @@ def speak(self, sandbox):
         function_code = re.findall(pattern, text, re.DOTALL)[0]
         qa_message["code"] = function_code
         qa_message["chain_of_thought"] = thought_process
+
+
+if __name__ == "__main__":
+    from amadeusgpt.config import Config   
+    from amadeusgpt.main import create_amadeus
+    config = Config("amadeusgpt/configs/EPM_template.yaml")
+
+    amadeus = create_amadeus(config)
+    sandbox = amadeus.sandbox
+    visualLLm = VisualLLM(config)
+    visualLLm.speak(sandbox)

amadeusgpt/app.py

@@ -23,6 +23,7 @@ def main():
             st.session_state["exist_valid_openai_api_key"] = True
         else:
             st.session_state["exist_valid_openai_api_key"] = False
+    
 
     example_to_page = {}
 

amadeusgpt/app_utils.py

@@ -548,25 +548,28 @@ def render_page_by_example(example):
 
     if example == "Custom":
         st.markdown("Provide your own video and keypoint file (in pairs)")
-        uploaded_keypoint_file = st.file_uploader(
-            "Choose keypoint files",
-            ["h5"],
-            accept_multiple_files=False,
-        )
-        uploaded_video_file = st.file_uploader(
-            "Choose video files",
-            VIDEO_EXTS,
-            accept_multiple_files=False,
-        )
-
-        if uploaded_keypoint_file is not None:
-            save_dir = os.path.join("examples", example)
-            config["keypoint_info"]["keypoint_file_path"] = save_uploaded_file(
-                uploaded_keypoint_file, save_dir
+        save_dir = os.path.join("examples", example)
+        if "uploaded_keypoint_file" not in st.session_state:
+            uploaded_keypoint_file = st.file_uploader(
+                "Choose keypoint files",
+                ["h5"],
+                accept_multiple_files=False,
             )
-            config["video_info"]["video_file_path"] = save_uploaded_file(
-                uploaded_video_file, save_dir
+            if uploaded_keypoint_file is not None:
+                path = save_uploaded_file(uploaded_keypoint_file, save_dir)
+                st.session_state["uploaded_keypoint_file"] = path
+                config["keypoint_info"]["keypoint_file_path"] = st.session_state["uploaded_keypoint_file"]
+
+        if "uploaded_video_file" not in st.session_state:
+            uploaded_video_file = st.file_uploader(
+                "Choose video files",
+                VIDEO_EXTS,
+                accept_multiple_files=False,
             )
+            if uploaded_video_file is not None:
+                path = save_uploaded_file(uploaded_video_file, save_dir)
+                st.session_state["uploaded_video_file"] = uploaded_video_file
+                config["video_info"]["video_file_path"] = st.session_state["uploaded_video_file"]      
 
         ###### USER INPUT PANEL ######
         # get user input once getting the uploaded files
@@ -648,7 +651,6 @@ def render_page_by_example(example):
     st.session_state["example"] = example
 
     scene_image_path = get_scene_image(config)
-    video_file = config["video_info"]["video_file_path"]
 
     if scene_image_path is not None:
         img_data = base64.b64decode(scene_image_path)
@@ -673,8 +675,12 @@ def render_page_by_example(example):
             st.caption("Raw video from Horse-30")
         else:
             st.caption("DeepLabCut-SuperAnimal tracked video")
-        if video_file:
-            st.video(video_file)
+        if config["video_info"]["video_file_path"] and config["video_info"]["video_file_path"] is not None: 
+            st.video(config["video_info"]["video_file_path"])
+
+        if "uploaded_video_file" in st.session_state:
+            st.video(st.session_state["uploaded_video_file"])
+        
         # we only show objects for MausHaus for demo
         # if sam_image is not None:
         #     st.caption("SAM segmentation results")

amadeusgpt/main.py

@@ -11,7 +11,7 @@
 import os
 
 from amadeusgpt.analysis_objects.llm import (CodeGenerationLLM, DiagnosisLLM,
-                                             SelfDebugLLM)
+                                             SelfDebugLLM, VisualLLM)
 from amadeusgpt.integration_module_hub import IntegrationModuleHub
 
 amadeus_fac = {}
@@ -31,6 +31,7 @@ def __init__(self, config: Dict[str, Any]):
         self.code_generator_llm = CodeGenerationLLM(config.get("llm_info", {}))
         self.self_debug_llm = SelfDebugLLM(config.get("llm_info", {}))
         self.diagnosis_llm = DiagnosisLLM(config.get("llm_info", {}))
+        self.visual_llm = VisualLLM(config.get("llm_info", {}))
         ### fields that decide the behavior of the application
         self.use_self_debug = True
         self.use_diagnosis = False
@@ -47,11 +48,15 @@ def __init__(self, config: Dict[str, Any]):
 
         ## register the llm to the sandbox
         self.sandbox.register_llm("code_generator", self.code_generator_llm)
+        self.sandbox.register_llm("visual_llm", self.visual_llm)
         if self.use_self_debug:
             self.sandbox.register_llm("self_debug", self.self_debug_llm)
         if self.use_diagnosis:
             self.sandbox.register_llm("diagnosis", self.diagnosis_llm)
 
+        # can only do this after the register process
+        self.sandbox.configure_using_vlm()
+            
     def match_integration_module(self, user_query: str):
         """
         Return a list of matched integration modules
@@ -80,21 +85,19 @@ def step(self, user_query):
         result = self.sandbox.llm_step(user_query)
         return result
 
+    def get_analysis(self):
+        sandbox = self.sandbox
+        analysis = sandbox.exec_namespace['behavior_analysis']
+        return analysis
 
-if __name__ == "__main__":
-    config = Config("amadeusgpt/configs/MausHaus_template.yaml")
 
-    # amadeus = AMADEUS(config)
-    # query = "Give me events when mice are close"
-    # amadeus.step(query)
+if __name__ == "__main__":
+    from amadeusgpt.config import Config   
+    from amadeusgpt.main import create_amadeus
+    from amadeusgpt.analysis_objects.llm import VisualLLM
+    config = Config("amadeusgpt/configs/EPM_template.yaml")
 
-    query = "Plot the trajectory with the keypoint butt"
     amadeus = create_amadeus(config)
     sandbox = amadeus.sandbox
-    analysis = sandbox.exec_namespace["behavior_analysis"]
-
-    analysis.object_manager.load_roi_objects("temp_roi_objects.pickle")
-
-    from amadeusgpt.programs.sandbox import render_temp_message
-
-    render_temp_message(query, sandbox)
+    visualLLm = VisualLLM(config)
+    visualLLm.speak(sandbox)

amadeusgpt/managers/animal_manager.py

@@ -80,9 +80,24 @@ def __init__(self, config: Dict[str, str], model_manager: ModelManager):
         self.model_manager = model_manager
         self.animals: List[AnimalSeq] = []
         self.full_keypoint_names = []
-
+        self.superanimal_predicted_video = None
         self.init_pose()
 
+    def configure_animal_from_meta(self, meta_info):
+        """
+        Set the max individuals here
+        Set the superanimal model here
+        """
+        self.max_individuals = int(meta_info['individuals'])
+        species =  meta_info['species']
+        if species == 'topview_mouse':
+            self.superanimal_name = 'superanimal_topviewmouse_hrnetw32'
+        elif species == 'sideview_quadruped':
+            self.superanimal_name = 'superanimal_quadruped_hrnetw32'
+        else:
+            self.superanimal_name = None
+
+
     def init_pose(self):
         keypoint_info = self.config["keypoint_info"]
 
@@ -224,18 +239,26 @@ def get_keypoints(self) -> ndarray:
         video_file_path = self.config['video_info']['video_file_path']
         if os.path.exists(video_file_path) and keypoint_file_path is None:
 
+            if self.superanimal_name is None:
+                raise ValueError("Couldn't determine the species of the animal from the image. Change the scene index")
+
+            # only import here because people who choose the minimal installation might not have deeplabcut 
             import deeplabcut
-            from deeplabcut.modelzoo.video_inference import video_inference_superanimal
-            superanimal_name = 'superanimal_topviewmouse_hrnetw32'
+            from deeplabcut.modelzoo.video_inference import video_inference_superanimal            
+            video_suffix = Path(video_file_path).suffix
+            
+            keypoint_file_path = video_file_path.replace(video_suffix, '_' + self.superanimal_name + '.h5')            
+            self.superanimal_predicted_video = keypoint_file_path.replace('.h5', '_labeled.mp4')
 
-            keypoint_file_path = video_file_path.replace('.mp4', '_' + superanimal_name + '.h5')
             if not os.path.exists(keypoint_file_path):
+                print (f"going to inference video with {self.superanimal_name}")
                 video_inference_superanimal(videos = [self.config['video_info']['video_file_path']],
-                                            superanimal_name = superanimal_name,
+                                            superanimal_name = self.superanimal_name,
+                                            max_individuals=self.max_individuals,
                                             video_adapt = False)
+
 
             if os.path.exists(keypoint_file_path):
-
                 self.config['keypoint_info']['keypoint_file_path'] = keypoint_file_path
                 self.init_pose()
 

amadeusgpt/managers/gui_manager.py

@@ -11,6 +11,7 @@
                                               register_core_api)
 
 from .base import Manager
+from amadeusgpt.analysis_objects.object import ROIObject
 from .object_manager import ObjectManager
 
 
@@ -50,9 +51,9 @@ def onselect(self, vertices):
 
         # Here you can add any further processing of the polygons
         self.object_manager.roi_objects = []
-        self.object_manager.add_roi_object(self.paths)
+        for idx, path in enumerate(self.paths):
+            self.object_manager.add_roi_object(ROIObject(f'ROI{idx}', path))
 
-        print(len(self.object_manager.roi_objects))
         # Assuming the object_manager's add_roi_object is meant to handle the completed polygons
 
 

amadeusgpt/managers/object_manager.py

@@ -60,6 +60,7 @@ def load_objects_from_disk(self):
         pass
 
     def get_roi_object_names(self) -> List[str]:
+        
         return [obj.name for obj in self.roi_objects]
 
     def get_roi_objects(self) -> List[Object]:
@@ -80,9 +81,7 @@ def add_roi_object(self, object) -> None:
         self.roi_objects = self.filter_duplicates(self.roi_objects)
 
     def save_roi_objects(self, path: str) -> None:
-        roi_obects = self.get_roi_objects()       
-        for roi in roi_obects:
-            print(roi.name)
+        roi_obects = self.get_roi_objects()               
         data = {}
         for obj in roi_obects:
             data[obj.name] = {"Path": obj.Path}