add random seed

app.py

@@ -11,6 +11,7 @@ import torch
 from torchvision import transforms
 import rembg
 import cv2
+from pytorch_lightning import seed_everything
 
 from src.visualizer import CameraVisualizer
 from src.pose_estimation import load_model_from_config, estimate_poses, estimate_elevs
@@ -113,7 +114,9 @@ def group_recenter(images, ratio=1.5, mask_thres=127, bkg_color=[255, 255, 255,
     return out_rgbs
 
 
-def run_preprocess(image1, image2, preprocess_chk):
+def run_preprocess(image1, image2, preprocess_chk, seed_value):
+
+    seed_everything(seed_value)
 
     if preprocess_chk:
         rembg_session = rembg.new_session()
@@ -138,7 +141,9 @@ def image_to_tensor(img, width=256, height=256):
 
 
 @spaces.GPU
-def run_pose_exploration_a(cam_vis, image1, image2):
+def run_pose_exploration_a(image1, image2, seed_value):
+
+    seed_everything(seed_value)
 
     image1 = image_to_tensor(image1).to(_device_)
     image2 = image_to_tensor(image2).to(_device_)
@@ -157,7 +162,9 @@ def run_pose_exploration_a(cam_vis, image1, image2):
 
 
 @spaces.GPU
-def run_pose_exploration_b(cam_vis, image1, image2, elevs, elev_ranges, probe_bsz, adj_bsz, adj_iters):
+def run_pose_exploration_b(cam_vis, image1, image2, elevs, elev_ranges, probe_bsz, adj_bsz, adj_iters, seed_value):
+
+    seed_everything(seed_value)
 
     noise = np.random.randn(probe_bsz, 4, 32, 32)
 
@@ -206,7 +213,9 @@ def run_pose_exploration_b(cam_vis, image1, image2, elevs, elev_ranges, probe_bs
 
 
 @spaces.GPU
-def run_pose_refinement(cam_vis, image1, image2, anchor_polar, explored_sph, refine_iters):
+def run_pose_refinement(cam_vis, image1, image2, anchor_polar, explored_sph, refine_iters, seed_value):
+
+    seed_everything(seed_value)
 
     cam_vis.set_images([np.asarray(image1, dtype=np.uint8), np.asarray(image2, dtype=np.uint8)])
 
@@ -295,7 +304,8 @@ def run_demo():
                 with gr.Accordion('Advanced options', open=False):
                     probe_bsz = gr.Slider(4, 32, value=16, step=4, label='Probe Batch Size')
                     adj_bsz = gr.Slider(1, 8, value=4, step=1, label='Adjust Batch Size')
-                    adj_iters = gr.Slider(1, 20, value=5, step=1, label='Adjust Iterations')
+                    adj_iters = gr.Slider(1, 20, value=8, step=1, label='Adjust Iterations')
+                    seed_value = gr.Number(value=0, label="Seed Value", precision=0)
 
                 with gr.Row():
                     run_btn = gr.Button('Estimate', variant='primary', interactive=True)
@@ -369,21 +379,21 @@ def run_demo():
 
         run_btn.click(
             fn=run_preprocess,
-            inputs=[input_image1, input_image2, preprocess_chk],
+            inputs=[input_image1, input_image2, preprocess_chk, seed_value],
             outputs=[processed_image1, processed_image2],
         ).success(
-            fn=partial(run_pose_exploration_a, cam_vis),
-            inputs=[processed_image1, processed_image2],
+            fn=run_pose_exploration_a,
+            inputs=[processed_image1, processed_image2, seed_value],
             outputs=[elevs, elev_ranges, vis_output]
         ).success(
             fn=partial(run_pose_exploration_b, cam_vis),
-            inputs=[processed_image1, processed_image2, elevs, elev_ranges, probe_bsz, adj_bsz, adj_iters],
+            inputs=[processed_image1, processed_image2, elevs, elev_ranges, probe_bsz, adj_bsz, adj_iters, seed_value],
             outputs=[anchor_polar, explored_sph, vis_output, refine_btn]
         )
         
         refine_btn.click(
             fn=partial(run_pose_refinement, cam_vis),
-            inputs=[processed_image1, processed_image2, anchor_polar, explored_sph, refine_iters],
+            inputs=[processed_image1, processed_image2, anchor_polar, explored_sph, refine_iters, seed_value],
             outputs=[refined_sph, vis_output]
         )