import os # import time # import shutil from pathlib import Path # from typing import Union # import atexit import spaces # from concurrent.futures import ThreadPoolExecutor import torch import open3d as o3d import trimesh import gradio as gr from gradio_imageslider import ImageSlider import cv2 import numpy as np import imageio # from promptda.promptda import PromptDA # from promptda.utils.io_wrapper import load_image, load_depth # from promptda.utils.depth_utils import visualize_depth, unproject_depth # import torch DEVICE = 'cuda' # if torch.cuda.is_available( # ) else 'mps' if torch.backends.mps.is_available() else 'cpu' # model = PromptDA.from_pretrained('depth-anything/promptda_vitl').to(DEVICE).eval() # model = PromptDA.from_pretrained('depth-anything/promptda_vitl').eval() # thread_pool_executor = ThreadPoolExecutor(max_workers=1) # def delete_later(path: Union[str, os.PathLike], delay: int = 300): # print(f"Deleting file: {path}") # def _delete(): # try: # if os.path.isfile(path): # os.remove(path) # print(f"Deleted file: {path}") # elif os.path.isdir(path): # shutil.rmtree(path) # print(f"Deleted directory: {path}") # except: # pass # def _wait_and_delete(): # time.sleep(delay) # _delete(path) # thread_pool_executor.submit(_wait_and_delete) # atexit.register(_delete) # @spaces.GPU def run_with_gpu(image, prompt_depth): image = image.to(DEVICE) prompt_depth = prompt_depth.to(DEVICE) model.to(DEVICE) depth = model.predict(image, prompt_depth) depth = depth[0, 0].detach().cpu().numpy() return depth def check_is_stray_scanner_app_capture(input_dir): assert os.path.exists(os.path.join(input_dir, 'rgb.mp4')), 'rgb.mp4 not found' pass @spaces.GPU def run(input_file, resolution): # unzip zip file input_file = input_file.name root_dir = os.path.dirname(input_file) scene_name = input_file.split('/')[-1].split('.')[0] input_dir = os.path.join(root_dir, scene_name) cmd = f'unzip -o {input_file} -d {root_dir}' os.system(cmd) check_is_stray_scanner_app_capture(input_dir) # extract rgb images os.makedirs(os.path.join(input_dir, 'rgb'), exist_ok=True) cmd = f'ffmpeg -i {input_dir}/rgb.mp4 -start_number 0 -frames:v 10 -q:v 2 {input_dir}/rgb/%06d.jpg' os.system(cmd) # Loading & Inference image_path = os.path.join(input_dir, 'rgb', '000000.jpg') image = load_image(image_path) prompt_depth_path = os.path.join(input_dir, 'depth/000000.png') prompt_depth = load_depth(prompt_depth_path) depth = run_with_gpu(image, prompt_depth) color = (image[0].permute(1,2,0).cpu().numpy() * 255.).astype(np.uint8) # Visualization file vis_depth, depth_min, depth_max = visualize_depth(depth, ret_minmax=True) vis_prompt_depth = visualize_depth(prompt_depth[0, 0].detach().cpu().numpy(), depth_min=depth_min, depth_max=depth_max) vis_prompt_depth = cv2.resize(vis_prompt_depth, (vis_depth.shape[1], vis_depth.shape[0]), interpolation=cv2.INTER_NEAREST) # PLY File ixt_path = os.path.join(input_dir, f'camera_matrix.csv') ixt = np.loadtxt(ixt_path, delimiter=',') orig_max = 1920 now_max = max(color.shape[1], color.shape[0]) scale = orig_max / now_max ixt[:2] = ixt[:2] / scale pcd = unproject_depth(depth, ixt=ixt, color=color, ret_pcd=True) ply_path = os.path.join(input_dir, f'pointcloud.ply') o3d.io.write_point_cloud(ply_path, pcd) glb_path = os.path.join(input_dir, f'pointcloud.glb') scene_3d = trimesh.Scene() glb_colors = np.asarray(pcd.colors).astype(np.float32) glb_colors = np.concatenate([glb_colors, np.ones_like(glb_colors[:, :1])], axis=1) # glb_colors = (np.asarray(pcd.colors) * 255).astype(np.uint8) pcd_data = trimesh.PointCloud( vertices=np.asarray(pcd.points) * np.array([[1, -1, -1]]), colors=glb_colors.astype(np.float64), ) scene_3d.add_geometry(pcd_data) scene_3d.export(file_obj=glb_path) # o3d.io.write_point_cloud(glb_path, pcd) # Depth Map Original Value depth_path = os.path.join(input_dir, f'depth.png') output_depth = (depth * 1000).astype(np.uint16) imageio.imwrite(depth_path, output_depth) delete_later(Path(input_dir)) delete_later(Path(input_file)) return color, (vis_depth, vis_prompt_depth), Path(glb_path), Path(ply_path).as_posix(), Path(depth_path).as_posix() DESCRIPTION = """ # Estimate accurate and high-resolution depth maps from your iPhone capture. Project Page: [Prompt Depth Anything](https://promptda.github.io/) ## Requirements: 1. iPhone 12 Pro or later Pro models, iPad 2020 Pro or later Pro models 2. Free iOS App: [Stray Scanner App](https://apps.apple.com/us/app/stray-scanner/id1557051662) ## Testing Steps: 1. Capture a scene with the Stray Scanner App. 2. Use the iPhone [Files App](https://apps.apple.com/us/app/files/id1232058109) to compress it into a zip file and transfer it to your computer. (Long press the capture folder to compress) 3. Upload the zip file and click "Submit" to get the depth map of the first frame. Note: - Currently, this demo only supports inference for the first frame. If you need to obtain all depth frames, please refer to our [GitHub repo](https://github.com/DepthAnything/PromptDA). - The depth map is stored as uint16, with a unit of millimeters. """ # @click.command() # @click.option('--share', is_flag=True, help='Whether to run the app in shared mode.') def main(): with gr.Blocks(theme=gr.themes.Soft()) as demo: gr.Markdown(DESCRIPTION) with gr.Row(): input_file = gr.File(type="filepath", label="Stray scanner app capture zip file") resolution = gr.Dropdown(choices=['756x1008', '1428x1904'], value='756x1008', label="Inference resolution") submit_btn = gr.Button("Submit") # gr.Examples(examples=[ # ["data/assets/example0_chair.zip", "756x1008"] # ], # inputs=[input_file, resolution], # label="Examples", # ) with gr.Row(): with gr.Column(): output_rgb = gr.Image(type="numpy", label="RGB Image") with gr.Column(): output_depths = ImageSlider(label="Output depth / prompt depth", position=0.5) with gr.Row(): with gr.Column(): output_3d_model = gr.Model3D(label="3D Viewer", display_mode='solid', clear_color=[1.0, 1.0, 1.0, 1.0]) with gr.Column(): output_ply = gr.File(type="filepath", label="Download the unprojected point cloud as .ply file", height=30) output_depth_map = gr.File(type="filepath", label="Download the depth map as .png file", height=30) outputs = [ output_rgb, output_depths, output_3d_model, output_ply, output_depth_map, ] gr.Examples(examples=[ ["data/assets/example0_chair.zip", "756x1008"] ], fn=run, inputs=[input_file, resolution], outputs=outputs, label="Examples", cache_examples=True, ) submit_btn.click(run, inputs=[input_file, resolution], outputs=outputs) demo.launch() # def main(): # gr.Interface( # fn=run, # inputs=[ # gr.File(type="filepath", label="Stray scanner app capture zip file"), # gr.Dropdown(choices=['756x1008', '1428x1904'], value='756x1008', label="Inference resolution") # ], # outputs=[ # gr.Image(type="numpy", label="RGB Image"), # ImageSlider(label="Depth map / prompt depth", position=0.5), # gr.Model3D(label="3D Viewer", display_mode='solid', clear_color=[1.0, 1.0, 1.0, 1.0]), # gr.File(type="filepath", label="Download the unprojected point cloud as .ply file"), # gr.File(type="filepath", label="Download the depth map as .png file"), # ], # title=None, # description=DESCRIPTION, # clear_btn=None, # allow_flagging="never", # theme=gr.themes.Soft(), # examples=[ # ["data/assets/example0_chair.zip"] # ] # ).launch() main()