mirror of
https://github.com/ultravideo/uvg266.git
synced 2024-11-23 18:14:06 +00:00
89 lines
3.7 KiB
Python
89 lines
3.7 KiB
Python
import sys
|
|
from pathlib import Path
|
|
from pprint import pprint
|
|
|
|
|
|
def get_ctx_count_and_names(cabac_source_file: Path):
|
|
count = 0
|
|
names = []
|
|
with open(cabac_source_file) as file:
|
|
for line in file:
|
|
line = line.strip().split(";")[0]
|
|
if line.startswith("cabac_ctx_t") and not "*" in line:
|
|
temp = 1
|
|
name = line.split()[1].split("[")[0]
|
|
for k in line.split("[")[1:]:
|
|
temp *= int(k[:-1])
|
|
count += temp
|
|
if temp == 1:
|
|
names.append(name)
|
|
else:
|
|
names.extend([f"{name}[{x}]" for x in range(temp)])
|
|
return count, names
|
|
|
|
|
|
def main(state_file: Path, ctx_names: list, ctx_count: int = 332, ctx_size: int = 6):
|
|
ctx_store = dict()
|
|
e_store = set()
|
|
was_zero_last = False
|
|
frame_num = -1
|
|
with open(state_file, "rb") as file:
|
|
try:
|
|
while True:
|
|
type_, x, y, depth, tree_type = file.read(23).decode().split()
|
|
# Reset stored data at the beginning of the frame
|
|
if x == '0' and y == '0' and type_ == "S" and tree_type != "2":
|
|
if not was_zero_last:
|
|
frame_num += 1
|
|
ctx_store = dict()
|
|
e_store = set()
|
|
was_zero_last = True
|
|
elif int(x) >= 64 and int(y) >= 64:
|
|
was_zero_last = False
|
|
|
|
ctx = file.read(ctx_count * ctx_size)
|
|
if type_ == "S":
|
|
# These shouldn't happen but just to make sure everything is working as intended
|
|
if ctx_store.get((x, y, depth, tree_type)):
|
|
raise RuntimeError
|
|
ctx_store[(x, y, depth, tree_type)] = ctx
|
|
else:
|
|
if (x, y, depth, tree_type) in e_store:
|
|
raise RuntimeError
|
|
e_store.add((x, y, depth, tree_type))
|
|
if (s_ctx := ctx_store[(x, y, depth, tree_type)]) != ctx:
|
|
actual_problem = False
|
|
|
|
for i in range(ctx_count):
|
|
temp_s = s_ctx[i * ctx_size: (i + 1) * ctx_size]
|
|
temp_e = ctx[i * ctx_size: (i + 1) * ctx_size]
|
|
if temp_s != temp_e:
|
|
if ctx_names[i] in ignore_list:
|
|
continue
|
|
actual_problem = True
|
|
print(f"MISSMATCH in {ctx_names[i]} {frame_num=} {x=} {y=} {depth=} {tree_type=}")
|
|
print(
|
|
f"GOT : {int.from_bytes(temp_s[0:2], 'little')}:"
|
|
f"{int.from_bytes(temp_s[2:4], 'little')} "
|
|
f"rate={int.from_bytes(temp_s[4:5], 'big')}")
|
|
print(
|
|
f"EXPECTED: {int.from_bytes(temp_e[0:2], 'little')}:"
|
|
f"{int.from_bytes(temp_e[2:4], 'little')} "
|
|
f"rate={int.from_bytes(temp_e[4:5], 'big')}")
|
|
if actual_problem:
|
|
exit(1)
|
|
except ValueError:
|
|
# EOF
|
|
pass
|
|
|
|
|
|
if __name__ == '__main__':
|
|
ignore_list = {"sao_type_idx_model", "sao_merge_flag_model"}
|
|
if len(sys.argv) < 2:
|
|
print("Usage: name of the file storing the cabac states")
|
|
exit(1)
|
|
path = Path(__file__) / "../.." / "src" / "cabac.h"
|
|
print(path.resolve())
|
|
counts, names_ = get_ctx_count_and_names(path.resolve())
|
|
main(Path(sys.argv[1]), names_, counts)
|