code pruning and change examples

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -1 +1,4 @@
-weights/
+weights/
+data/
+*.ipynb
+demo_examples_*

app.py

@@ -1,8 +1,3 @@
-'''
-conda activate zero123
-cd stable-diffusion
-python gradio_new.py 0
-'''
 import os, sys
 from huggingface_hub import snapshot_download
 
@@ -66,6 +61,71 @@ _DONE = "Done! Mesh is shown on the right. <br> If it is not satisfactory, pleas
 _REGEN_1 = "Selected view(s) are regenerated. You can click **Regenerate nearby views and mesh**. <br> Alternatively, if the regenerated view(s) are still not satisfactory, you can repeat the previous step (select the view and regenerate)."
 _REGEN_2 = "Regeneration done. Mesh is shown on the right."
 
+
+def calc_cam_cone_pts_3d(polar_deg, azimuth_deg, radius_m, fov_deg):
+    '''
+    :param polar_deg (float).
+    :param azimuth_deg (float).
+    :param radius_m (float).
+    :param fov_deg (float).
+    :return (5, 3) array of float with (x, y, z).
+    '''
+    polar_rad = np.deg2rad(polar_deg)
+    azimuth_rad = np.deg2rad(azimuth_deg)
+    fov_rad = np.deg2rad(fov_deg)
+    polar_rad = -polar_rad  # NOTE: Inverse of how used_x relates to x.
+
+    # Camera pose center:
+    cam_x = radius_m * np.cos(azimuth_rad) * np.cos(polar_rad)
+    cam_y = radius_m * np.sin(azimuth_rad) * np.cos(polar_rad)
+    cam_z = radius_m * np.sin(polar_rad)
+
+    # Obtain four corners of camera frustum, assuming it is looking at origin.
+    # First, obtain camera extrinsics (rotation matrix only):
+    camera_R = np.array([[np.cos(azimuth_rad) * np.cos(polar_rad),
+                          -np.sin(azimuth_rad),
+                          -np.cos(azimuth_rad) * np.sin(polar_rad)],
+                         [np.sin(azimuth_rad) * np.cos(polar_rad),
+                          np.cos(azimuth_rad),
+                          -np.sin(azimuth_rad) * np.sin(polar_rad)],
+                         [np.sin(polar_rad),
+                          0.0,
+                          np.cos(polar_rad)]])
+
+    # Multiply by corners in camera space to obtain go to space:
+    corn1 = [-1.0, np.tan(fov_rad / 2.0), np.tan(fov_rad / 2.0)]
+    corn2 = [-1.0, -np.tan(fov_rad / 2.0), np.tan(fov_rad / 2.0)]
+    corn3 = [-1.0, -np.tan(fov_rad / 2.0), -np.tan(fov_rad / 2.0)]
+    corn4 = [-1.0, np.tan(fov_rad / 2.0), -np.tan(fov_rad / 2.0)]
+    corn1 = np.dot(camera_R, corn1)
+    corn2 = np.dot(camera_R, corn2)
+    corn3 = np.dot(camera_R, corn3)
+    corn4 = np.dot(camera_R, corn4)
+
+    # Now attach as offset to actual 3D camera position:
+    corn1 = np.array(corn1) / np.linalg.norm(corn1, ord=2)
+    corn_x1 = cam_x + corn1[0]
+    corn_y1 = cam_y + corn1[1]
+    corn_z1 = cam_z + corn1[2]
+    corn2 = np.array(corn2) / np.linalg.norm(corn2, ord=2)
+    corn_x2 = cam_x + corn2[0]
+    corn_y2 = cam_y + corn2[1]
+    corn_z2 = cam_z + corn2[2]
+    corn3 = np.array(corn3) / np.linalg.norm(corn3, ord=2)
+    corn_x3 = cam_x + corn3[0]
+    corn_y3 = cam_y + corn3[1]
+    corn_z3 = cam_z + corn3[2]
+    corn4 = np.array(corn4) / np.linalg.norm(corn4, ord=2)
+    corn_x4 = cam_x + corn4[0]
+    corn_y4 = cam_y + corn4[1]
+    corn_z4 = cam_z + corn4[2]
+
+    xs = [cam_x, corn_x1, corn_x2, corn_x3, corn_x4]
+    ys = [cam_y, corn_y1, corn_y2, corn_y3, corn_y4]
+    zs = [cam_z, corn_z1, corn_z2, corn_z3, corn_z4]
+
+    return np.array([xs, ys, zs]).T
+
 class CameraVisualizer:
     def __init__(self, gradio_plot):
         self._gradio_plot = gradio_plot
@@ -77,18 +137,6 @@ class CameraVisualizer:
         self._8bit_image = None
         self._image_colorscale = None
 
