Spaces:
Runtime error
Runtime error
Rawal Khirodkar commited on
Commit ·
2593450
1
Parent(s): 9705303
Pointmap: drop bg-removal toggle; fg-bg is mandatory
Browse files
app.py
CHANGED
|
@@ -59,8 +59,6 @@ DEFAULT_SIZE = "0.4B" # iteration mode — only this is preloaded; others lazy-
|
|
| 59 |
|
| 60 |
FG_REPO = "facebook/sapiens-seg-foreground-1b-torchscript"
|
| 61 |
FG_FILENAME = "sapiens_1b_seg_foreground_epoch_8_torchscript.pt2"
|
| 62 |
-
BG_OPTIONS = ["fg-bg", "no-bg-removal"]
|
| 63 |
-
DEFAULT_BG = "fg-bg"
|
| 64 |
|
| 65 |
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
| 66 |
|
|
@@ -198,7 +196,7 @@ def _make_ply(image_rgb: np.ndarray, pointmap_hwc: np.ndarray, mask_hw: np.ndarr
|
|
| 198 |
# Gradio handler
|
| 199 |
|
| 200 |
@spaces.GPU(duration=180)
|
| 201 |
-
def predict(image: Image.Image, size: str
|
| 202 |
if image is None:
|
| 203 |
return None, None
|
| 204 |
|
|
@@ -210,7 +208,9 @@ def predict(image: Image.Image, size: str, bg_mode: str):
|
|
| 210 |
model = _get_pointmap_model(size)
|
| 211 |
pointmap = _estimate_pointmap(image_bgr, model)
|
| 212 |
|
| 213 |
-
|
|
|
|
|
|
|
| 214 |
ply_path = _make_ply(image_rgb, pointmap, mask)
|
| 215 |
|
| 216 |
return ply_path, ply_path
|
|
@@ -295,13 +295,7 @@ with gr.Blocks(title="Sapiens2 Pointmap", theme=gr.themes.Soft(), css=CUSTOM_CSS
|
|
| 295 |
choices=["0.4B"], # iteration mode — re-add other sizes when shipping
|
| 296 |
value=DEFAULT_SIZE,
|
| 297 |
label="Model",
|
| 298 |
-
scale=
|
| 299 |
-
)
|
| 300 |
-
bg = gr.Radio(
|
| 301 |
-
choices=BG_OPTIONS,
|
| 302 |
-
value=DEFAULT_BG,
|
| 303 |
-
label="Background",
|
| 304 |
-
scale=2,
|
| 305 |
)
|
| 306 |
run = gr.Button("Run", variant="primary", size="lg", scale=1)
|
| 307 |
|
|
@@ -310,7 +304,7 @@ with gr.Blocks(title="Sapiens2 Pointmap", theme=gr.themes.Soft(), css=CUSTOM_CSS
|
|
| 310 |
with gr.Accordion("Raw Pointmap", open=False):
|
| 311 |
out_ply_file = gr.File(label="Point cloud (.ply — open in MeshLab/CloudCompare/Blender)")
|
| 312 |
|
| 313 |
-
run.click(predict, inputs=[inp, size
|
| 314 |
|
| 315 |
|
| 316 |
if __name__ == "__main__":
|
|
|
|
| 59 |
|
| 60 |
FG_REPO = "facebook/sapiens-seg-foreground-1b-torchscript"
|
| 61 |
FG_FILENAME = "sapiens_1b_seg_foreground_epoch_8_torchscript.pt2"
|
|
|
|
|
|
|
| 62 |
|
| 63 |
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
| 64 |
|
|
|
|
| 196 |
# Gradio handler
|
| 197 |
|
| 198 |
@spaces.GPU(duration=180)
|
| 199 |
+
def predict(image: Image.Image, size: str):
|
| 200 |
if image is None:
|
| 201 |
return None, None
|
| 202 |
|
|
|
|
| 208 |
model = _get_pointmap_model(size)
|
| 209 |
pointmap = _estimate_pointmap(image_bgr, model)
|
| 210 |
|
| 211 |
+
# Foreground masking is mandatory — keeps the cloud clean and the camera
|
| 212 |
+
# marker meaningful (background depth is unreliable).
|
| 213 |
+
mask = _foreground_mask(image_pil, h0, w0)
|
| 214 |
ply_path = _make_ply(image_rgb, pointmap, mask)
|
| 215 |
|
| 216 |
return ply_path, ply_path
|
|
|
|
| 295 |
choices=["0.4B"], # iteration mode — re-add other sizes when shipping
|
| 296 |
value=DEFAULT_SIZE,
|
| 297 |
label="Model",
|
| 298 |
+
scale=4,
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 299 |
)
|
| 300 |
run = gr.Button("Run", variant="primary", size="lg", scale=1)
|
| 301 |
|
|
|
|
| 304 |
with gr.Accordion("Raw Pointmap", open=False):
|
| 305 |
out_ply_file = gr.File(label="Point cloud (.ply — open in MeshLab/CloudCompare/Blender)")
|
| 306 |
|
| 307 |
+
run.click(predict, inputs=[inp, size], outputs=[out_ply, out_ply_file])
|
| 308 |
|
| 309 |
|
| 310 |
if __name__ == "__main__":
|