Files
coven/modules/ufbx/misc/check_dataset.py

227 lines
6.4 KiB
Python

import os
import json
from typing import NamedTuple, Optional, List
import subprocess
import glob
import re
import urllib.parse
class TestModel(NamedTuple):
fbx_path: str
obj_path: Optional[str]
mtl_path: Optional[str]
mat_path: Optional[str]
frame: Optional[int]
class TestCase(NamedTuple):
root: str
json_path: str
title: str
author: str
license: str
url: str
extra_files: List[str]
models: List[TestModel]
def log(message=""):
print(message, flush=True)
def single_file(path):
if os.path.exists(path):
return [path]
else:
return []
def strip_ext(path):
if path.endswith(".gz"):
path = path[:-3]
base, _ = os.path.splitext(path)
return base
def get_fbx_files(json_path):
base_path = strip_ext(json_path)
yield from single_file(f"{base_path}.fbx")
yield from single_file(f"{base_path}.ufbx.obj")
yield from glob.glob(f"{glob.escape(base_path)}/*.fbx")
def get_obj_files(fbx_path):
base_path = strip_ext(fbx_path)
yield from single_file(f"{base_path}.obj.gz")
yield from single_file(f"{base_path}.obj")
yield from glob.glob(f"{glob.escape(base_path)}_*.obj.gz")
yield from glob.glob(f"{glob.escape(base_path)}_*.obj")
def get_mtl_files(obj_path):
base_path = strip_ext(obj_path)
yield from single_file(f"{base_path}.mtl")
def get_mat_files(obj_path):
base_path = strip_ext(obj_path)
yield from single_file(f"{base_path}.mat")
def remove_duplicate_files(paths):
seen = set()
for path in paths:
base = strip_ext(path)
if base in seen: continue
seen.add(base)
yield path
def gather_case_models(json_path):
for fbx_path in get_fbx_files(json_path):
for obj_path in remove_duplicate_files(get_obj_files(fbx_path)):
mtl_path = next(get_mtl_files(obj_path), None)
mat_path = next(get_mat_files(fbx_path), None)
fbx_base = strip_ext(fbx_path)
obj_base = strip_ext(obj_path)
flags = obj_base[len(fbx_base):].split("_")
# Parse flags
frame = None
for flag in flags:
m = re.match(r"frame(\d+)", flag)
if m:
frame = int(m.group(1))
yield TestModel(
fbx_path=fbx_path,
obj_path=obj_path,
mtl_path=mtl_path,
mat_path=mat_path,
frame=frame)
else:
# TODO: Handle objless fbx
pass
def gather_dataset_tasks(root_dir):
for root, _, files in os.walk(root_dir):
for filename in files:
if not filename.endswith(".json"):
continue
path = os.path.join(root, filename)
with open(path, "rt", encoding="utf-8") as f:
desc = json.load(f)
models = list(gather_case_models(path))
if not models:
raise RuntimeError(f"No models found for {path}")
extra_files = [os.path.join(root, ex) for ex in desc.get("extra-files", [])]
yield TestCase(
root=root_dir,
json_path=path,
title=desc["title"],
author=desc["author"],
license=desc["license"],
url=desc["url"],
extra_files=extra_files,
models=models,
)
if __name__ == "__main__":
from argparse import ArgumentParser
parser = ArgumentParser("check_dataset.py --root <root>")
parser.add_argument("--root", help="Root directory to search for .json files")
parser.add_argument("--host-url", help="URL where the files are hosted")
parser.add_argument("--exe", help="check_fbx.c executable")
parser.add_argument("--verbose", action="store_true", help="Print verbose information")
argv = parser.parse_args()
cases = list(gather_dataset_tasks(root_dir=argv.root))
def fmt_url(path, root=""):
if root:
path = os.path.relpath(path, root)
path = path.replace("\\", "/")
safe_path = urllib.parse.quote(path)
return f"{argv.host_url}/{safe_path}"
def fmt_rel(path, root=""):
if root:
path = os.path.relpath(path, root)
path = path.replace("\\", "/")
return f"{path}"
ok_count = 0
test_count = 0
case_ok_count = 0
for case in cases:
log(f"== '{case.title}' by '{case.author}' ({case.license}) ==")
log()
if case.url:
log(f" source url: {case.url}")
log(f" .json url: {fmt_url(case.json_path, case.root)}")
for extra in case.extra_files:
log(f" extra url: {fmt_url(extra, case.root)}")
log()
case_ok = True
for model in case.models:
test_count += 1
args = [argv.exe]
args.append(model.fbx_path)
extra = []
if model.obj_path:
args += ["--obj", model.obj_path]
if model.mat_path:
args += ["--mat", model.mat_path]
if model.frame is not None:
extra.append(f"frame {model.frame}")
args += ["--frame", str(model.frame)]
name = fmt_rel(model.fbx_path, case.root)
extra_str = ""
if extra:
extra_str = " [" + ", ".join(extra) + "]"
log(f"-- {name}{extra_str} --")
log()
if argv.host_url:
log(f" .fbx url: {fmt_url(model.fbx_path, case.root)}")
if model.obj_path:
log(f" .obj url: {fmt_url(model.obj_path, case.root)}")
if model.mtl_path:
log(f" .mtl url: {fmt_url(model.mtl_path, case.root)}")
if model.mat_path:
log(f" .mat url: {fmt_url(model.mat_path, case.root)}")
log()
log("$ " + " ".join(args))
log()
try:
subprocess.check_call(args)
log()
log("-- PASS --")
ok_count += 1
except subprocess.CalledProcessError:
log()
log("-- FAIL --")
case_ok = False
log()
if case_ok:
case_ok_count += 1
log(f"{ok_count}/{test_count} files passed ({case_ok_count}/{len(cases)} test cases)")
if ok_count < test_count:
exit(1)