import dataclasses, enum, argparse, os, itertools, time, ctypes from typing import Any from tinygrad import Tensor, dtypes, Device, TinyJit from tinygrad.helpers import DEBUG, round_up, ceildiv, Timing, prod from tinygrad.runtime.autogen import avcodec, nv_570 as nv_gpu class BitReader: def __init__(self, data:bytes): self.reader, self.current_bits, self.bits, self.read_bits, self.total = iter(data), 0, 0, 0, len(data) * 8 def empty(self): return self.read_bits == self.total and self.current_bits == 0 def peak_bits(self, n): while self.current_bits < n: self.bits = (self.bits << 8) | next(self.reader) self.current_bits += 8 self.read_bits += 8 return (self.bits >> (self.current_bits - n)) & ((1 << n) - 1) def _next_bits(self, n): val = self.peak_bits(n) self.bits &= (1 << (self.current_bits - n)) - 1 self.current_bits -= n return val def u(self, n): return self._next_bits(n) # 9.2 Parsing process for 0-th order Exp-Golomb codes def ue_v(self): leading_zero_bits = -1 while True: bit = self.u(1) leading_zero_bits += 1 if bit == 1: break part = self.u(leading_zero_bits) if leading_zero_bits == 0: return 0 return (1 << leading_zero_bits) - 1 + part # 9.2.2 Mapping process for signed Exp-Golomb codes def se_v(self): k = self.ue_v() return (-1 ** (k + 1)) * (k // 2) # 7.3.1.1 General NAL unit syntax def _hevc_get_rbsp(dat:bytes, off=0) -> bytes: rbsp = bytes() while off < len(dat): if off + 2 < len(dat) and dat[off:off+3] == b'\x00\x00\x03': rbsp += bytes([0, 0]) off += 3 else: rbsp += bytes([dat[off]]) off += 1 return rbsp class HevcSlice: # 7.3.3 Profile, tier and level syntax def profile_tier_level(self, r:BitReader, enable:bool, max_sub_layers:int): assert enable and max_sub_layers == 0, "no sublayers supported" self._notimpl_profile_tier_level = r.u(88) self.general_level_idc = r.u(8) # 7.3.7 Short-term reference picture set syntax def st_ref_pic_set(self, r:BitReader, stRpsIdx:int, num_short_term_ref_pic_sets:int=0, sps=None): inter_ref_pic_set_prediction_flag = r.u(1) if stRpsIdx != 0 else 0 if inter_ref_pic_set_prediction_flag: if stRpsIdx == num_short_term_ref_pic_sets: delta_idx_minus1 = r.ue_v() delta_rps_sign = r.u(1) abs_delta_rps_minus1 = r.ue_v() NumDeltaPocs = sps.num_negative_pics + sps.num_positive_pics for i in range(NumDeltaPocs + 1): used_by_curr_pic_flag = r.u(1) if not used_by_curr_pic_flag: use_delta_flag = r.u(1) else: self.num_negative_pics = r.ue_v() self.num_positive_pics = r.ue_v() for i in range(self.num_negative_pics): delta_poc_s0_minus1 = r.ue_v() used_by_curr_pic_s0_flag = r.u(1) for i in range(self.num_positive_pics): delta_poc_s1_minus1 = r.ue_v() used_by_curr_pic_s1_flag = r.u(1) # 7.3.2.2 Sequence parameter set RBSP syntax class SPS(HevcSlice): def __init__(self, r:BitReader): self.sps_video_parameter_set_id = r.u(4) self.sps_max_sub_layers_minus1 = r.u(3) self.sps_temporal_id_nesting_flag = r.u(1) self.profile_tier_level(r, True, self.sps_max_sub_layers_minus1) self.sps_seq_parameter_set_id = r.ue_v() self.chroma_format_idc = r.ue_v() self.separate_colour_plane_flag = r.u(1) if self.chroma_format_idc == 3 else 0 self.pic_width_in_luma_samples = r.ue_v() self.pic_height_in_luma_samples = r.ue_v() self.conformance_window_flag = r.u(1) if self.conformance_window_flag: self.conf_win_left_offset = r.ue_v() self.conf_win_right_offset = r.ue_v() self.conf_win_top_offset = r.ue_v() self.conf_win_bottom_offset = r.ue_v() else: self.conf_win_left_offset = self.conf_win_right_offset = self.conf_win_top_offset = self.conf_win_bottom_offset = 0 self.bit_depth_luma = r.ue_v() + 8 self.bit_depth_chroma = r.ue_v() + 8 self.log2_max_pic_order_cnt_lsb_minus4 = r.ue_v() self.sps_sub_layer_ordering_info_present_flag = r.u(1) self.sps_max_dec_pic_buffering, self.sps_max_num_reorder_pics, self.sps_max_latency_increase_plus1 = [], [], [] for i in range((0 if self.sps_sub_layer_ordering_info_present_flag else self.sps_max_sub_layers_minus1), self.sps_max_sub_layers_minus1 + 1): self.sps_max_dec_pic_buffering.append(r.ue_v() + 1) self.sps_max_num_reorder_pics.append(r.ue_v()) self.sps_max_latency_increase_plus1.append(r.ue_v()) self.log2_min_luma_coding_block_size = r.ue_v() + 3 self.log2_max_luma_coding_block_size = self.log2_min_luma_coding_block_size + r.ue_v() self.log2_min_transform_block_size = r.ue_v() + 2 self.log2_max_transform_block_size = self.log2_min_transform_block_size + r.ue_v() self.max_transform_hierarchy_depth_inter = r.ue_v() self.max_transform_hierarchy_depth_intra = r.ue_v() if scaling_list_enabled_flag := r.u(1): if sps_scaling_list_data_present_flag := r.u(1): assert False, "scaling_list_data parsing not implemented" self.amp_enabled_flag = r.u(1) self.sample_adaptive_offset_enabled_flag = r.u(1) self.pcm_enabled_flag = r.u(1) assert self.pcm_enabled_flag == 0, "pcm not implemented" self.num_short_term_ref_pic_sets = r.ue_v() for i in range(self.num_short_term_ref_pic_sets): self.st_ref_pic_set(r, i, self.num_short_term_ref_pic_sets) self.long_term_ref_pics_present_flag = r.u(1) if self.long_term_ref_pics_present_flag: assert False, "long_term_ref_pics parsing not implemented" self.sps_temporal_mvp_enabled_flag = r.u(1) self.strong_intra_smoothing_enabled_flag = r.u(1) # 7.3.2.3 Picture parameter set RBSP syntax class PPS(HevcSlice): def __init__(self, r:BitReader): self.pps_pic_parameter_set_id = r.ue_v() self.pps_seq_parameter_set_id = r.ue_v() self.dependent_slice_segments_enabled_flag = r.u(1) self.output_flag_present_flag = r.u(1) self.num_extra_slice_header_bits = r.u(3) self.sign_data_hiding_enabled_flag = r.u(1) self.cabac_init_present_flag = r.u(1) self.num_ref_idx_l0_default_active = r.ue_v() + 1 self.num_ref_idx_l1_default_active = r.ue_v() + 1 self.init_qp = r.se_v() + 26 self.constrained_intra_pred_flag = r.u(1) self.transform_skip_enabled_flag = r.u(1) self.cu_qp_delta_enabled_flag = r.u(1) if self.cu_qp_delta_enabled_flag: self.diff_cu_qp_delta_depth = r.ue_v() self.pps_cb_qp_offset = r.se_v() self.pps_cr_qp_offset = r.se_v() self.pps_slice_chroma_qp_offsets_present_flag = r.u(1) self.weighted_pred_flag = r.u(1) self.weighted_bipred_flag = r.u(1) self.transquant_bypass_enabled_flag = r.u(1) self.tiles_enabled_flag = r.u(1) self.entropy_coding_sync_enabled_flag = r.u(1) if self.tiles_enabled_flag: self.num_tile_columns_minus1 = r.ue_v() self.num_tile_rows_minus1 = r.ue_v() self.uniform_spacing_flag = r.u(1) self.column_width_minus1, self.row_height_minus1 = [], [] if not self.uniform_spacing_flag: for i in range(self.num_tile_columns_minus1): self.column_width_minus1.append(r.ue_v()) for i in range(self.num_tile_rows_minus1): self.row_height_minus1.append(r.ue_v()) self.loop_filter_across_tiles_enabled_flag = r.u(1) self.loop_filter_across_slices_enabled_flag = r.u(1) self.deblocking_filter_control_present_flag = r.u(1) if self.deblocking_filter_control_present_flag: assert False, "deblocking_filter parsing not implemented" self.scaling_list_data_present_flag = r.u(1) if self.scaling_list_data_present_flag: assert False, "scaling_list_data parsing not implemented" self.lists_modification_present_flag = r.u(1) self.log2_parallel_merge_level = r.ue_v() + 2 # 7.3.6 Slice segment header syntax class SliceSegment(HevcSlice): def __init__(self, r:BitReader, nal_unit_type:int, sps:SPS, pps:PPS): self.first_slice_segment_in_pic_flag = r.u(1) if nal_unit_type >= avcodec.HEVC_NAL_BLA_W_LP and nal_unit_type <= avcodec.HEVC_NAL_RSV_IRAP_VCL23: self.no_output_of_prior_pics_flag = r.u(1) self.slice_pic_parameter_set_id = r.ue_v() if not self.first_slice_segment_in_pic_flag: if pps.dependent_slice_segments_enabled_flag: self.dependent_slice_segment_flag = r.u(1) self.slice_segment_address = r.ue_v() self.dependent_slice_segment_flag = 0 if not self.dependent_slice_segment_flag: r.u(pps.num_extra_slice_header_bits) # extra bits ignored self.slice_type = r.ue_v() self.sw_skip_start = r.read_bits - r.current_bits self.pic_output_flag = r.u(1) if pps.output_flag_present_flag else 0 self.colour_plane_id = r.u(2) if sps.separate_colour_plane_flag else 0 if nal_unit_type != avcodec.HEVC_NAL_IDR_W_RADL and nal_unit_type != avcodec.HEVC_NAL_IDR_N_LP: self.slice_pic_order_cnt_lsb = r.u(sps.log2_max_pic_order_cnt_lsb_minus4 + 4) self.short_term_ref_pic_set_sps_flag = r.u(1) if not self.short_term_ref_pic_set_sps_flag: self.short_term_ref_pics_in_slice_start = r.read_bits - r.current_bits self.st_ref_pic_set(r, sps.num_short_term_ref_pic_sets, sps=sps) self.short_term_ref_pics_in_slice_end = r.read_bits - r.current_bits elif sps.num_short_term_ref_pic_sets > 1: assert False, "short_term_ref_pic_set parsing not implemented" if sps.long_term_ref_pics_present_flag: assert False, "long_term_ref_pics parsing not implemented" self.sw_skip_end = r.read_bits - r.current_bits self.slice_temporal_mvp_enabled_flag = r.u(1) if sps.sps_temporal_mvp_enabled_flag else 0 else: self.slice_pic_order_cnt_lsb, self.sw_skip_end = 0, self.sw_skip_start if sps.sample_adaptive_offset_enabled_flag: slice_sao_luma_flag = r.u(1) ChromaArrayType = sps.chroma_format_idc if sps.separate_colour_plane_flag == 0 else 0 slice_sao_chroma_flag = r.u(1) if ChromaArrayType != 0 else 0 if self.slice_type in {avcodec.HEVC_SLICE_B, avcodec.HEVC_SLICE_B}: if num_ref_idx_active_override_flag := r.u(1): num_ref_idx_l0_active_minus1 = r.ue_v() num_ref_idx_l1_active_minus1 = r.ue_v() if self.slice_type == avcodec.HEVC_SLICE_B else 0 def fill_sps_into_dev_context(device_ctx, sps:SPS): device_ctx.chroma_format_idc = sps.chroma_format_idc device_ctx.pic_width_in_luma_samples = sps.pic_width_in_luma_samples device_ctx.pic_height_in_luma_samples = sps.pic_height_in_luma_samples device_ctx.bit_depth_luma = sps.bit_depth_luma device_ctx.bit_depth_chroma = sps.bit_depth_chroma device_ctx.log2_max_pic_order_cnt_lsb_minus4 = sps.log2_max_pic_order_cnt_lsb_minus4 device_ctx.log2_min_luma_coding_block_size = sps.log2_min_luma_coding_block_size device_ctx.log2_max_luma_coding_block_size = sps.log2_max_luma_coding_block_size device_ctx.log2_min_transform_block_size = sps.log2_min_transform_block_size device_ctx.log2_max_transform_block_size = sps.log2_max_transform_block_size device_ctx.amp_enabled_flag = sps.amp_enabled_flag device_ctx.pcm_enabled_flag = sps.pcm_enabled_flag device_ctx.sample_adaptive_offset_enabled_flag = sps.sample_adaptive_offset_enabled_flag device_ctx.sps_temporal_mvp_enabled_flag = sps.sps_temporal_mvp_enabled_flag device_ctx.strong_intra_smoothing_enabled_flag = sps.strong_intra_smoothing_enabled_flag def fill_pps_into_dev_context(device_ctx, pps:PPS): device_ctx.sign_data_hiding_enabled_flag = pps.sign_data_hiding_enabled_flag device_ctx.cabac_init_present_flag = pps.cabac_init_present_flag device_ctx.num_ref_idx_l0_default_active = pps.num_ref_idx_l0_default_active device_ctx.num_ref_idx_l1_default_active = pps.num_ref_idx_l1_default_active device_ctx.init_qp = pps.init_qp device_ctx.cu_qp_delta_enabled_flag = pps.cu_qp_delta_enabled_flag device_ctx.diff_cu_qp_delta_depth = getattr(pps, 'diff_cu_qp_delta_depth', 0) device_ctx.pps_cb_qp_offset = pps.pps_cb_qp_offset device_ctx.pps_cr_qp_offset = pps.pps_cr_qp_offset device_ctx.pps_slice_chroma_qp_offsets_present_flag = pps.pps_slice_chroma_qp_offsets_present_flag device_ctx.weighted_pred_flag = pps.weighted_pred_flag device_ctx.weighted_bipred_flag = pps.weighted_bipred_flag device_ctx.transquant_bypass_enabled_flag = pps.transquant_bypass_enabled_flag device_ctx.tiles_enabled_flag = pps.tiles_enabled_flag device_ctx.entropy_coding_sync_enabled_flag = pps.entropy_coding_sync_enabled_flag device_ctx.loop_filter_across_slices_enabled_flag = pps.loop_filter_across_slices_enabled_flag device_ctx.deblocking_filter_control_present_flag = pps.deblocking_filter_control_present_flag device_ctx.scaling_list_data_present_flag = pps.scaling_list_data_present_flag device_ctx.lists_modification_present_flag = pps.lists_modification_present_flag device_ctx.log2_parallel_merge_level = pps.log2_parallel_merge_level device_ctx.loop_filter_across_tiles_enabled_flag = getattr(pps, 'loop_filter_across_tiles_enabled_flag', 0) def parse_hevc_file_headers(dat:bytes, device="NV"): res = [] nal_unit_start = 1 history:list[tuple[int, int, int]] = [] device_ctx = nv_gpu.nvdec_hevc_pic_s(gptimer_timeout_value=92720000, tileformat=1, sw_start_code_e=1, pattern_id=2) nal_infos = [] ctx_bytes = bytes() align_ctx_bytes_size = 0x300 def _flush_picture(): nonlocal res, history, device_ctx, nal_infos, ctx_bytes, align_ctx_bytes_size if not len(nal_infos): return hdr, nal_unit_type = nal_infos[0][0] assert all(nal_unit_type == x[0][1] for x in nal_infos), "all NAL units in a picture must be of the same type" device_ctx.curr_pic_idx = next(i for i in range(16) if all(d[0] != i for d in history)) if nal_unit_type in {avcodec.HEVC_NAL_IDR_W_RADL, avcodec.HEVC_NAL_IDR_N_LP}: history = [] device_ctx.num_ref_frames = len(history) device_ctx.IDR_picture_flag = int(nal_unit_type in {avcodec.HEVC_NAL_IDR_W_RADL, avcodec.HEVC_NAL_IDR_N_LP}) device_ctx.RAP_picture_flag = int(nal_unit_type >= avcodec.HEVC_NAL_BLA_W_LP and nal_unit_type <= avcodec.HEVC_NAL_RSV_IRAP_VCL23) device_ctx.RefDiffPicOrderCnts=(ctypes.c_int16 * 16)() device_ctx.colMvBuffersize = (round_up(sps.pic_width_in_luma_samples, 64) * round_up(sps.pic_height_in_luma_samples, 64) // 16) // 256 device_ctx.framestride=(ctypes.c_uint32 * 2)(round_up(sps.pic_width_in_luma_samples, 64), round_up(sps.pic_width_in_luma_samples, 64)) device_ctx.sw_hdr_skip_length = hdr.sw_skip_end - hdr.sw_skip_start device_ctx.num_bits_short_term_ref_pics_in_slice = max(0, device_ctx.sw_hdr_skip_length - 9) device_ctx.stream_len = sum(x[2] for x in nal_infos) if pps.tiles_enabled_flag: device_ctx.num_tile_columns = pps.num_tile_columns_minus1 + 1 device_ctx.num_tile_rows = pps.num_tile_rows_minus1 + 1 device_ctx.num_short_term_ref_pic_sets = sps.num_short_term_ref_pic_sets luma_h_rounded = round_up(sps.pic_height_in_luma_samples, 64) device_ctx.HevcSaoBufferOffset = (608 * luma_h_rounded) >> 8 device_ctx.HevcBsdCtrlOffset = ((device_ctx.HevcSaoBufferOffset<<8) + 4864 * luma_h_rounded) >> 8 device_ctx.v1.hevc_main10_444_ext.HevcFltAboveOffset = ((device_ctx.HevcBsdCtrlOffset<<8) + 152 * luma_h_rounded) >> 8 device_ctx.v1.hevc_main10_444_ext.HevcSaoAboveOffset = ((device_ctx.v1.hevc_main10_444_ext.HevcFltAboveOffset<<8) + 2000 * luma_h_rounded) >> 8 device_ctx.v3.HevcSliceEdgeOffset = device_ctx.v1.hevc_main10_444_ext.HevcSaoAboveOffset before_list, after_list = [], [] for pic_idx, poc, _ in history: device_ctx.RefDiffPicOrderCnts[pic_idx] = hdr.slice_pic_order_cnt_lsb - poc if hdr.slice_pic_order_cnt_lsb < poc: after_list.append((poc - hdr.slice_pic_order_cnt_lsb, pic_idx)) else: before_list.append((hdr.slice_pic_order_cnt_lsb - poc, pic_idx)) before_list.sort() after_list.sort() device_ctx.initreflistidxl0 = (ctypes.c_uint8 * 16)(*[idx for _,idx in before_list + after_list]) if hdr.slice_type == avcodec.HEVC_SLICE_B: device_ctx.initreflistidxl1 = (ctypes.c_uint8 * 16)(*[idx for _,idx in after_list + before_list]) locl_ctx_bytes = bytes(device_ctx) locl_ctx_bytes += b'\x00\x00\x00\x00\x00\x00\x00\x00\x10\x00\x00\x00' # blackwell extension locl_ctx_bytes += bytes(0x200 - len(locl_ctx_bytes)) # pad to 512 bytes pic_width_in_ctbs = ceildiv(sps.pic_width_in_luma_samples, (1 << sps.log2_max_luma_coding_block_size)) pic_height_in_ctbs = ceildiv(sps.pic_height_in_luma_samples, (1 << sps.log2_max_luma_coding_block_size)) # append tile sizes 0x200 if pps.tiles_enabled_flag and pps.uniform_spacing_flag: assert device_ctx.num_tile_columns == 1 and device_ctx.num_tile_rows == 1, "not implemented: uniform spacing with multiple tiles" locl_ctx_bytes += pic_width_in_ctbs.to_bytes(2, "little") + pic_height_in_ctbs.to_bytes(2, "little") else: if pps.tiles_enabled_flag and not getattr(pps, 'uniform_spacing_flag', 0): column_width = [cw_minus1 + 1 for cw_minus1 in pps.column_width_minus1[0:pps.num_tile_columns_minus1]] row_height = [rh_minus1 + 1 for rh_minus1 in pps.row_height_minus1[0:pps.num_tile_rows_minus1]] else: column_width = [] row_height = [] column_width.append(pic_width_in_ctbs - sum(column_width)) row_height.append(pic_height_in_ctbs - sum(row_height)) for c in column_width: for r in row_height: locl_ctx_bytes += c.to_bytes(2, "little") + r.to_bytes(2, "little") luma_size = round_up(sps.pic_width_in_luma_samples, 64) * round_up(sps.pic_height_in_luma_samples, 64) chroma_size = round_up(sps.pic_width_in_luma_samples, 64) * round_up((sps.pic_height_in_luma_samples + 1) // 2, 64) is_hist = nal_unit_type in {avcodec.HEVC_NAL_TRAIL_R, avcodec.HEVC_NAL_IDR_N_LP, avcodec.HEVC_NAL_IDR_W_RADL} res.append((nal_infos[0][1], device_ctx.stream_len, device_ctx.curr_pic_idx, len(history), is_hist)) locl_ctx_bytes += (align_ctx_bytes_size - len(locl_ctx_bytes)) * b'\x00' ctx_bytes += locl_ctx_bytes if nal_unit_type in {avcodec.HEVC_NAL_TRAIL_R, avcodec.HEVC_NAL_IDR_N_LP, avcodec.HEVC_NAL_IDR_W_RADL}: history.append((device_ctx.curr_pic_idx, hdr.slice_pic_order_cnt_lsb, None)) if len(history) >= sps.sps_max_dec_pic_buffering[0]: # remove the oldest poc history.pop(0) nal_infos = [] cnt = 0 while nal_unit_start < len(dat): assert dat[nal_unit_start:nal_unit_start+3] == b"\x00\x00\x01", "NAL unit start code not found" pos = dat.find(b"\x00\x00\x01", nal_unit_start + 3) nal_unit_len = (pos if pos != -1 else len(dat)) - nal_unit_start # 7.3.1.1 General NAL unit syntax nal_unit_type = (dat[nal_unit_start+3] >> 1) & 0x3F slice_dat = dat[nal_unit_start+5:nal_unit_start+nal_unit_len] if nal_unit_type == avcodec.HEVC_NAL_SPS: sps = SPS(BitReader(_hevc_get_rbsp(slice_dat))) fill_sps_into_dev_context(device_ctx, sps) elif nal_unit_type == avcodec.HEVC_NAL_PPS: pps = PPS(BitReader(_hevc_get_rbsp(slice_dat))) fill_pps_into_dev_context(device_ctx, pps) elif nal_unit_type in {avcodec.HEVC_NAL_IDR_N_LP, avcodec.HEVC_NAL_IDR_W_RADL, avcodec.HEVC_NAL_TRAIL_R, avcodec.HEVC_NAL_TRAIL_N}: hdr = SliceSegment(BitReader(slice_dat), nal_unit_type, sps, pps) if hdr.first_slice_segment_in_pic_flag == 1: _flush_picture() nal_infos.append(((hdr, nal_unit_type), nal_unit_start, nal_unit_len)) nal_unit_start += nal_unit_len _flush_picture() w = sps.pic_width_in_luma_samples - 2 * (sps.conf_win_left_offset + sps.conf_win_right_offset) h = sps.pic_height_in_luma_samples - 2 * (sps.conf_win_top_offset + sps.conf_win_bottom_offset) chroma_off = round_up(sps.pic_width_in_luma_samples, 64) * round_up(sps.pic_height_in_luma_samples, 64) opaque = Tensor(ctx_bytes, device=device).reshape(len(res), align_ctx_bytes_size) return opaque, res, w, h, sps.pic_width_in_luma_samples, sps.pic_height_in_luma_samples, chroma_off def _addr_table(h, w, w_aligned): GOB_W, GOB_H = 64, 8 GOB_SIZE = GOB_W * GOB_H BLOCK_H_GOBS = 2 xs = Tensor.arange(w, dtype=dtypes.uint32).reshape(1, w) ys = Tensor.arange(h, dtype=dtypes.uint32).reshape(h, 1) gob_x = xs // GOB_W gob_y = ys // GOB_H super_block_y = gob_y // BLOCK_H_GOBS gob_y_in_block = gob_y % BLOCK_H_GOBS stride_gobs = w_aligned // GOB_W base = ((super_block_y * stride_gobs + gob_x) * BLOCK_H_GOBS + gob_y_in_block) * GOB_SIZE lx, ly = xs % GOB_W, ys % GOB_H swiz = (lx & 0x0F) | ((ly & 0x03) << 4) | ((lx & 0x10) << 2) | ((ly & 0x04) << 5) | ((lx & 0x20) << 3) return (base + swiz).reshape(-1) def nv12_to_bgr_from_planes(luma: Tensor, chroma: Tensor, h: int, w: int) -> Tensor: Y = luma.reshape(h, w).cast(dtypes.float32) uv = chroma.reshape(h // 2, w // 2, 2).cast(dtypes.float32) U_small = uv[..., 0] V_small = uv[..., 1] U = U_small.reshape(h // 2, 1, w // 2, 1).expand(h // 2, 2, w // 2, 2).reshape(h, w) V = V_small.reshape(h // 2, 1, w // 2, 1).expand(h // 2, 2, w // 2, 2).reshape(h, w) C = Y - 16.0 D = U - 128.0 E = V - 128.0 R = 1.1643835616438356 * C + 1.5960267857142858 * E G = 1.1643835616438356 * C - 0.39176229009491365 * D - 0.8129676472377708 * E B = 1.1643835616438356 * C + 2.017232142857143 * D R = R.maximum(0.0).minimum(255.0) G = G.maximum(0.0).minimum(255.0) B = B.maximum(0.0).minimum(255.0) return Tensor.stack([B, G, R], dim=2).cast(dtypes.uint8) def untile_nv12(src:Tensor, h:int, w:int, luma_w:int, chroma_off:int) -> Tensor: luma = src.reshape(-1)[_addr_table(h, w, round_up(luma_w, 64))] chroma = src.reshape(-1)[chroma_off:][_addr_table((h + 1) // 2, w, round_up(luma_w, 64))] return luma.cat(chroma).realize() def to_bgr(tensor:Tensor, h:int, w:int, luma_w:int, chroma_off:int) -> Tensor: luma = tensor.reshape(-1)[_addr_table(h, w, round_up(luma_w, 64))] chroma = tensor.reshape(-1)[chroma_off:][_addr_table((h + 1) // 2, w, round_up(luma_w, 64))] return nv12_to_bgr_from_planes(luma, chroma, h, w).realize()