Support API calls

app.py

@@ -12,17 +12,13 @@ elev_est_dir = os.path.abspath(os.path.join(code_dir, "one2345_elev_est"))
 sys.path.append(elev_est_dir)
 
 if not is_local_run:
-    # import pip
-    # pip.main(['install', elev_est_dir])
     # export TORCH_CUDA_ARCH_LIST="7.0;7.2;8.0;8.6"
     # export IABN_FORCE_CUDA=1
     os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6"
     os.environ["IABN_FORCE_CUDA"] = "1"
     os.environ["FORCE_CUDA"] = "1"
-    # pip.main(["install", "inplace_abn"]) 
     subprocess.run(['pip', 'install', 'inplace_abn'])
     # FORCE_CUDA=1 pip install --no-cache-dir git+https://github.com/mit-han-lab/[email protected]
-    # pip.main(["install", "--no-cache-dir", "git+https://github.com/mit-han-lab/[email protected]"])
     subprocess.run(['pip', 'install', '--no-cache-dir', 'git+https://github.com/mit-han-lab/[email protected]'])
 
 import shutil
@@ -154,12 +150,9 @@ class CameraVisualizer:
 
         self._raw_image = raw_image
         self._8bit_image = Image.fromarray(raw_image).convert('P', palette='WEB', dither=None)
-        # self._8bit_image = Image.fromarray(raw_image.clip(0, 254)).convert(
-        #     'P', palette='WEB', dither=None)
         self._image_colorscale = [
             [i / 255.0, 'rgb({}, {}, {})'.format(*rgb)] for i, rgb in enumerate(idx_to_color)]
         self._elev = elev
-        # return self.update_figure()
 
     def update_figure(self):
         fig = go.Figure()
