modeld: use const border mode for dm warp (#37986)

This commit is contained in:
ZwX1616
2026-05-13 16:52:08 -07:00
committed by GitHub
parent b88adeb4ad
commit f24ad7e27a
2 changed files with 13 additions and 5 deletions
+1 -1
View File
@@ -17,7 +17,7 @@ def make_warp_dm(nv12: NV12Frame, dm_w, dm_h):
def warp_dm(input_frame, M_inv):
M_inv = M_inv.to(Device.DEFAULT).realize()
return warp_perspective_tinygrad(input_frame[:cam_h*stride], M_inv,
(dm_w, dm_h), (cam_h, cam_w), stride_pad).reshape(-1, dm_h * dm_w)
(dm_w, dm_h), (cam_h, cam_w), stride_pad, border_fill_val=0).reshape(-1, dm_h * dm_w)
return warp_dm
+12 -4
View File
@@ -19,7 +19,7 @@ UV_SCALE_MATRIX = np.array([[0.5, 0, 0], [0, 0.5, 0], [0, 0, 1]], dtype=np.float
UV_SCALE_MATRIX_INV = np.linalg.inv(UV_SCALE_MATRIX)
def warp_perspective_tinygrad(src_flat, M_inv, dst_shape, src_shape, stride_pad):
def warp_perspective_tinygrad(src_flat, M_inv, dst_shape, src_shape, stride_pad, border_fill_val=None):
w_dst, h_dst = dst_shape
h_src, w_src = src_shape
@@ -34,11 +34,19 @@ def warp_perspective_tinygrad(src_flat, M_inv, dst_shape, src_shape, stride_pad)
src_x = src_x / src_w
src_y = src_y / src_w
x_nn_clipped = Tensor.round(src_x).clip(0, w_src - 1).cast('int')
y_nn_clipped = Tensor.round(src_y).clip(0, h_src - 1).cast('int')
x_round = Tensor.round(src_x)
y_round = Tensor.round(src_y)
x_nn_clipped = x_round.clip(0, w_src - 1).cast('int')
y_nn_clipped = y_round.clip(0, h_src - 1).cast('int')
idx = y_nn_clipped * (w_src + stride_pad) + x_nn_clipped
sampled = src_flat[idx]
return src_flat[idx]
if border_fill_val is None:
return sampled
in_bounds = ((x_round >= 0) & (x_round <= w_src - 1) &
(y_round >= 0) & (y_round <= h_src - 1)).cast(sampled.dtype)
return sampled * in_bounds + Tensor(border_fill_val, dtype=sampled.dtype) * (1 - in_bounds)
def frames_to_tensor(frames):