ZiyuG/SAM2Point
16
1import spaces2from pickle import FALSE3import gradio as gr4import numpy as np5import plotly.graph_objects as go6from sam2point import dataset7import sam2point.configs as configs8from demo_utils import run_demo, create_box9 10samples = {11 "3D Indoor Scene - S3DIS": ["Conference Room", "Restroom", "Lobby", "Office1", "Office2"],12 "3D Indoor Scene - ScanNet": ["Scene1", "Scene2", "Scene3", "Scene4", "Scene5", "Scene6"],13 "3D Raw LiDAR - KITTI": ["Scene1", "Scene2", "Scene3", "Scene4", "Scene5", "Scene6"],14 "3D Outdoor Scene - Semantic3D": ["Scene1", "Scene2", "Scene3", "Scene4", "Scene5", "Scene6", "Scene7"],15 "3D Object - Objaverse": ["Plant", "Lego", "Lock", "Eleplant", "Knife Rest", "Skateboard", "Popcorn Machine", "Stove", "Bus Shelter", "Thor Hammer", "Horse"],16}17 18PATH = {19 "S3DIS": ['Area_1_conferenceRoom_1.txt', 'Area_2_WC_1.txt', 'Area_4_lobby_2.txt', 'Area_5_office_3.txt', 'Area_6_office_9.txt'],20 "ScanNet": ['scene0005_01.pth', 'scene0010_01.pth', 'scene0016_02.pth', 'scene0019_01.pth', 'scene0000_00.pth', 'scene0002_00.pth'],21 "Objaverse": ["plant.npy", "human.npy", "lock.npy", "elephant.npy", "knife_rest.npy", "skateboard.npy", "popcorn_machine.npy", "stove.npy", "bus_shelter.npy", "thor_hammer.npy", "horse.npy"],22 "KITTI": ["scene1.npy", "scene2.npy", "scene3.npy", "scene4.npy", "scene5.npy", "scene6.npy"],23 "Semantic3D": ["scene1.npy", "scene2.npy", "patch19.npy", "patch0.npy", "patch1.npy", "patch50.npy", "patch62.npy"]24}25 26prompt_types = ["Point", "Box", "Mask"]27 28def load_3d_scene(name, sample_idx=-1, type_=None, prompt=None, final=False, new_color=None):29 DATASET = name.split('-')[1].replace(" ", "")30 path = 'data/' + DATASET + '/' + PATH[DATASET][sample_idx]31 asp, SIZE = 1., 132 33 print(path)34 if DATASET == 'S3DIS':35 point, color = dataset.load_S3DIS_sample(path, sample=True)36 alpha = 137 elif DATASET == 'ScanNet':38 point, color = dataset.load_ScanNet_sample(path)39 alpha = 140 elif DATASET == 'Objaverse':41 point, color = dataset.load_Objaverse_sample(path)42 alpha = 143 SIZE = 244 elif DATASET == 'KITTI':45 point, color = dataset.load_KITTI_sample(path)46 asp = 0.347 alpha = 0.748 elif DATASET == 'Semantic3D':49 point, color = dataset.load_Semantic3D_sample(path, sample_idx, sample=True)50 alpha = 0.251 print("Loading Dataset:", DATASET, "Point Cloud Size:", point.shape, "Path:", path)52 53 ##### Initial Show #####54 if not type_:55 if point.shape[0] > 100000: # sample points for speeding up56 indices = np.random.choice(point.shape[0], 100000, replace=False)57 point = point[indices]58 color = color[indices]59 fig = go.Figure(60 data=[61 go.Scatter3d(62 x=point[:,0], y=point[:,1], z=point[:,2],63 mode='markers',64 marker=dict(size=SIZE, color=color, opacity=alpha),65 name=""66 )67 ],68 layout=dict(69 scene=dict(70 xaxis=dict(visible=False),71 yaxis=dict(visible=False),72 zaxis=dict(visible=False),73 aspectratio=dict(x=1, y=1, z=asp), 74 camera=dict(eye=dict(x=1.5, y=1.5, z=1.5))75 )76 )77 )78 return fig79 ##### Final Results #####80 if final:81 color = new_color82 green = np.array([[0.1, 0.1, 0.1]])83 add_green = go.Scatter3d(84 x=green[:,0], y=green[:,1], z=green[:,2],85 mode='markers',86 marker=dict(size=0.0001, color='green', opacity=1),87 name="Segmentation Results"88 )89 if type_ == "box": 90 if point.shape[0] > 100000:91 indices = np.random.choice(point.shape[0], 100000, replace=False)92 point = point[indices]93 color = color[indices]94 scatter = go.Scatter3d(95 x=point[:,0], y=point[:,1], z=point[:,2],96 mode='markers',97 marker=dict(size=SIZE, color=color, opacity=alpha),98 name="3D Object/Scene"99 )100 if final: scatter = [scatter, add_green] + create_box(prompt)101 else: scatter = [scatter] + create_box(prompt)102 elif type_ == "point":103 prompt = np.array([prompt])104 new = go.Scatter3d(105 x=prompt[:,0], y=prompt[:,1], z=prompt[:,2],106 mode='markers',107 marker=dict(size=5, color='red', opacity=1),108 name="Point Prompt"109 )110 if point.shape[0] > 100000:111 indices = np.random.choice(point.shape[0], 100000, replace=False)112 point = point[indices]113 color = color[indices]114 scatter = go.Scatter3d(115 x=point[:,0], y=point[:,1], z=point[:,2],116 mode='markers',117 marker=dict(size=SIZE, color=color, opacity=alpha),118 name="3D Object/Scene"119 )120 if final: scatter = [scatter, new, add_green]121 else: scatter = [scatter, new]122 elif type_ == 'mask' and not final:123 color = np.clip(prompt * 255, 0, 255).astype(np.uint8)124 if point.shape[0] > 100000:125 indices = np.random.choice(point.shape[0], 100000, replace=False)126 point = point[indices]127 color = color[indices]128 scatter = go.Scatter3d(129 x=point[:,0], y=point[:,1], z=point[:,2],130 mode='markers',131 marker=dict(size=SIZE, color=color, opacity=alpha),132 name="3D Object/Scene"133 )134 red = np.array([[0.1, 0.1, 0.1]])135 add_red = go.Scatter3d(136 x=red[:,0], y=red[:,1], z=red[:,2],137 mode='markers',138 marker=dict(size=0.0001, color='red', opacity=1),139 name="Mask Prompt"140 )141 scatter = [scatter, add_red]142 elif type_ == 'mask' and final:143 if point.shape[0] > 100000:144 indices = np.random.choice(point.shape[0], 100000, replace=False)145 point = point[indices]146 color = color[indices]147 scatter = go.Scatter3d(148 x=point[:,0], y=point[:,1], z=point[:,2],149 mode='markers',150 marker=dict(size=SIZE, color=color, opacity=alpha),151 name="3D Object/Scene"152 )153 scatter = [scatter, add_green]154 else: 155 print("Wrong Prompt Type")156 exit(1)157 158 fig = go.Figure(159 data=scatter,160 layout=dict(161 scene=dict(162 xaxis=dict(visible=False),163 yaxis=dict(visible=False),164 zaxis=dict(visible=False),165 aspectratio=dict(x=1, y=1, z=asp), 166 camera=dict(eye=dict(x=1.5, y=1.5, z=1.5))167 )168 )169 )170 return fig171 172@spaces.GPU()173def show_prompt_in_3d(name, sample_idx, prompt_type, prompt_idx):174 if name == None or sample_idx == None or prompt_type == None or prompt_idx == None:175 return gr.Plot(), gr.Textbox(label="Response", value="Please ensure all options are selected.", visible=True)176 177 DATASET = name.split('-')[1].replace(" ", "")178 TYPE = prompt_type.lower()179 theta = 0. if DATASET in "S3DIS ScanNet" else 0.5180 mode = "bilinear" if DATASET in "S3DIS ScanNet" else 'nearest'181 182 prompt = run_demo(DATASET, TYPE, sample_idx, prompt_idx, 0.02, theta, mode, ret_prompt=True)183 fig = load_3d_scene(name, sample_idx, TYPE, prompt)184 return fig, gr.Textbox(label="Response", value="Prompt has been shown in 3D Object/Scene!", visible=True)185 186@spaces.GPU()187def start_segmentation(name=None, sample_idx=None, prompt_type=None, prompt_idx=None, vx=0.02):188 if name == None or sample_idx == None or prompt_type == None or prompt_idx == None:189 return gr.Plot(), gr.Textbox(label="Response", value="Please ensure all options are selected.", visible=True)190 191 DATASET = name.split('-')[1].replace(" ", "")192 TYPE = prompt_type.lower()193 theta = 0. if DATASET in "S3DIS ScanNet" else 0.5194 mode = "bilinear" if DATASET in "S3DIS ScanNet" else 'nearest'195 196 new_color, prompt = run_demo(DATASET, TYPE, sample_idx, prompt_idx, vx, theta, mode, ret_prompt=False)197 fig = load_3d_scene(name, sample_idx, TYPE, prompt, final=True, new_color=new_color)198 return fig, gr.Textbox(label="Response", value="Segmentation completed successfully!", visible=True)199 200def update1(datasets):201 if 'Objaverse' in datasets:202 return gr.Radio(label="Select 3D Object", choices=samples[datasets]), gr.Textbox(label="Response", value="", visible=True) 203 return gr.Radio(label="Select 3D Scene", choices=samples[datasets]), gr.Textbox(label="Response", value="", visible=True) 204 205def update2(name, sample_idx, prompt_type):206 if name == None or sample_idx == None or prompt_type == None:207 return gr.Radio(label="Select Prompt Example", choices=[]), gr.Textbox(label="Response", value="", visible=True) 208 DATASET = name.split('-')[1].replace(" ", "")209 TYPE = prompt_type.lower() + '_prompts'210 211 if DATASET == 'S3DIS': 212 info = configs.S3DIS_samples[sample_idx][TYPE]213 elif DATASET == 'ScanNet': 214 info = configs.ScanNet_samples[sample_idx][TYPE]215 elif DATASET == 'Objaverse': 216 info = configs.Objaverse_samples[sample_idx][TYPE]217 elif DATASET == 'KITTI': 218 info = configs.KITTI_samples[sample_idx][TYPE]219 elif DATASET == 'Semantic3D': 220 info = configs.Semantic3D_samples[sample_idx][TYPE]221 222 cur = ['Example ' + str(i) for i in range(1, len(info) + 1)]223 return gr.Radio(label="Select Prompt Example", choices=cur), gr.Textbox(label="Response", value="", visible=True) 224 225def update3(name, sample_idx, prompt_type, prompt_idx):226 if name == None or sample_idx == None or prompt_type == None:227 return gr.Textbox(label="Response", value="", visible=True), gr.Slider(minimum=0.01, maximum=0.15, step=0.001, label="Voxel Size", value=0.02)228 DATASET = name.split('-')[1].replace(" ", "")229 TYPE = configs.VOXEL[prompt_type.lower()]230 231 if DATASET in "S3DIS ScanNet": 232 vx_ = 0.02233 elif DATASET == 'Objaverse': 234 vx_ = configs.Objaverse_samples[sample_idx][TYPE][prompt_idx]235 elif DATASET == 'KITTI': 236 vx_ = configs.KITTI_samples[sample_idx][TYPE][prompt_idx]237 elif DATASET == 'Semantic3D': 238 vx_ = configs.Semantic3D_samples[sample_idx][TYPE][prompt_idx]239 240 return gr.Textbox(label="Response", value="", visible=True), gr.Slider(minimum=0.01, maximum=0.15, step=0.001, label="Voxel Size", value=vx_)241 242def main():243 title = """<h1 style="text-align: center;">244 <div style="width: 1.2em; height: 1.2em; display: inline-block;"><img src="https://github.com/ZiyuGuo99/ZiyuGuo99.github.io/blob/main/assets/img/logo.png?raw=true" style='width: 100%; height: 100%; object-fit: contain;' /></div>245 <span style="font-variant: small-caps; font-weight: bold;">Sam2Point</span>246 </h1>247 <h3 align="center"><span style="font-variant: small-caps; ">Segment Any 3D as Videos in Zero-shot and Promptable Manners248 </span></h3>249 250 <div style="text-align: center;">251 <div style="display: flex; align-items: center; justify-content: center; gap: 0.5rem; margin-bottom: 0.5rem; font-size: 1rem; flex-wrap: wrap;">252 <a href="https://sam2point.github.io/" target="_blank">[Webpage]</a>253 <a href="https://arxiv.org/pdf/2408.16768" target="_blank">[Paper]</a>254 <a href="https://github.com/ZiyuGuo99/SAM2Point" target="_blank">[Code]</a>255 </div>256 </div>257 <p style="text-align: center;">258 Select an example and a 3D prompt to start segmentation using <span style="font-variant: small-caps;">Sam2Point</span>. 259 </p>260 <p style="text-align: center;">261 Custom 3D input and prompts will be supported soon.262 </p>263 """264 265 with gr.Blocks(266 css="""267 .contain { display: flex; flex-direction: column; }268 .gradio-container { height: 100vh !important; }269 #col_container { height: 100%; }270 pre {271 white-space: pre-wrap; /* Since CSS 2.1 */272 white-space: -moz-pre-wrap; /* Mozilla, since 1999 */273 white-space: -pre-wrap; /* Opera 4-6 */274 white-space: -o-pre-wrap; /* Opera 7 */275 word-wrap: break-word; /* Internet Explorer 5.5+ */276 }""",277 js="""278 function refresh() {279 const url = new URL(window.location);280 if (url.searchParams.get('__theme') !== 'light') {281 url.searchParams.set('__theme', 'light');282 window.location.href = url.href;283 }284 }""",285 title="SAM2Point: Segment Any 3D as Videos in Zero-shot and Promptable Manners",286 theme=gr.themes.Soft()287 ) as app:288 gr.HTML(title)289 with gr.Row():290 with gr.Column(elem_id="col_container"):291 sample_dropdown = gr.Dropdown(label="Select 3D Data Type", choices=samples, type="value")292 scene_dropdown = gr.Radio(label="Select 3D Object/Scene", choices=[], type="index")293 show_button = gr.Button("Show 3D Scene/Object")294 prompt_type_dropdown = gr.Radio(label="Select Prompt Type", choices=prompt_types)295 prompt_sample_dropdown = gr.Radio(label="Select Prompt Example", choices=[], type="index")296 show_prompt_button = gr.Button("Show Prompt in 3D Scene/Object")297 with gr.Column():298 start_segment_button = gr.Button("Start Segmentation")299 plot1 = gr.Plot()300 301 response = gr.Textbox(label="Response")302 303 sample_dropdown.change(update1, sample_dropdown, [scene_dropdown, response])304 sample_dropdown.change(update2, [sample_dropdown, scene_dropdown, prompt_type_dropdown], [prompt_sample_dropdown, response])305 scene_dropdown.change(update2, [sample_dropdown, scene_dropdown, prompt_type_dropdown], [prompt_sample_dropdown, response])306 prompt_type_dropdown.change(update2, [sample_dropdown, scene_dropdown, prompt_type_dropdown], [prompt_sample_dropdown, response])307 308 show_button.click(load_3d_scene, inputs=[sample_dropdown, scene_dropdown], outputs=plot1)309 show_prompt_button.click(show_prompt_in_3d, inputs=[sample_dropdown, scene_dropdown, prompt_type_dropdown, prompt_sample_dropdown], outputs=[plot1, response])310 start_segment_button.click(start_segmentation, inputs=[sample_dropdown, scene_dropdown, prompt_type_dropdown, prompt_sample_dropdown], outputs=[plot1, response])311 312 app.queue(max_size=20, api_open=False)313 app.launch(max_threads=400)314 315if __name__ == "__main__":316 main()