@@ -243,9 +236,6 @@ class CameraVisualizer:
 
             # look at center of scene
             fig.update_layout(
-                # width=640,
-                # height=480,
-                # height=400,
                 height=450,
                 autosize=True,
                 hovermode=False,
@@ -312,7 +302,7 @@ def stage1_run(models, device, cam_vis, tmp_dir,
         stage2_steps = 50 # ddim_steps
         zero123_infer(model, tmp_dir, indices=[0], device=device, ddim_steps=stage2_steps, scale=scale)
         try:
-            elev_output = estimate_elev(tmp_dir)
+            elev_output = int(estimate_elev(tmp_dir))
         except:
             print("Failed to estimate polar angle")
             elev_output = 90
@@ -459,6 +449,119 @@ def init_bbox(image):
             gr.update(value=x_max, maximum=width),
             gr.update(value=y_max, maximum=height)]
 
+### API functions
+def preprocess_api(predictor, raw_im):
+    raw_im.thumbnail([512, 512], Image.Resampling.LANCZOS)
+    image_rem = raw_im.convert('RGBA')
+    image_nobg = remove(image_rem, alpha_matting=True)
+    arr = np.asarray(image_nobg)[:,:,-1]
+    x_nonzero = np.nonzero(arr.sum(axis=0))
+    y_nonzero = np.nonzero(arr.sum(axis=1))
+    x_min = int(x_nonzero[0].min())
+    y_min = int(y_nonzero[0].min())
+    x_max = int(x_nonzero[0].max())
+    y_max = int(y_nonzero[0].max())
+    image_sam = sam_out_nosave(predictor, raw_im.convert("RGB"), x_min, y_min, x_max, y_max)
+    input_256 = image_preprocess_nosave(image_sam, lower_contrast=False, rescale=True)
+    torch.cuda.empty_cache()
+    return input_256
+
+def estimate_elev_api(models, device, predictor,
+               input_im, preprocess=True, scale=3, ddim_steps=50):
+    model = models['turncam'].half()
+    tmp_dir = tempfile.TemporaryDirectory(dir=os.path.join(os.path.dirname(__file__), 'demo_tmp')).name
+    stage1_dir = os.path.join(tmp_dir, "stage1_8")
+    os.makedirs(stage1_dir, exist_ok=True)
+    if preprocess:
+        input_im = preprocess_api(predictor, input_im)
+    input_image = input_im.convert("RGB")
+    output_ims = predict_stage1_gradio(model, input_image, save_path=stage1_dir, adjust_set=[0], device=device, ddim_steps=ddim_steps, scale=scale)
+    stage2_steps = 50 # ddim_steps
+    zero123_infer(model, tmp_dir, indices=[0], device=device, ddim_steps=stage2_steps, scale=scale)
+    try:
+        polar_angle = int(estimate_elev(tmp_dir))
+    except:
+        print("Failed to estimate polar angle")
+        polar_angle = 90
+    print("Estimated polar angle:", polar_angle)
+    return 90-polar_angle
+
+def convert_mesh_format(exp_dir, output_format=".obj"):
+    ply_path = os.path.join(exp_dir, f"meshes_val_bg/lod0/mesh_00215000_gradio_lod0.ply")
+    mesh_path = os.path.join(exp_dir, f"mesh{output_format}")
+    mesh = trimesh.load_mesh(ply_path)
+    rotation_matrix = trimesh.transformations.rotation_matrix(np.pi/2, [1, 0, 0]) @ trimesh.transformations.rotation_matrix(np.pi, [0, 0, 1])
+    mesh.apply_transform(rotation_matrix)
+    mesh.vertices[:, 0] = -mesh.vertices[:, 0]
+    mesh.faces = np.fliplr(mesh.faces)
+    if output_format == ".obj":
+        # Export the mesh as .obj file with colors
+        mesh.export(mesh_path, file_type='obj', include_color=True)
+    else:
+        mesh.export(mesh_path, file_type='glb')
+    return mesh_path
+    
+def reconstruct(exp_dir, output_format=".ply", device_idx=0):
+
+    main_dir_path = os.path.dirname(__file__)
+    torch.cuda.empty_cache()
+    os.chdir(os.path.join(code_dir, 'SparseNeuS_demo_v1/'))
+
+    bash_script = f'CUDA_VISIBLE_DEVICES={device_idx} python exp_runner_generic_blender_val.py \
+                    --specific_dataset_name {exp_dir} \
+                    --mode export_mesh \
+                    --conf confs/one2345_lod0_val_demo.conf'
+    print(bash_script)
+    os.system(bash_script)
+    os.chdir(main_dir_path)
+
+    ply_path = os.path.join(exp_dir, f"meshes_val_bg/lod0/mesh_00215000_gradio_lod0.ply")
+    if output_format == ".ply":
+        return ply_path
+    if output_format not in [".obj", ".glb"]:
+        print("Invalid output format, must be one of .ply, .obj, .glb")
+        return ply_path
+    return convert_mesh_format(exp_dir, output_format=output_format)
+
+def gen_mesh_api(models, predictor, device,
+               input_im, preprocess=True, scale=3, ddim_steps=75, stage2_steps=50):
+    if preprocess:
+        input_im = preprocess_api(predictor, input_im)
+    model = models['turncam'].half()
+    # folder to save the stage 1 images
+    exp_dir = tempfile.TemporaryDirectory(dir=os.path.join(os.path.dirname(__file__), 'demo_tmp')).name
+    stage1_dir = os.path.join(exp_dir, "stage1_8")
+    os.makedirs(stage1_dir, exist_ok=True)
+
+    # stage 1: generate 4 views at the same elevation as the input
+    output_ims = predict_stage1_gradio(model, input_im, save_path=stage1_dir, adjust_set=list(range(4)), device=device, ddim_steps=ddim_steps, scale=scale)
+    
+    # stage 2 for the first image
+    # infer 4 nearby views for an image to estimate the polar angle of the input
+    stage2_steps = 50 # ddim_steps
+    zero123_infer(model, exp_dir, indices=[0], device=device, ddim_steps=stage2_steps, scale=scale)
+    # estimate the camera pose (elevation) of the input image.
+    try:
+        polar_angle = int(estimate_elev(exp_dir))
+    except:
+        print("Failed to estimate polar angle")
+        polar_angle = 90
+    print("Estimated polar angle:", polar_angle)
+    gen_poses(exp_dir, polar_angle)
+
+    # stage 1: generate another 4 views at a different elevation
+    if polar_angle <= 75:
+        output_ims_2 = predict_stage1_gradio(model, input_im, save_path=stage1_dir, adjust_set=list(range(4,8)), device=device, ddim_steps=ddim_steps, scale=scale)
+    else:
+        output_ims_2 = predict_stage1_gradio(model, input_im, save_path=stage1_dir, adjust_set=list(range(8,12)), device=device, ddim_steps=ddim_steps, scale=scale)
+    torch.cuda.empty_cache()
+    # stage 2 for the remaining 7 images, generate 7*4=28 views
+    if polar_angle <= 75:
+        zero123_infer(model, exp_dir, indices=list(range(1,8)), device=device, ddim_steps=stage2_steps, scale=scale)
+    else:
+        zero123_infer(model, exp_dir, indices=list(range(1,4))+list(range(8,12)), device=device, ddim_steps=stage2_steps, scale=scale)
+    return reconstruct(exp_dir)
+
 
 def run_demo(
         device_idx=_GPU_INDEX,
@@ -473,7 +576,6 @@ def run_demo(
     with open('instructions_12345.md', 'r') as f:
         article = f.read()
 
-    # NOTE: Examples must match inputs
     example_folder = os.path.join(os.path.dirname(__file__), 'demo_examples')
     example_fns = os.listdir(example_folder)
     example_fns.sort()
@@ -494,7 +596,7 @@ def run_demo(
                 image_block = gr.Image(type='pil', image_mode='RGBA', height=290, label='Input image', tool=None)
 
                 gr.Examples(
-                    examples=examples_full,  # NOTE: elements must match inputs list!
+                    examples=examples_full,
                     inputs=[image_block],
                     outputs=[image_block],
                     cache_examples=False,
@@ -569,6 +671,31 @@ def run_demo(
             </div>
         """)
 
+        # hidden buttons for supporting API calls
+        elev_est_btn = gr.Button('Run API', variant='primary', visible=False)
+        elev_est_out = gr.Number(value=0, visible=False)
+        elev_preprocess_chk = gr.Checkbox(value=True, visible=False)
+
+        elev_est_btn.click(fn=partial(estimate_elev_api, models, device, predictor), 
+                           inputs=[image_block, elev_preprocess_chk], 
+                           outputs=[elev_est_out], 
+                           api_name='estimate_elevation',
+                           queue=True)
+        
+        preprocess_btn = gr.Button('Run API', variant='primary', visible=False)
+        preprocess_btn.click(fn=partial(preprocess_api, predictor), 
+                             inputs=[image_block], 
+                             outputs=[sam_block], 
+                             api_name='preprocess',
+                             queue=True)
+        
+        gen_mesh_btn = gr.Button('Run API', variant='primary', visible=False)
+        gen_mesh_btn.click(fn=partial(gen_mesh_api, models, predictor, device),
+                            inputs=[image_block, elev_preprocess_chk], 
+                            outputs=[mesh_output], 
+                            api_name='generate_mesh',
+                            queue=True)
+
         update_guide = lambda GUIDE_TEXT: gr.update(value=GUIDE_TEXT)
 
         views = [view_1, view_2, view_3, view_4, view_5, view_6, view_7, view_8]

instructions_12345.md

@@ -1,3 +1,35 @@
+## APIs:
+<details>
+<summary>We offer handy APIs for our pipeline and its components.</summary>
+
+```python
+from gradio_client import Client
+client = Client("https://one-2-3-45-one-2-3-45.hf.space/")
+input_img_path = "https://huggingface.co/spaces/One-2-3-45/One-2-3-45/resolve/main/demo_examples/01_wild_hydrant.png"
+
+### Single image to 3D mesh
+generated_mesh_filepath = client.predict(
+	input_img_path,	
+	True,		# image preprocessing
+	api_name="/generate_mesh"
+)
+
+### Elevation estimation 
+# DON'T TO ASK USERS TO ESTIMATE ELEVATION! This OFF-THE-SHELF algorithm is ALL YOU NEED!
+elevation_angle_deg = client.predict(
+	input_img_path,
+	True,		# image preprocessing
+	api_name="/estimate_elevation"
+)
+
+### Image preprocessing: segment, rescale, and recenter
+segmented_img_filepath = client.predict(
+	input_img_path,	
+	api_name="/preprocess"
+)
+```
+</details>
+
 ## Tuning Tips:
 
 1. The multi-view prediction module (Zero123) operates probabilistically. If some of the predicted views are not satisfactory, you may select and regenerate them.