pablovela5620 commited on
Commit
9ec56ae
1 Parent(s): 20fff88

update app to use new version of mini-dust3r

Browse files
Files changed (2) hide show
  1. app.py +25 -16
  2. requirements.txt +1 -1
app.py CHANGED
@@ -3,40 +3,49 @@ import spaces
3
  import torch
4
  from gradio_rerun import Rerun
5
  import rerun as rr
 
6
  from pathlib import Path
7
-
8
 
9
  from mini_dust3r.api import OptimizedResult, inferece_dust3r, log_optimized_result
10
  from mini_dust3r.model import AsymmetricCroCo3DStereo
11
 
12
- DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
13
  model = AsymmetricCroCo3DStereo.from_pretrained(
14
- "naver/DUSt3R_ViTLarge_BaseDecoder_512_dpt"
15
- ).to(DEVICE)
 
16
 
17
  @spaces.GPU
18
- def predict(image_dir: str):
19
- rr.init("my data")
 
 
 
20
  optimized_results: OptimizedResult = inferece_dust3r(
21
- image_dir=image_dir,
22
  model=model,
23
  device=DEVICE,
24
  batch_size=1,
25
  )
26
- log_optimized_result(optimized_results, Path("world"))
27
- rr.save("dust3r.rrd")
28
- return "dust3r.rrd"
 
 
29
 
30
- with gr.Blocks(css=""".gradio-container {margin: 0 !important; min-width: 100%};""", title="DUSt3R Demo") as demo:
 
 
 
 
31
  # scene state is save so that you can change conf_thr, cam_size... without rerunning the inference
32
- gr.HTML('<h2 style="text-align: center;">DUSt3R Demo</h2>')
33
  with gr.Column():
34
  inputfiles = gr.File(file_count="multiple")
35
  rerun_viewer = Rerun(height=900)
36
 
37
  run_btn = gr.Button("Run")
38
- run_btn.click(fn=predict,
39
- inputs=[inputfiles],
40
- outputs=[rerun_viewer])
41
-
42
  demo.launch()
 
3
  import torch
4
  from gradio_rerun import Rerun
5
  import rerun as rr
6
+ import rerun.blueprint as rrb
7
  from pathlib import Path
8
+ import uuid
9
 
10
  from mini_dust3r.api import OptimizedResult, inferece_dust3r, log_optimized_result
11
  from mini_dust3r.model import AsymmetricCroCo3DStereo
12
 
13
+ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
14
  model = AsymmetricCroCo3DStereo.from_pretrained(
15
+ "naver/DUSt3R_ViTLarge_BaseDecoder_512_dpt"
16
+ ).to(DEVICE)
17
+
18
 
19
  @spaces.GPU
20
+ def predict(image_name_list: list[str]):
21
+ uuid_str = str(uuid.uuid4())
22
+ filename = Path(f"/tmp/gradio/{uuid_str}.rrd")
23
+ rr.init(f"{uuid_str}")
24
+ log_path = Path("world")
25
  optimized_results: OptimizedResult = inferece_dust3r(
26
+ image_dir_or_list=image_name_list,
27
  model=model,
28
  device=DEVICE,
29
  batch_size=1,
30
  )
31
+ rr.set_time_sequence("sequence", 0)
32
+ log_optimized_result(optimized_results, log_path)
33
+ # blueprint = rrb.Spatial3DView(origin="cube")
34
+ rr.save(filename.as_posix())
35
+ return filename.as_posix()
36
 
37
+
38
+ with gr.Blocks(
39
+ css=""".gradio-container {margin: 0 !important; min-width: 100%};""",
40
+ title="Mini-DUSt3R Demo",
41
+ ) as demo:
42
  # scene state is save so that you can change conf_thr, cam_size... without rerunning the inference
43
+ gr.HTML('<h2 style="text-align: center;">Mini-DUSt3R Demo</h2>')
44
  with gr.Column():
45
  inputfiles = gr.File(file_count="multiple")
46
  rerun_viewer = Rerun(height=900)
47
 
48
  run_btn = gr.Button("Run")
49
+ run_btn.click(fn=predict, inputs=[inputfiles], outputs=[rerun_viewer])
50
+
 
 
51
  demo.launch()
requirements.txt CHANGED
@@ -1 +1 @@
1
- mini-dust3r==0.1.0
 
1
+ mini-dust3r==0.1.1