-    def polar_change(self, value):
-        self._polar = value
-        # return self.update_figure()
-
-    def azimuth_change(self, value):
-        self._azimuth = value
-        # return self.update_figure()
-
-    def radius_change(self, value):
-        self._radius = value
-        # return self.update_figure()
-
     def encode_image(self, raw_image, elev=90):
         '''
         :param raw_image (H, W, 3) array of uint8 in [0, 255].
@@ -261,7 +309,7 @@ def stage1_run(models, device, cam_vis, tmp_dir,
         os.makedirs(stage1_dir, exist_ok=True)
         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)
         stage2_steps = 50 # ddim_steps
-        zero123_infer(models['turncam'], tmp_dir, indices=[0], device=device, ddim_steps=stage2_steps, scale=scale)
+        zero123_infer(model, tmp_dir, indices=[0], device=device, ddim_steps=stage2_steps, scale=scale)
         elev_output = estimate_elev(tmp_dir)
         gen_poses(tmp_dir, elev_output)
         show_in_im1 = np.asarray(input_im, dtype=np.uint8)
@@ -277,8 +325,6 @@ def stage1_run(models, device, cam_vis, tmp_dir,
         return (90-elev_output, new_fig, *output_ims, *output_ims_2)
     else:
         rerun_idx = [i for i in range(len(btn_retrys)) if btn_retrys[i]]
-        # elev_output = estimate_elev(tmp_dir)
-        # if elev_output > 75:
         if 90-int(elev["label"]) > 75:
             rerun_idx_in = [i if i < 4 else i+4 for i in rerun_idx]
         else:
@@ -297,7 +343,6 @@ def stage1_run(models, device, cam_vis, tmp_dir,
     
 def stage2_run(models, device, tmp_dir,
                elev, scale, rerun_all=[], stage2_steps=50):
-    # print("elev", elev)
     flag_lower_cam = 90-int(elev["label"]) <= 75
     is_rerun = True if rerun_all else False
     model = models['turncam'].half()
@@ -308,7 +353,7 @@ def stage2_run(models, device, tmp_dir,
             zero123_infer(model, tmp_dir, indices=list(range(1,4))+list(range(8,12)), device=device, ddim_steps=stage2_steps, scale=scale)
     else:
         print("rerun_idx", rerun_all)
-        zero123_infer(models['turncam'], tmp_dir, indices=rerun_all, device=device, ddim_steps=stage2_steps, scale=scale)
+        zero123_infer(model, tmp_dir, indices=rerun_all, device=device, ddim_steps=stage2_steps, scale=scale)
     
     dataset = tmp_dir
     main_dir_path = os.path.dirname(os.path.abspath(
@@ -321,7 +366,7 @@ def stage2_run(models, device, tmp_dir,
     os.system(bash_script)
     os.chdir(main_dir_path)
 
-    ply_path = os.path.join(tmp_dir, f"meshes_val_bg/lod0/mesh_00340000_gradio_lod0.ply")
+    ply_path = os.path.join(tmp_dir, f"meshes_val_bg/lod0/mesh_00215000_gradio_lod0.ply")
     mesh_path = os.path.join(tmp_dir, "mesh.obj")
     # Read the textured mesh from .ply file
     mesh = trimesh.load_mesh(ply_path)
@@ -382,77 +427,6 @@ def preprocess_run(predictor, models, raw_im, preprocess, *bbox_sliders):
     torch.cuda.empty_cache()
     return input_256
 
-def calc_cam_cone_pts_3d(polar_deg, azimuth_deg, radius_m, fov_deg):
-    '''
-    :param polar_deg (float).
-    :param azimuth_deg (float).
-    :param radius_m (float).
-    :param fov_deg (float).
-    :return (5, 3) array of float with (x, y, z).
-    '''
-    polar_rad = np.deg2rad(polar_deg)
-    azimuth_rad = np.deg2rad(azimuth_deg)
-    fov_rad = np.deg2rad(fov_deg)
-    polar_rad = -polar_rad  # NOTE: Inverse of how used_x relates to x.
-
-    # Camera pose center:
-    cam_x = radius_m * np.cos(azimuth_rad) * np.cos(polar_rad)
-    cam_y = radius_m * np.sin(azimuth_rad) * np.cos(polar_rad)
-    cam_z = radius_m * np.sin(polar_rad)
-
-    # Obtain four corners of camera frustum, assuming it is looking at origin.
-    # First, obtain camera extrinsics (rotation matrix only):
-    camera_R = np.array([[np.cos(azimuth_rad) * np.cos(polar_rad),
-                          -np.sin(azimuth_rad),
-                          -np.cos(azimuth_rad) * np.sin(polar_rad)],
-                         [np.sin(azimuth_rad) * np.cos(polar_rad),
-                          np.cos(azimuth_rad),
-                          -np.sin(azimuth_rad) * np.sin(polar_rad)],
-                         [np.sin(polar_rad),
-                          0.0,
-                          np.cos(polar_rad)]])
-    # print('camera_R:', lo(camera_R).v)
-
-    # Multiply by corners in camera space to obtain go to space:
-    corn1 = [-1.0, np.tan(fov_rad / 2.0), np.tan(fov_rad / 2.0)]
-    corn2 = [-1.0, -np.tan(fov_rad / 2.0), np.tan(fov_rad / 2.0)]
-    corn3 = [-1.0, -np.tan(fov_rad / 2.0), -np.tan(fov_rad / 2.0)]
-    corn4 = [-1.0, np.tan(fov_rad / 2.0), -np.tan(fov_rad / 2.0)]
-    corn1 = np.dot(camera_R, corn1)
-    corn2 = np.dot(camera_R, corn2)
-    corn3 = np.dot(camera_R, corn3)
-    corn4 = np.dot(camera_R, corn4)
-
-    # Now attach as offset to actual 3D camera position:
-    corn1 = np.array(corn1) / np.linalg.norm(corn1, ord=2)
-    corn_x1 = cam_x + corn1[0]
-    corn_y1 = cam_y + corn1[1]
-    corn_z1 = cam_z + corn1[2]
-    corn2 = np.array(corn2) / np.linalg.norm(corn2, ord=2)
-    corn_x2 = cam_x + corn2[0]
-    corn_y2 = cam_y + corn2[1]
-    corn_z2 = cam_z + corn2[2]
-    corn3 = np.array(corn3) / np.linalg.norm(corn3, ord=2)
-    corn_x3 = cam_x + corn3[0]
-    corn_y3 = cam_y + corn3[1]
-    corn_z3 = cam_z + corn3[2]
-    corn4 = np.array(corn4) / np.linalg.norm(corn4, ord=2)
-    corn_x4 = cam_x + corn4[0]
-    corn_y4 = cam_y + corn4[1]
-    corn_z4 = cam_z + corn4[2]
-
-    xs = [cam_x, corn_x1, corn_x2, corn_x3, corn_x4]
-    ys = [cam_y, corn_y1, corn_y2, corn_y3, corn_y4]
-    zs = [cam_z, corn_z1, corn_z2, corn_z3, corn_z4]
-
-    return np.array([xs, ys, zs]).T
-
-def save_bbox(dir, x_min, y_min, x_max, y_max):
-    box = np.array([x_min, y_min, x_max, y_max])
-    # save the box to a file
-    bbox_path = os.path.join(dir, "bbox.txt")
-    np.savetxt(bbox_path, box)
-
 def on_coords_slider(image, x_min, y_min, x_max, y_max, color=(88, 191, 131, 255)):
     """Draw a bounding box annotation for an image."""
     print("on_coords_slider, drawing bbox...")
@@ -464,12 +438,11 @@ def on_coords_slider(image, x_min, y_min, x_max, y_max, color=(88, 191, 131, 255
         y_min = int(y_min * shrink_ratio)
         x_max = int(x_max * shrink_ratio)
         y_max = int(y_max * shrink_ratio)
-    print("on_coords_slider, image_size:", np.array(image).shape)
     image = cv2.cvtColor(np.array(image), cv2.COLOR_RGBA2BGRA)
     image = cv2.rectangle(image, (x_min, y_min), (x_max, y_max), color, int(max(max(image.shape) / 400*2, 2)))
     return cv2.cvtColor(image, cv2.COLOR_BGRA2RGBA) # image[:, :, ::-1]
 
-def save_img(image):
+def init_bbox(image):
     image.thumbnail([512, 512], Image.Resampling.LANCZOS)
     width, height = image.size
     image_rem = image.convert('RGBA')
@@ -500,12 +473,6 @@ def run_demo(
         device_idx=_GPU_INDEX,
         ckpt='zero123-xl.ckpt'):
 
-    print('sys.argv:', sys.argv)
-    if len(sys.argv) > 1:
-        print('old device_idx:', device_idx)
-        device_idx = int(sys.argv[1])
-        print('new device_idx:', device_idx)
-
     device = f"cuda:{device_idx}" if torch.cuda.is_available() else "cpu"
     models = init_model(device, os.path.join(code_dir, ckpt))
     # model = models['turncam']
@@ -523,9 +490,8 @@ def run_demo(
     example_fns.sort()
     examples_full = [os.path.join(example_folder, x) for x in example_fns if x.endswith('.png')]
 
-
     # Compose demo layout & data flow.
-    css = "#model-3d-out {height: 400px;} #plot-out {height: 425px;}"
+    css = "#model-3d-out {height: 400px;} #plot-out {height: 450px;}"
     with gr.Blocks(title=_TITLE, css=css) as demo:
         gr.Markdown('# ' + _TITLE)
         gr.Markdown(_DESCRIPTION)
@@ -625,7 +591,7 @@ def run_demo(
                            ).success(disable_func, inputs=run_btn, outputs=run_btn
                            ).success(fn=tmp_func, inputs=[image_block], outputs=[placeholder]
                            ).success(fn=partial(update_guide, _BBOX_1), outputs=[guide_text]
-                           ).success(fn=save_img,
+                           ).success(fn=init_bbox,
                                      inputs=[image_block],
                                      outputs=[bbox_block, *bbox_sliders]
                            ).success(fn=partial(update_guide, _BBOX_3), outputs=[guide_text]