update global variance traj_list and camera_dict

app.py

@@ -13,9 +13,8 @@ from omegaconf import OmegaConf
 from PIL import Image
 from pytorch_lightning import seed_everything
 
-from gradio_utils.camera_utils import CAMERA_MOTION_MODE, process_camera, create_relative
-from gradio_utils.traj_utils import (OBJECT_MOTION_MODE, get_provided_traj,
-                                     process_points, process_traj)
+from gradio_utils.camera_utils import CAMERA_MOTION_MODE, process_camera
+from gradio_utils.traj_utils import (OBJECT_MOTION_MODE, process_traj)
 from gradio_utils.utils import vis_camera
 from lvdm.models.samplers.ddim import DDIMSampler
 from main.evaluation.motionctrl_inference import (DEFAULT_NEGATIVE_PROMPT,
@@ -23,7 +22,8 @@ from main.evaluation.motionctrl_inference import (DEFAULT_NEGATIVE_PROMPT,
                                                   post_prompt)
 from utils.utils import instantiate_from_config
 
-from gradio_utils.page_control import (MODE, BASE_MODEL, traj_list, camera_dict, 
+from gradio_utils.page_control import (MODE, BASE_MODEL, 
+                                       get_camera_dict, get_traj_list,
                                        reset_camera, 
                                        visualized_step1, visualized_step2,
                                        visualized_camera_poses, visualized_traj_poses,
@@ -31,7 +31,8 @@ from gradio_utils.page_control import (MODE, BASE_MODEL, traj_list, camera_dict,
                                        input_raw_camera_pose,
                                        change_camera_mode, change_camera_speed,
                                        add_traj_point, add_provided_traj, 
-                                       fn_traj_droplast, fn_traj_reset)
+                                       fn_traj_droplast, fn_traj_reset,
+                                       fn_vis_camera, fn_vis_traj,)
 
 os.environ['KMP_DUPLICATE_LIB_OK']='True'
 SPACE_ID = os.environ.get('SPACE_ID', '')
@@ -140,107 +141,6 @@ for i in range(0, 16):
     res.append(RT)
     
 fig = vis_camera(res)
-    
-def fn_vis_camera(info_mode, camera_args=None):
-    global camera_dict
-    RT = process_camera(camera_dict, camera_args) # [t, 3, 4]
-
-    rescale_T = 1.0
-    rescale_T = max(rescale_T, np.max(np.abs(RT[:,:,-1])) / 1.9)
-
-    fig = vis_camera(create_relative(RT), rescale_T=rescale_T)
-
-    if info_mode == MODE[0]:
-        vis_step3_prompt_generate = True
-        vis_prompt = True
-        vis_num_samples = True
-        vis_seed = True
-        vis_start = True
-        vis_gen_video = True
-
-        vis_object_mode = False
-        vis_object_info = False
-
-    else:
-        vis_step3_prompt_generate = False
-        vis_prompt = False
-        vis_num_samples = False
-        vis_seed = False
-        vis_start = False
-        vis_gen_video = False
-
-        vis_object_mode = True
-        vis_object_info = True
-
-    return fig, \
-            gr.update(visible=vis_object_mode), \
-            gr.update(visible=vis_object_info), \
-            gr.update(visible=vis_step3_prompt_generate), \
-            gr.update(visible=vis_prompt), \
-            gr.update(visible=vis_num_samples), \
-            gr.update(visible=vis_seed), \
-            gr.update(visible=vis_start), \
-            gr.update(visible=vis_gen_video, value=None)
-
-def fn_vis_traj():
-    global traj_list
-    global exp_no
-    xy_range = 1024
-    points = process_points(traj_list)
-    imgs = []
-    for idx in range(16):
-        bg_img = np.ones((1024, 1024, 3), dtype=np.uint8) * 255
-        for i in range(15):
-            p = points[i]
-            p1 = points[i+1]
-            cv2.line(bg_img, p, p1, (255, 0, 0), 2)
-
-            if i == idx:
-                cv2.circle(bg_img, p, 2, (0, 255, 0), 20)
-
-        if idx==(15):
-            cv2.circle(bg_img, points[-1], 2, (0, 255, 0), 20)
-        
-        imgs.append(bg_img.astype(np.uint8))
-
-    # size = (512, 512)
-    fps = 10
-
-    out_dir = f'./results_trajs/{exp_no}'
-    os.makedirs(out_dir, exist_ok=True)
-    exp_no += 1
-
-    traj_flow = process_traj(traj_list).transpose(3,0,1,2)
-
-    np.save(f'{out_dir}/traj_flow.npy', traj_flow)
-    with open(f'{out_dir}/traj_list.txt', 'w') as f:
-        for item in traj_list:
-            f.write(f"{item[0]}, {item[1]}\n")
-
-    if out_dir is None:
-        path = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False).name
-    else:
-        path = os.path.join(out_dir, 'traj.mp4')
-    writer = imageio.get_writer(path, format='mp4', mode='I', fps=fps)
-    for img in imgs:
-        writer.append_data(img)
-
-    writer.close()
-
-    vis_step3_prompt_generate = True
-    vis_prompt = True
-    vis_num_samples = True
-    vis_seed = True
-    vis_start = True
-    vis_gen_video = True
-    return path, gr.update(visible=vis_step3_prompt_generate), \
-                gr.update(visible=vis_prompt), \
-                gr.update(visible=vis_num_samples), \
-                gr.update(visible=vis_seed), \
-                gr.update(visible=vis_start), \
-                gr.update(visible=vis_gen_video, value=None)
-
-
 
 
 ###########################################
@@ -274,8 +174,8 @@ if torch.cuda.is_available():
 model_v2.eval()
 
 def model_run(prompts, choose_model, infer_mode, seed, n_samples, camera_args=None):
-    global traj_list
-    global camera_dict
+    traj_list = get_traj_list()
+    camera_dict = get_camera_dict()
 
     RT = process_camera(camera_dict, camera_args).reshape(-1,12)
     traj_flow = process_traj(traj_list).transpose(3,0,1,2)

gradio_utils/page_control.py

@@ -1,6 +1,11 @@
+import cv2
+import numpy as np
+import tempfile
+import imageio
 import gradio as gr
-from gradio_utils.camera_utils import CAMERA_MOTION_MODE
-from gradio_utils.traj_utils import get_provided_traj
+from gradio_utils.camera_utils import CAMERA_MOTION_MODE, process_camera, create_relative
+from gradio_utils.traj_utils import get_provided_traj, process_points
+from gradio_utils.utils import vis_camera
 
 MODE = ["control camera poses", "control object trajectory", "control both camera and object motion"]
 
@@ -17,6 +22,8 @@ def display_camera_info(camera_dict, camera_mode=None):
             res += f"mode : {camera_dict['mode']}. "
     return res
 
+global traj_list, camera_dict
+
 traj_list = [] 
 camera_dict = {
                 "motion":[],
@@ -25,8 +32,16 @@ camera_dict = {
                 "complex": None
                 }
 
+def get_traj_list():
+    global traj_list
+    return traj_list
+
+def get_camera_dict():
+    global camera_dict
+    return camera_dict
+
 def reset_camera():
-    # global camera_dict
+    global camera_dict
     camera_dict = {
                     "motion":[],
                     "mode": "Customized Mode 1: First A then B",
@@ -36,7 +51,7 @@ def reset_camera():
     return display_camera_info(camera_dict)
 
 def fn_traj_reset():
-    # global traj_list
+    global traj_list
     traj_list = []
     return "Click to specify trajectory"
 
@@ -478,7 +493,7 @@ def visualized_traj_poses(step2_object_motion):
             gr.update(visible=vis_gen_video)
 
 def add_camera_motion(camera_motion, camera_mode):  
-    # global camera_dict
+    global camera_dict
     if camera_dict['complex'] is not None:
         camera_dict['complex'] = None
     if camera_mode == CAMERA_MOTION_MODE[2] and len(camera_dict['motion']) <2:
@@ -489,7 +504,7 @@ def add_camera_motion(camera_motion, camera_mode):
     return display_camera_info(camera_dict, camera_mode)
 
 def add_complex_camera_motion(camera_motion):
-    # global camera_dict
+    global camera_dict
     camera_dict['complex']=camera_motion
     return display_camera_info(camera_dict)
 
@@ -522,7 +537,7 @@ def change_camera_mode(combine_type, camera_mode):
             gr.update(visible=vis_combine3_des)
 
 def input_raw_camera_pose(combine_type, camera_mode):
-    # global camera_dict
+    global camera_dict
     camera_dict['mode'] = combine_type
 
     vis_U = False
@@ -549,26 +564,26 @@ def input_raw_camera_pose(combine_type, camera_mode):
             gr.update(visible=vis_combine3_des)
 
 def change_camera_speed(camera_speed):
-    # global camera_dict
+    global camera_dict
     camera_dict['speed'] = camera_speed
     return display_camera_info(camera_dict)
 
 def add_traj_point(evt: gr.SelectData, ):
-    # global traj_list
+    global traj_list
     traj_list.append(evt.index)
     traj_str = [f"{traj}" for traj in traj_list]
     return ", ".join(traj_str)
 
 def add_provided_traj(traj_name):
-    # global traj_list
+    global traj_list
+    # import pdb; pdb.set_trace()
     traj_list = get_provided_traj(traj_name)
     traj_str = [f"{traj}" for traj in traj_list]
     return ", ".join(traj_str)
 
 
 def fn_traj_droplast():
-    # global traj_list
-
+    global traj_list
     if traj_list:
         traj_list.pop()
 
@@ -578,3 +593,87 @@ def fn_traj_droplast():
     else:   
         return "Click to specify trajectory"
 
+def fn_vis_camera(info_mode, camera_args=None):
+    global camera_dict
+    RT = process_camera(camera_dict, camera_args) # [t, 3, 4]
+
+    rescale_T = 1.0
+    rescale_T = max(rescale_T, np.max(np.abs(RT[:,:,-1])) / 1.9)
+
+    fig = vis_camera(create_relative(RT), rescale_T=rescale_T)
+
+    if info_mode == MODE[0]:
+        vis_step3_prompt_generate = True
+        vis_prompt = True
+        vis_num_samples = True
+        vis_seed = True
+        vis_start = True
+        vis_gen_video = True
+
+        vis_object_mode = False
+        vis_object_info = False
+
+    else:
+        vis_step3_prompt_generate = False
+        vis_prompt = False
+        vis_num_samples = False
+        vis_seed = False
+        vis_start = False
+        vis_gen_video = False
+
+        vis_object_mode = True
+        vis_object_info = True
+
+    return fig, \
+            gr.update(visible=vis_object_mode), \
+            gr.update(visible=vis_object_info), \
+            gr.update(visible=vis_step3_prompt_generate), \
+            gr.update(visible=vis_prompt), \
+            gr.update(visible=vis_num_samples), \
+            gr.update(visible=vis_seed), \
+            gr.update(visible=vis_start), \
+            gr.update(visible=vis_gen_video, value=None)
+
+def fn_vis_traj():
+    # import pdb; pdb.set_trace()
+    # global traj_list
+    # xy_range = 1024
+    # print(traj_list)
+    global traj_list
+    print(traj_list)
+    points = process_points(traj_list)
+    imgs = []
+    for idx in range(16):
+        bg_img = np.ones((1024, 1024, 3), dtype=np.uint8) * 255
+        for i in range(15):
+            p = points[i]
+            p1 = points[i+1]
+            cv2.line(bg_img, p, p1, (255, 0, 0), 2)
+
+            if i == idx:
+                cv2.circle(bg_img, p, 2, (0, 255, 0), 20)
+
+        if idx==(15):
+            cv2.circle(bg_img, points[-1], 2, (0, 255, 0), 20)
+        
+        imgs.append(bg_img.astype(np.uint8))
+
+    path = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False).name
+    writer = imageio.get_writer(path, format='mp4', mode='I', fps=10)
+    for img in imgs:
+        writer.append_data(img)
+    writer.close()
+
+    vis_step3_prompt_generate = True
+    vis_prompt = True
+    vis_num_samples = True
+    vis_seed = True
+    vis_start = True
+    vis_gen_video = True
+    return path, gr.update(visible=vis_step3_prompt_generate), \
+                gr.update(visible=vis_prompt), \
+                gr.update(visible=vis_num_samples), \
+                gr.update(visible=vis_seed), \
+                gr.update(visible=vis_start), \
+                gr.update(visible=vis_gen_video, value=None)
+