cf402e322c
Movie SBS converges on the screen plane, so real content sits at NEGATIVE disparity; minDisparity=0 clipped everything behind convergence to black. Window is now ±num_disp/2 with 2-98 percentile normalization and a validity mask. Verified on Big Buck Bunny (half-OU) + One Night In Hell (half-SBS): structured depth + colored clouds that render in Blender. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
113 lines
4.8 KiB
Python
113 lines
4.8 KiB
Python
"""stereo_depth — SBS stereoscopic video/frame → SGBM disparity maps + mid-frame point cloud.
|
|
|
|
True geometric depth from the built-in L/R parallax (cv2.StereoSGBM). ffmpeg (system)
|
|
decodes/splits; depth PNGs + preview mp4 + a colored .ply land as the outputs.
|
|
ponytail: uncalibrated disparity (SBS rips have no camera intrinsics) — z is relative,
|
|
perfect for Blender displace planes; not metric reconstruction.
|
|
"""
|
|
import argparse
|
|
import glob
|
|
import json
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import zipfile
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
FFMPEG = shutil.which("ffmpeg") or "/opt/homebrew/bin/ffmpeg" # headless ssh has no brew PATH
|
|
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--input", action="append", default=[])
|
|
ap.add_argument("--outdir", required=True)
|
|
ap.add_argument("--params", default="{}")
|
|
a = ap.parse_args()
|
|
p = json.loads(a.params)
|
|
|
|
if not a.input:
|
|
print("ERROR: no input")
|
|
sys.exit(1)
|
|
|
|
src = a.input[0]
|
|
outdir = a.outdir
|
|
max_frames = int(p.get("max_frames", 120))
|
|
num_disp = max(16, (int(p.get("num_disp", 128)) // 16) * 16)
|
|
block = int(p.get("block", 5)) | 1
|
|
tab = p.get("layout", "sbs") == "tab"
|
|
|
|
work = tempfile.mkdtemp(prefix="stereo_")
|
|
frames_dir = os.path.join(work, "frames")
|
|
depth_dir = os.path.join(work, "depth")
|
|
os.makedirs(frames_dir)
|
|
os.makedirs(depth_dir)
|
|
|
|
# decode → frames (a single image just becomes frame 1)
|
|
subprocess.run([FFMPEG, "-y", "-loglevel", "error", "-i", src,
|
|
"-frames:v", str(max_frames), os.path.join(frames_dir, "f%05d.png")], check=True)
|
|
frames = sorted(glob.glob(os.path.join(frames_dir, "*.png")))
|
|
print(f"[stereo] {len(frames)} frames, num_disp={num_disp} block={block}", flush=True)
|
|
|
|
min_disp = -(num_disp // 2) # movie stereo converges on the screen plane — content
|
|
sgbm = cv2.StereoSGBM_create( # lives BOTH sides of it, so the window must go negative
|
|
minDisparity=min_disp, numDisparities=num_disp, blockSize=block,
|
|
P1=8 * 3 * block * block, P2=32 * 3 * block * block,
|
|
disp12MaxDiff=1, uniquenessRatio=10, speckleWindowSize=100, speckleRange=2,
|
|
mode=cv2.STEREO_SGBM_MODE_SGBM_3WAY)
|
|
|
|
mid = frames[len(frames) // 2]
|
|
for i, f in enumerate(frames):
|
|
img = cv2.imread(f)
|
|
h, w = img.shape[:2]
|
|
if tab:
|
|
L, R = img[: h // 2], img[h // 2:]
|
|
else:
|
|
L, R = img[:, : w // 2], img[:, w // 2:]
|
|
gl, gr = cv2.cvtColor(L, cv2.COLOR_BGR2GRAY), cv2.cvtColor(R, cv2.COLOR_BGR2GRAY)
|
|
disp = sgbm.compute(gl, gr).astype(np.float32) / 16.0
|
|
valid = disp > (min_disp - 0.5) # cv2 marks unmatched pixels min_disp-1
|
|
if valid.sum() > 100:
|
|
lo, hi = np.percentile(disp[valid], [2, 98])
|
|
else:
|
|
lo, hi = 0.0, 1.0
|
|
norm = np.zeros(disp.shape, np.uint8)
|
|
norm[valid] = (np.clip((disp[valid] - lo) / max(1e-6, hi - lo), 0, 1) * 255).astype(np.uint8)
|
|
cv2.imwrite(os.path.join(depth_dir, os.path.basename(f)), norm)
|
|
if f == mid: # colored point cloud of the middle frame
|
|
ys, xs = np.where(valid & (norm > 2))
|
|
step = max(1, len(ys) // 400_000) # cap cloud size
|
|
ys, xs = ys[::step], xs[::step]
|
|
z = disp[ys, xs]
|
|
rgb = L[ys, xs][:, ::-1] # BGR→RGB
|
|
ply = os.path.join(outdir, "midframe_cloud.ply")
|
|
with open(ply, "w") as fh:
|
|
fh.write("ply\nformat ascii 1.0\n"
|
|
f"element vertex {len(ys)}\n"
|
|
"property float x\nproperty float y\nproperty float z\n"
|
|
"property uchar red\nproperty uchar green\nproperty uchar blue\nend_header\n")
|
|
for (yy, xx, zz, (r, g, b)) in zip(ys, xs, z, rgb):
|
|
fh.write(f"{xx} {-yy} {zz:.1f} {r} {g} {b}\n")
|
|
print(f"[stereo] wrote midframe_cloud.ply ({len(ys)} pts)", flush=True)
|
|
if i % 24 == 0:
|
|
print(f"[stereo] frame {i + 1}/{len(frames)}", flush=True)
|
|
|
|
outputs = [{"path": "midframe_cloud.ply", "meta": {"tool": "stereo_sgbm"}}]
|
|
if len(frames) > 1: # depth preview video
|
|
prev = os.path.join(outdir, "depth_preview.mp4")
|
|
subprocess.run([FFMPEG, "-y", "-loglevel", "error", "-framerate", "24",
|
|
"-i", os.path.join(depth_dir, "f%05d.png"),
|
|
"-pix_fmt", "yuv420p", prev], check=True)
|
|
outputs.append({"path": "depth_preview.mp4", "meta": {"tool": "stereo_sgbm"}})
|
|
|
|
zpath = os.path.join(outdir, "depth_frames.zip")
|
|
with zipfile.ZipFile(zpath, "w", zipfile.ZIP_DEFLATED) as zf:
|
|
for f in sorted(glob.glob(os.path.join(depth_dir, "*.png"))):
|
|
zf.write(f, os.path.basename(f))
|
|
outputs.append({"path": "depth_frames.zip", "meta": {"tool": "stereo_sgbm"}})
|
|
|
|
json.dump({"outputs": outputs}, open(os.path.join(outdir, "result.json"), "w"))
|
|
shutil.rmtree(work, ignore_errors=True)
|
|
print("done:", ", ".join(o["path"] for o in outputs), flush=True)
|