Spaces:

Tremick
/

SoftwareTOF

Running on Zero

Perf: binary PLY, downsample input, subsample 200k points max

7012f96 3 days ago

7.54 kB

import spaces  # MUST be the very first import for ZeroGPU to work

import io
import numpy as np
from PIL import Image, UnidentifiedImageError
import torch
import gradio as gr
from fastapi import File, UploadFile, HTTPException
from fastapi.responses import Response
from transformers import DepthProImageProcessor, DepthProForDepthEstimation


# =====================================================================
# 1. Load Apple Depth Pro at module level.
#    ZeroGPU safely emulates CUDA during startup so .to("cuda") works here.
# =====================================================================
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype  = torch.float16 if torch.cuda.is_available() else torch.float32

print("Loading Apple Depth Pro components...")
image_processor = DepthProImageProcessor.from_pretrained("apple/DepthPro-hf")
model = DepthProForDepthEstimation.from_pretrained(
    "apple/DepthPro-hf",
    torch_dtype=dtype,
).to(device)
model.eval()
print(f"Model successfully loaded on target device: {device}")


# =====================================================================
# 2. Core inference helper (shared by Gradio UI and the REST endpoint)
# =====================================================================
def _run_depth_inference(image: Image.Image):
    """Returns post-processed depth output list. Caller handles indexing."""
    inputs = image_processor(images=image, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = model(**inputs)
    return image_processor.post_process_depth_estimation(
        outputs,
        target_sizes=[(image.height, image.width)],
    )


# =====================================================================
# 3. Gradio function - @spaces.GPU tells ZeroGPU to allocate a GPU here.
#    This MUST be bound to a Gradio event handler for ZeroGPU to detect it.
# =====================================================================
@spaces.GPU(duration=120)
def create_3d_point_cloud(input_pil_image):
    """Processes the image, calculates 3D vertices, and writes a binary .ply file."""
    if input_pil_image is None:
        return None

    # Step A: Downscale large images before inference.
    # DepthPro quality is stable at 512px; this dramatically cuts runtime.
    image = input_pil_image.convert("RGB")
    MAX_DIM = 512
    if max(image.width, image.height) > MAX_DIM:
        image.thumbnail((MAX_DIM, MAX_DIM), Image.LANCZOS)

    post_processed = _run_depth_inference(image)
    depth_matrix = (
        post_processed[0]["predicted_depth"]
        .detach().cpu().numpy().astype(np.float32)
    )

    # Step B: Pinhole camera projection (u, v, depth) -> (X, Y, Z)
    height, width = depth_matrix.shape
    color_np = np.array(image, dtype=np.float32) / 255.0

    focal_length = float(max(width, height))
    cx, cy = width / 2.0, height / 2.0
    u, v = np.meshgrid(np.arange(width, dtype=np.float32),
                        np.arange(height, dtype=np.float32))

    z = depth_matrix.flatten()
    x = ((u.flatten() - cx) * z) / focal_length
    y = ((v.flatten() - cy) * z) / focal_length

    # Keep objects between 10 cm and 30 m; discard sky / infinite background
    valid_mask = (z > 0.1) & (z < 30.0)
    x, y, z = x[valid_mask], y[valid_mask], z[valid_mask]
    rgb_ints = (color_np.reshape(-1, 3)[valid_mask] * 255).astype(np.uint8)

    # Step C: Subsample to keep the browser renderer fast (<= 200k points)
    MAX_POINTS = 200_000
    if len(z) > MAX_POINTS:
        idx = np.random.choice(len(z), MAX_POINTS, replace=False)
        x, y, z = x[idx], y[idx], z[idx]
        rgb_ints = rgb_ints[idx]

    # Step D: Write BINARY PLY via numpy (100x faster than an ASCII for-loop)
    num_points = len(z)
    vertex_dtype = np.dtype([
        ('x', '<f4'), ('y', '<f4'), ('z', '<f4'),
        ('red', 'u1'), ('green', 'u1'), ('blue', 'u1'),
    ])
    vertices = np.empty(num_points, dtype=vertex_dtype)
    vertices['x']     = x.astype(np.float32)
    vertices['y']     = (-y).astype(np.float32)   # invert Y so it loads right-side up
    vertices['z']     = z.astype(np.float32)
    vertices['red']   = rgb_ints[:, 0]
    vertices['green'] = rgb_ints[:, 1]
    vertices['blue']  = rgb_ints[:, 2]

    ply_header = (
        "ply\n"
        "format binary_little_endian 1.0\n"
        f"element vertex {num_points}\n"
        "property float x\n"
        "property float y\n"
        "property float z\n"
        "property uchar red\n"
        "property uchar green\n"
        "property uchar blue\n"
        "end_header\n"
    )

    ply_filename = "/tmp/output_scene.ply"
    with open(ply_filename, "wb") as f:
        f.write(ply_header.encode("ascii"))
        f.write(vertices.tobytes())

    return ply_filename


# =====================================================================
# 4. Gradio UI - uses demo.launch() (NOT gr.mount_gradio_app) so that
#    ZeroGPU can properly detect the @spaces.GPU decorated function above.
# =====================================================================
with gr.Blocks(title="Software-ToF 3D Cloud Visualizer") as demo:
    gr.Markdown("# 📷 Software-ToF Camera Simulation")
    gr.Markdown(
        "Upload any image to calculate absolute metric depths and generate "
        "an interactive 3D Point Cloud visualization."
    )

    with gr.Row():
        with gr.Column():
            input_img  = gr.Image(type="pil", label="Input Scene Photo")
            submit_btn = gr.Button("Generate ToF 3D Structure", variant="primary")
        with gr.Column():
            output_3d = gr.Model3D(
                label="Interactive 3D Point Cloud Canvas",
                display_mode="point_cloud",
            )

    submit_btn.click(fn=create_3d_point_cloud, inputs=[input_img], outputs=[output_3d])


# =====================================================================
# 5. Add the REST /get_metric_depth endpoint via Gradio's internal
#    FastAPI app - this avoids gr.mount_gradio_app which breaks ZeroGPU.
# =====================================================================
@demo.app.get("/status")
def status():
    return {"status": "running", "device": str(device)}


@demo.app.post("/get_metric_depth")
async def get_metric_depth(file: UploadFile = File(...)):
    if not (file.content_type or "").startswith("image/"):
        raise HTTPException(status_code=400, detail="Uploaded file must be an image.")

    try:
        image_bytes = await file.read()
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    except (UnidentifiedImageError, OSError):
        raise HTTPException(status_code=400, detail="Could not decode image.")

    try:
        post_processed = _run_depth_inference(image)
        metric_depth_map = (
            post_processed[0]["predicted_depth"]
            .detach().cpu().numpy().astype(np.float32)
        )
        buffer = io.BytesIO()
        np.save(buffer, metric_depth_map)
        return Response(content=buffer.getvalue(), media_type="application/octet-stream")

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Inference error: {str(e)}")


# =====================================================================
# 6. Launch - Hugging Face auto-detects the `demo` variable when
#    sdk: gradio is set and launches it. The explicit launch() call
#    here ensures it also works when run locally.
# =====================================================================
if __name__ == "__main__":
    demo.launch()

	import spaces # MUST be the very first import for ZeroGPU to work

	import io
	import numpy as np
	from PIL import Image, UnidentifiedImageError
	import torch
	import gradio as gr
	from fastapi import File, UploadFile, HTTPException
	from fastapi.responses import Response
	from transformers import DepthProImageProcessor, DepthProForDepthEstimation


	# =====================================================================
	# 1. Load Apple Depth Pro at module level.
	# ZeroGPU safely emulates CUDA during startup so .to("cuda") works here.
	# =====================================================================
	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
	dtype = torch.float16 if torch.cuda.is_available() else torch.float32

	print("Loading Apple Depth Pro components...")
	image_processor = DepthProImageProcessor.from_pretrained("apple/DepthPro-hf")
	model = DepthProForDepthEstimation.from_pretrained(
	"apple/DepthPro-hf",
	torch_dtype=dtype,
	).to(device)
	model.eval()
	print(f"Model successfully loaded on target device: {device}")


	# =====================================================================
	# 2. Core inference helper (shared by Gradio UI and the REST endpoint)
	# =====================================================================
	def _run_depth_inference(image: Image.Image):
	"""Returns post-processed depth output list. Caller handles indexing."""
	inputs = image_processor(images=image, return_tensors="pt").to(device)
	with torch.no_grad():
	outputs = model(**inputs)
	return image_processor.post_process_depth_estimation(
	outputs,
	target_sizes=[(image.height, image.width)],
	)


	# =====================================================================
	# 3. Gradio function - @spaces.GPU tells ZeroGPU to allocate a GPU here.
	# This MUST be bound to a Gradio event handler for ZeroGPU to detect it.
	# =====================================================================
	@spaces.GPU(duration=120)
	def create_3d_point_cloud(input_pil_image):
	"""Processes the image, calculates 3D vertices, and writes a binary .ply file."""
	if input_pil_image is None:
	return None

	# Step A: Downscale large images before inference.
	# DepthPro quality is stable at 512px; this dramatically cuts runtime.
	image = input_pil_image.convert("RGB")
	MAX_DIM = 512
	if max(image.width, image.height) > MAX_DIM:
	image.thumbnail((MAX_DIM, MAX_DIM), Image.LANCZOS)

	post_processed = _run_depth_inference(image)
	depth_matrix = (
	post_processed[0]["predicted_depth"]
	.detach().cpu().numpy().astype(np.float32)
	)

	# Step B: Pinhole camera projection (u, v, depth) -> (X, Y, Z)
	height, width = depth_matrix.shape
	color_np = np.array(image, dtype=np.float32) / 255.0

	focal_length = float(max(width, height))
	cx, cy = width / 2.0, height / 2.0
	u, v = np.meshgrid(np.arange(width, dtype=np.float32),
	np.arange(height, dtype=np.float32))

	z = depth_matrix.flatten()
	x = ((u.flatten() - cx) * z) / focal_length
	y = ((v.flatten() - cy) * z) / focal_length

	# Keep objects between 10 cm and 30 m; discard sky / infinite background
	valid_mask = (z > 0.1) & (z < 30.0)
	x, y, z = x[valid_mask], y[valid_mask], z[valid_mask]
	rgb_ints = (color_np.reshape(-1, 3)[valid_mask] * 255).astype(np.uint8)

	# Step C: Subsample to keep the browser renderer fast (<= 200k points)
	MAX_POINTS = 200_000
	if len(z) > MAX_POINTS:
	idx = np.random.choice(len(z), MAX_POINTS, replace=False)
	x, y, z = x[idx], y[idx], z[idx]
	rgb_ints = rgb_ints[idx]

	# Step D: Write BINARY PLY via numpy (100x faster than an ASCII for-loop)
	num_points = len(z)
	vertex_dtype = np.dtype([
	('x', '<f4'), ('y', '<f4'), ('z', '<f4'),
	('red', 'u1'), ('green', 'u1'), ('blue', 'u1'),
	])
	vertices = np.empty(num_points, dtype=vertex_dtype)
	vertices['x'] = x.astype(np.float32)
	vertices['y'] = (-y).astype(np.float32) # invert Y so it loads right-side up
	vertices['z'] = z.astype(np.float32)
	vertices['red'] = rgb_ints[:, 0]
	vertices['green'] = rgb_ints[:, 1]
	vertices['blue'] = rgb_ints[:, 2]

	ply_header = (
	"ply\n"
	"format binary_little_endian 1.0\n"
	f"element vertex {num_points}\n"
	"property float x\n"
	"property float y\n"
	"property float z\n"
	"property uchar red\n"
	"property uchar green\n"
	"property uchar blue\n"
	"end_header\n"
	)

	ply_filename = "/tmp/output_scene.ply"
	with open(ply_filename, "wb") as f:
	f.write(ply_header.encode("ascii"))
	f.write(vertices.tobytes())

	return ply_filename


	# =====================================================================
	# 4. Gradio UI - uses demo.launch() (NOT gr.mount_gradio_app) so that
	# ZeroGPU can properly detect the @spaces.GPU decorated function above.
	# =====================================================================
	with gr.Blocks(title="Software-ToF 3D Cloud Visualizer") as demo:
	gr.Markdown("# 📷 Software-ToF Camera Simulation")
	gr.Markdown(
	"Upload any image to calculate absolute metric depths and generate "
	"an interactive 3D Point Cloud visualization."
	)

	with gr.Row():
	with gr.Column():
	input_img = gr.Image(type="pil", label="Input Scene Photo")
	submit_btn = gr.Button("Generate ToF 3D Structure", variant="primary")
	with gr.Column():
	output_3d = gr.Model3D(
	label="Interactive 3D Point Cloud Canvas",
	display_mode="point_cloud",
	)

	submit_btn.click(fn=create_3d_point_cloud, inputs=[input_img], outputs=[output_3d])


	# =====================================================================
	# 5. Add the REST /get_metric_depth endpoint via Gradio's internal
	# FastAPI app - this avoids gr.mount_gradio_app which breaks ZeroGPU.
	# =====================================================================
	@demo.app.get("/status")
	def status():
	return {"status": "running", "device": str(device)}


	@demo.app.post("/get_metric_depth")
	async def get_metric_depth(file: UploadFile = File(...)):
	if not (file.content_type or "").startswith("image/"):
	raise HTTPException(status_code=400, detail="Uploaded file must be an image.")

	try:
	image_bytes = await file.read()
	image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
	except (UnidentifiedImageError, OSError):
	raise HTTPException(status_code=400, detail="Could not decode image.")

	try:
	post_processed = _run_depth_inference(image)
	metric_depth_map = (
	post_processed[0]["predicted_depth"]
	.detach().cpu().numpy().astype(np.float32)
	)
	buffer = io.BytesIO()
	np.save(buffer, metric_depth_map)
	return Response(content=buffer.getvalue(), media_type="application/octet-stream")

	except Exception as e:
	raise HTTPException(status_code=500, detail=f"Inference error: {str(e)}")


	# =====================================================================
	# 6. Launch - Hugging Face auto-detects the `demo` variable when
	# sdk: gradio is set and launches it. The explicit launch() call
	# here ensures it also works when run locally.
	# =====================================================================
	if __name__ == "__main__":
	demo.launch()