steam-frame/its-ven--FrameDepth
1
0 Fork 0
mirror of https://github.com/its-ven/FrameDepth.git synced 2026-10-17 16:01:39 +02:00
Code Issues Packages Projects Releases Wiki Activity
Files
a561b5b9f0accd0b3dd08b294bbba33cd3ee70e3
its-ven--FrameDepth/depth3d.py
T
its-ven 40c11fdb2f Add files via upload
2026-10-11 22:06:14 +02:00

448 lines
17 KiB
Python
Raw Blame History

"""
depth3d.py - screen capture + depth estimation core, and a simple side-by-side (SBS) window.
Two roles:
* Library: `Grabber` (Windows Graphics Capture) and `DepthWorker` (Depth Anything V2 on the
GPU, scene-cut aware) are used by depth3d_vr.py, the main SteamVR player.
* Standalone: `python depth3d.py --source-monitor 1` opens a plain SBS window that you can
feed to a 3D-capable viewer (Desktop+ HSBS, Bigscreen...). Simpler and less capable than
the VR player: older warp shader, no menu, no environments.
Controls of the standalone SBS window:
Up / Down depth strength
Left / Right convergence (which depth sits at the screen plane)
X swap eyes
D toggle depth-map debug view
F toggle fullscreen
Esc quit
"""
import argparse
import threading
import time
import glfw
import moderngl
import numpy as np
import torch
import torch.nn.functional as F
from transformers import AutoModelForDepthEstimation
from windows_capture import Frame, InternalCaptureControl, WindowsCapture
MODELS = {
"small": "depth-anything/Depth-Anything-V2-Small-hf",
"base": "depth-anything/Depth-Anything-V2-Base-hf",
"large": "depth-anything/Depth-Anything-V2-Large-hf",
}
# --------------------------------------------------------------------------
# Window listing helper
# --------------------------------------------------------------------------
def list_windows():
import ctypes
from ctypes import wintypes
user32 = ctypes.windll.user32
titles = []
@ctypes.WINFUNCTYPE(wintypes.BOOL, wintypes.HWND, wintypes.LPARAM)
def cb(hwnd, _):
if user32.IsWindowVisible(hwnd):
n = user32.GetWindowTextLengthW(hwnd)
if n:
buf = ctypes.create_unicode_buffer(n + 1)
user32.GetWindowTextW(hwnd, buf, n + 1)
titles.append(buf.value)
return True
user32.EnumWindows(cb, 0)
return titles
# --------------------------------------------------------------------------
# Capture
# --------------------------------------------------------------------------
class Grabber:
"""Keeps only the newest captured frame (BGRA, HxWx4 uint8)."""
def __init__(self, window_name=None, monitor_index=None, cursor=True):
self.lock = threading.Lock()
self.frame = None
self.seq = 0
kwargs = dict(cursor_capture=cursor, draw_border=False)
if window_name:
kwargs["window_name"] = window_name
else:
kwargs["monitor_index"] = monitor_index # 1-based
cap = WindowsCapture(**kwargs)
@cap.event
def on_frame_arrived(frame: Frame, ctrl: InternalCaptureControl):
buf = np.array(frame.frame_buffer, copy=True)
with self.lock:
self.frame = buf
self.seq += 1
@cap.event
def on_closed():
print("Capture session closed.")
self.control = cap.start_free_threaded()
def latest(self):
with self.lock:
return self.frame, self.seq
def stop(self):
try:
self.control.stop()
except Exception:
pass
# --------------------------------------------------------------------------
# Depth estimation worker
# --------------------------------------------------------------------------
class DepthWorker(threading.Thread):
def __init__(self, grabber, width, smooth, model_id, cut_threshold=0.12):
super().__init__(daemon=True)
self.grabber = grabber
self.width = max(112, (width // 14) * 14)
self.smooth = smooth
self.cut_threshold = cut_threshold
self.active = True # set False to idle the model (e.g. when 3D is switched off)
self.frame = None # the captured frame the current depth belongs to
self.stop_evt = threading.Event()
self.lock = threading.Lock()
self.depth = None # float32 HxW in 0..1 (1 = near)
self.seq = 0
self.fps = 0.0
self.device = "cuda" if torch.cuda.is_available() else "cpu"
self.dtype = torch.float16 if self.device == "cuda" else torch.float32
if self.device == "cpu":
print("WARNING: CUDA not available, running on CPU (will be slow).")
print(f"Loading {model_id} on {self.device} ...")
self.model = (
AutoModelForDepthEstimation.from_pretrained(model_id)
.to(self.device, self.dtype)
.eval()
)
self.mean = torch.tensor([0.485, 0.456, 0.406], device=self.device).view(1, 3, 1, 1)
self.std = torch.tensor([0.229, 0.224, 0.225], device=self.device).view(1, 3, 1, 1)
def latest(self):
with self.lock:
return self.depth, self.seq
def latest_pair(self):
"""Depth plus the exact captured frame it was computed from (keeps them in sync)."""
with self.lock:
return self.depth, self.frame, self.seq
def run(self):
last_seq = -1
lo = hi = None
prev = None
prev_thumb = None
n, t0 = 0, time.time()
while not self.stop_evt.is_set():
if not self.active:
time.sleep(0.02)
continue
frame, seq = self.grabber.latest()
if frame is None or seq == last_seq:
time.sleep(0.002)
continue
last_seq = seq
with torch.inference_mode():
h, w = frame.shape[:2]
mw = self.width
mh = max(14, int(round(mw * h / w / 14)) * 14)
t = torch.from_numpy(frame).to(self.device, non_blocking=True)
t = t[..., [2, 1, 0]].permute(2, 0, 1).unsqueeze(0).float() / 255.0
t = F.interpolate(t, size=(mh, mw), mode="bilinear",
antialias=True, align_corners=False)
# Tiny luminance thumbnail -> scene-cut / motion detection
thumb = F.adaptive_avg_pool2d(t, (18, 32)).mean(1)
if prev_thumb is None:
diff, cut = 1.0, True
else:
diff = (thumb - prev_thumb).abs().mean().item()
cut = diff > self.cut_threshold
prev_thumb = thumb
if cut: # hard cut: forget everything from the old shot
lo = hi = None
prev = None
t = ((t - self.mean) / self.std).to(self.dtype)
d = self.model(pixel_values=t).predicted_depth[0].float() # HxW, larger = nearer
# Robust, temporally smoothed normalisation (avoids pumping/flicker)
sample = d.flatten()[::7]
q = torch.quantile(sample, torch.tensor([0.02, 0.98], device=self.device))
l, hgh = q[0].item(), q[1].item()
if lo is None:
lo, hi = l, hgh
else:
lo += 0.08 * (l - lo)
hi += 0.08 * (hgh - hi)
dn = ((d - lo) / max(hi - lo, 1e-6)).clamp(0, 1)
if prev is not None and prev.shape == dn.shape:
# calm shot -> smooth heavily; lots of motion -> trust the new frame
alpha = self.smooth + (1 - self.smooth) * min(1.0, diff / 0.06)
dn = prev * (1 - alpha) + dn * alpha
prev = dn
out = dn.cpu().numpy().astype(np.float32)
with self.lock:
self.depth = out
self.frame = frame
self.seq += 1
n += 1
if time.time() - t0 >= 2.0:
self.fps = n / (time.time() - t0)
n, t0 = 0, time.time()
# --------------------------------------------------------------------------
# Rendering
# --------------------------------------------------------------------------
VERT = """
#version 330
in vec2 pos;
out vec2 uv;
void main() {
uv = vec2(pos.x * 0.5 + 0.5, 1.0 - (pos.y * 0.5 + 0.5));
gl_Position = vec4(pos, 0.0, 1.0);
}
"""
FRAG = """
#version 330
uniform sampler2D color_tex;
uniform sampler2D depth_tex;
uniform float strength; // max horizontal shift per eye, as fraction of width
uniform float convergence; // depth value (0..1) that sits at the screen plane
uniform float eye; // -1 = left, +1 = right
uniform int steps;
uniform int show_depth;
in vec2 uv;
out vec4 fragColor;
void main() {
// Backward warp via fixed-point iteration: find the source pixel whose
// own depth-based shift lands on this output pixel.
vec2 src = uv;
for (int i = 0; i < steps; i++) {
float d = texture(depth_tex, src).r;
src.x = uv.x + eye * strength * (d - convergence);
}
if (show_depth == 1) {
fragColor = vec4(vec3(texture(depth_tex, uv).r), 1.0);
} else {
fragColor = vec4(texture(color_tex, src).rgb, 1.0);
}
}
"""
def make_tex(ctx, w, h, comps, dtype, data=None):
tex = ctx.texture((w, h), comps, data, dtype=dtype)
tex.filter = (moderngl.LINEAR, moderngl.LINEAR)
tex.repeat_x = False
tex.repeat_y = False
return tex
def main():
ap = argparse.ArgumentParser(description="Window -> depth -> SBS 3D")
ap.add_argument("--window", help="Title (or part of it) of the window to capture")
ap.add_argument("--source-monitor", type=int, default=None,
help="Capture a whole monitor instead of a window (1 = first monitor)")
ap.add_argument("--list", action="store_true",
help="List visible window titles and monitors, then exit")
ap.add_argument("--model", choices=list(MODELS), default="base",
help="Depth model size (small = fastest, large = best quality)")
ap.add_argument("--depth-width", type=int, default=700,
help="Model input width (lower = faster). Try 518 or 392.")
ap.add_argument("--strength", type=float, default=0.025)
ap.add_argument("--convergence", type=float, default=0.4)
ap.add_argument("--steps", type=int, default=4, help="Warp iterations (quality/cost)")
ap.add_argument("--smooth", type=float, default=0.5,
help="Depth temporal blend, 1 = no smoothing, lower = smoother")
ap.add_argument("--out-monitor", type=int, default=None,
help="Start the SBS output fullscreen on this monitor (0-based). "
"Must NOT be the monitor you are capturing.")
ap.add_argument("--width", type=int, default=1600)
ap.add_argument("--height", type=int, default=900)
args = ap.parse_args()
if args.list:
print("Windows:")
for t in sorted(set(list_windows())):
print(" ", t)
if glfw.init():
print("Monitors (--source-monitor is 1-based, --out-monitor is 0-based):")
for i, m in enumerate(glfw.get_monitors()):
vm = glfw.get_video_mode(m)
x, y = glfw.get_monitor_pos(m)
print(f" out={i} source={i + 1} {glfw.get_monitor_name(m).decode(errors='ignore')}"
f" {vm.size.width}x{vm.size.height} at ({x},{y})")
glfw.terminate()
return
if not args.window and args.source_monitor is None:
ap.error("give --window TITLE or --source-monitor N (use --list to see options)")
if args.window and args.source_monitor is not None:
ap.error("use either --window or --source-monitor, not both")
if args.source_monitor is not None and args.out_monitor is not None \
and args.source_monitor - 1 == args.out_monitor:
ap.error("output monitor and captured monitor are the same: that would feed back on itself")
grabber = Grabber(args.window, args.source_monitor)
worker = DepthWorker(grabber, args.depth_width, args.smooth, MODELS[args.model])
worker.start()
if not glfw.init():
raise RuntimeError("GLFW init failed")
glfw.window_hint(glfw.CONTEXT_VERSION_MAJOR, 3)
glfw.window_hint(glfw.CONTEXT_VERSION_MINOR, 3)
glfw.window_hint(glfw.OPENGL_PROFILE, glfw.OPENGL_CORE_PROFILE)
monitors = glfw.get_monitors()
mon = None
w, h = args.width, args.height
if args.out_monitor is not None and args.out_monitor < len(monitors):
mon = monitors[args.out_monitor]
mode = glfw.get_video_mode(mon)
w, h = mode.size.width, mode.size.height
win = glfw.create_window(w, h, "depth3d SBS", mon, None)
glfw.make_context_current(win)
glfw.swap_interval(1)
ctx = moderngl.create_context()
prog = ctx.program(vertex_shader=VERT, fragment_shader=FRAG)
vbo = ctx.buffer(np.array([-1, -1, 1, -1, -1, 1, 1, 1], dtype="f4").tobytes())
vao = ctx.vertex_array(prog, [(vbo, "2f", "pos")])
state = {
"strength": args.strength,
"conv": args.convergence,
"swap": False,
"debug": False,
"fullscreen": mon is not None,
"windowed_geom": (100, 100, args.width, args.height),
}
def on_key(window, key, scancode, action, mods):
if action not in (glfw.PRESS, glfw.REPEAT):
return
if key == glfw.KEY_ESCAPE:
glfw.set_window_should_close(window, True)
elif key == glfw.KEY_UP:
state["strength"] = min(0.1, state["strength"] + 0.0025)
elif key == glfw.KEY_DOWN:
state["strength"] = max(0.0, state["strength"] - 0.0025)
elif key == glfw.KEY_RIGHT:
state["conv"] = min(1.0, state["conv"] + 0.02)
elif key == glfw.KEY_LEFT:
state["conv"] = max(0.0, state["conv"] - 0.02)
elif key == glfw.KEY_X and action == glfw.PRESS:
state["swap"] = not state["swap"]
elif key == glfw.KEY_D and action == glfw.PRESS:
state["debug"] = not state["debug"]
elif key == glfw.KEY_F and action == glfw.PRESS:
if state["fullscreen"]:
x, y, ww, hh = state["windowed_geom"]
glfw.set_window_monitor(window, None, x, y, ww, hh, 0)
else:
x, y = glfw.get_window_pos(window)
ww, hh = glfw.get_window_size(window)
state["windowed_geom"] = (x, y, ww, hh)
m = glfw.get_primary_monitor()
vm = glfw.get_video_mode(m)
glfw.set_window_monitor(window, m, 0, 0, vm.size.width, vm.size.height,
vm.refresh_rate)
state["fullscreen"] = not state["fullscreen"]
glfw.set_key_callback(win, on_key)
color_tex = None
color_size = None
depth_tex = make_tex(ctx, 2, 2, 1, "f4", np.full((2, 2), 0.5, dtype="f4").tobytes())
depth_size = (2, 2)
last_color_seq = -1
last_depth_seq = -1
last_title = 0.0
prog["color_tex"].value = 0
prog["depth_tex"].value = 1
prog["steps"].value = args.steps
try:
while not glfw.window_should_close(win):
glfw.poll_events()
frame, cseq = grabber.latest()
if frame is not None and cseq != last_color_seq:
last_color_seq = cseq
fh, fw = frame.shape[:2]
if color_size != (fw, fh):
if color_tex is not None:
color_tex.release()
color_tex = make_tex(ctx, fw, fh, 4, "f1")
color_tex.swizzle = "BGRA"
color_size = (fw, fh)
color_tex.write(np.ascontiguousarray(frame))
depth, dseq = worker.latest()
if depth is not None and dseq != last_depth_seq:
last_depth_seq = dseq
dh, dw = depth.shape
if depth_size != (dw, dh):
depth_tex.release()
depth_tex = make_tex(ctx, dw, dh, 1, "f4")
depth_size = (dw, dh)
depth_tex.write(depth.tobytes())
fbw, fbh = glfw.get_framebuffer_size(win)
ctx.viewport = (0, 0, fbw, fbh)
ctx.clear(0.0, 0.0, 0.0)
if color_tex is not None:
color_tex.use(0)
depth_tex.use(1)
prog["strength"].value = state["strength"]
prog["convergence"].value = state["conv"]
prog["show_depth"].value = 1 if state["debug"] else 0
half = fbw // 2
sign = -1.0 if state["swap"] else 1.0
for i, eye in enumerate((-1.0 * sign, 1.0 * sign)):
ctx.viewport = (i * half, 0, half, fbh)
prog["eye"].value = eye
vao.render(moderngl.TRIANGLE_STRIP)
glfw.swap_buffers(win)
now = time.time()
if now - last_title > 1.0:
last_title = now
glfw.set_window_title(
win,
f"depth3d SBS | depth {worker.fps:.0f} fps | "
f"strength {state['strength']:.3f} | conv {state['conv']:.2f}",
)
finally:
worker.stop_evt.set()
grabber.stop()
glfw.terminate()
if __name__ == "__main__":
main()
Reference in new issue View Git Blame Copy Permalink