| """Checkpoint-compatible v61 fine tuning with pixel supervision and context regularization. |
| |
| Run through run_local.py. Real-image generalization is an experimental objective; |
| synthetic appearance stress tests do not establish real-image accuracy. |
| """ |
| import argparse |
| import hashlib |
| import inspect |
| import json |
| import math |
| import os |
| import shutil |
| import sys |
| import time |
| from contextlib import contextmanager |
| from pathlib import Path |
| from types import SimpleNamespace |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch.utils.data import DataLoader, Subset, random_split |
|
|
| from SkingToolkit.dense_uv_parser import train as core |
| from SkingToolkit.dense_uv_parser.model import DenseUVParserNet |
| from SkingToolkit.dense_uv_parser.losses import DenseUVParserLoss |
| from SkingToolkit.dense_uv_parser.semantic import attach_semantic_runtime |
| from SkingToolkit.dense_uv_parser.skin_dataset import SkinUVDataset |
| from SkingToolkit.dense_uv_parser.utils import IGNORE_INDEX |
| from SkingToolkit.dense_uv_parser.semantic_targets import ( |
| build_part_layer_masks, build_semantic_attribute_targets, |
| ) |
| from SkingToolkit.renderer import DifferentiableRenderer |
|
|
|
|
| def write_json(path, value): |
| path = Path(path) |
| temp = path.with_suffix(path.suffix + ".tmp") |
| temp.write_text(json.dumps(value, indent=2, allow_nan=False)) |
| os.replace(temp, path) |
|
|
|
|
| @torch.no_grad() |
| def appearance_augment(images, foreground, strength=1.0, generator=None): |
| """Change appearance without moving pixels, alpha, background, or GT labels.""" |
| rgb = images[:, :3].float() |
| fg = foreground.float() |
| n, _, h, w = rgb.shape |
| def rand(*shape): |
| return torch.rand(shape, device=rgb.device, generator=generator) |
| def randn(*shape): |
| return torch.randn(shape, device=rgb.device, generator=generator) |
| |
| smooth = F.avg_pool2d(rgb * fg, 3, 1, 1) / F.avg_pool2d(fg, 3, 1, 1).clamp_min(1e-6) |
| mix = rand(n, 1, 1, 1) * (0.55 * strength) |
| altered = rgb * (1 - mix) + smooth * mix |
| illumination = F.interpolate( |
| randn(n, 1, 8, 4), size=(h, w), mode="bilinear", align_corners=False |
| ) * (0.045 * strength) |
| contrast = 1 + (rand(n, 1, 1, 1) - 0.5) * (0.30 * strength) |
| tint = (rand(n, 3, 1, 1) - 0.5) * (0.08 * strength) |
| altered = (altered - 0.5) * contrast + 0.5 + illumination + tint |
| altered = altered + randn(n, 3, h, w) * (0.008 * strength) |
| result = images.clone() |
| result[:, :3] = torch.where(fg.bool(), altered.clamp(0, 1), rgb) |
| return result |
|
|
|
|
| def flat_foreground_mask(images, targets): |
| rgb = images[:, :3].float() |
| variance = (F.avg_pool2d(rgb.square(), 3, 1, 1) |
| - F.avg_pool2d(rgb, 3, 1, 1).square()).clamp_min(0) |
| interior = F.avg_pool2d(targets["foreground"].float(), 3, 1, 1) > 0.999 |
| return (variance.amax(1) < (2 / 255) ** 2) & interior[:, 0] |
|
|
|
|
| def flat_route_loss(logits, labels, flat): |
| per_pixel = F.cross_entropy(logits.float(), labels, ignore_index=IGNORE_INDEX, reduction="none") |
| terms = [per_pixel[flat & (labels == role)].mean() |
| for role in (0, 1) if (flat & (labels == role)).any()] |
| return torch.stack(terms).mean() if terms else logits.sum() * 0 |
|
|
|
|
| def head_outer_hard_positive_loss(logits, targets): |
| """Mine the worst head-outer pixels using actual renderer labels.""" |
| head_outer = (targets["route_role"] == 1) & (targets["part"] == 0) |
| if not head_outer.any(): |
| return logits.sum() * 0 |
| errors = -F.log_softmax(logits.float(), dim=1)[:, 1][head_outer] |
| return errors.topk(max(1, math.ceil(errors.numel() * 0.20))).values.mean() |
|
|
|
|
| def skip_regularizer(model, probability): |
| """Drop entire shallow skip inputs per image only during training.""" |
| def hook(module, inputs): |
| if not model.training or probability == 0: |
| return inputs |
| x, skip = inputs |
| keep = (torch.rand(skip.shape[0], 1, 1, 1, device=skip.device) >= probability) |
| return x, skip * keep.to(skip.dtype) |
| return [block.register_forward_pre_hook(hook) for block in (model.up0, model.up1)] |
|
|
|
|
| @contextmanager |
| def semantics_disabled(model): |
| handles = [] |
| if model.semantic_spatial_fusion is not None: |
| handles.append(model.semantic_spatial_fusion.register_forward_hook( |
| lambda module, inputs, output: torch.zeros_like(output))) |
| if model.semantic_fusion is not None: |
| handles.append(model.semantic_fusion.register_forward_hook( |
| lambda module, inputs, output: (torch.zeros_like(output[0]), output[1]))) |
| try: |
| yield |
| finally: |
| for handle in handles: |
| handle.remove() |
|
|
|
|
| def build_split(dataset, seed, val_split): |
| val_count = int(len(dataset) * val_split) |
| return random_split(dataset, [len(dataset) - val_count, val_count], |
| generator=torch.Generator().manual_seed(seed)) |
|
|
|
|
| def count_routes(output, target, flat): |
| labels = target["route_role"] |
| valid = labels != IGNORE_INDEX |
| pred = output["layer"].argmax(1) |
| counts = {} |
| for role, name in ((0, "inner"), (1, "outer"), (2, "secondary")): |
| gt, selected = labels == role, (pred == role) & valid |
| counts[f"{name}_tp"] = int((gt & selected).sum()) |
| counts[f"{name}_fp"] = int((~gt & selected).sum()) |
| counts[f"{name}_fn"] = int((gt & ~selected).sum()) |
| inner_flat = flat & (labels == 0) |
| counts["flat_inner_count"] = int(inner_flat.sum()) |
| counts["flat_inner_false_outer"] = int((inner_flat & (pred == 1)).sum()) |
| head = (target["part"] == 0) & valid |
| gt_head_outer, pred_head_outer = head & (labels == 1), head & (pred == 1) |
| counts["head_outer_tp"] = int((gt_head_outer & pred_head_outer).sum()) |
| counts["head_outer_fp"] = int((~gt_head_outer & pred_head_outer).sum()) |
| counts["head_outer_fn"] = int((gt_head_outer & ~pred_head_outer).sum()) |
| return counts |
|
|
|
|
| def summarize_counts(counts): |
| values = dict(counts) |
| for name in ("inner", "outer", "secondary", "head_outer"): |
| tp, fp, fn = [counts[f"{name}_{k}"] for k in ("tp", "fp", "fn")] |
| values[f"{name}_precision"] = tp / max(1, tp + fp) |
| values[f"{name}_recall"] = tp / max(1, tp + fn) |
| values[f"{name}_iou"] = tp / max(1, tp + fp + fn) |
| values["flat_inner_false_outer_rate"] = ( |
| counts["flat_inner_false_outer"] / max(1, counts["flat_inner_count"])) |
| return values |
|
|
|
|
| @torch.no_grad() |
| def evaluate(model, renderer, loader, args, device, ablation=True, hard=True): |
| model.eval() |
| totals = {name: {} for name in ("clean", "appearance")} |
| if ablation: |
| totals["no_semantics"] = {} |
| hard_totals = {} |
| rng = torch.Generator(device=device).manual_seed(90451) |
| for batch in loader: |
| uv = batch["uv"].to(device) |
| images, targets, _, ids = core.build_parser_inputs(uv, renderer, args.views, False, args) |
| flat = flat_foreground_mask(images, targets) |
| with core.autocast_context(device, "bf16"): |
| |
| clean_features = model._runtime_semantic_features(images, targets["foreground"]) |
| for mode in totals: |
| if mode == "appearance": |
| actual = appearance_augment(images, targets["foreground"], 1.25, rng) |
| output = model(actual, view_ids=ids, semantic_foreground=targets["foreground"]) |
| elif mode == "no_semantics": |
| with semantics_disabled(model): |
| output = model(images, view_ids=ids, semantic_features=clean_features) |
| else: |
| output = model(images, view_ids=ids, semantic_features=clean_features) |
| counts = count_routes(output, targets, flat) |
| for key, value in counts.items(): |
| totals[mode][key] = totals[mode].get(key, 0) + value |
| if mode == "clean" and hard: |
| hm = core.hard_uv_conditioning_metrics(images, output, targets, renderer, args.views, args) |
| for key, value in hm.items(): |
| if key.startswith("count_"): |
| hard_totals[key] = hard_totals.get(key, 0.0) + float(value) |
| results = {k: summarize_counts(v) for k, v in totals.items()} |
| if hard_totals: |
| hard_metrics = {} |
| for name in ("inner", "outer"): |
| tp, fp, fn = [hard_totals[f"count_hard_{name}_{k}"] for k in ("tp", "fp", "fn")] |
| for metric, value in (("precision", tp / max(1, tp + fp)), |
| ("recall", tp / max(1, tp + fn)), |
| ("iou", tp / max(1, tp + fp + fn))): |
| hard_metrics[f"{name}_{metric}"] = value |
| hard_metrics[f"{name}_rgb_mae"] = ( |
| hard_totals[f"count_hard_{name}_rgb_abs"] / |
| max(1, hard_totals[f"count_hard_{name}_rgb_values"])) |
| results["hard_uv"] = hard_metrics |
| return results |
|
|
|
|
| def eligible(metrics, baseline, tolerance): |
| |
| for section in ("clean", "hard_uv"): |
| for name in ("inner_iou", "outer_iou", "outer_precision", "outer_recall"): |
| if metrics[section][name] < baseline[section][name] - tolerance: |
| return False |
| return True |
|
|
|
|
| def robustness_score(metrics): |
| appearance = metrics["appearance"] |
| return (0.40 * appearance["head_outer_recall"] |
| + 0.25 * appearance["outer_recall"] |
| + 0.25 * appearance["outer_iou"] |
| + 0.10 * appearance["inner_iou"]) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| root = Path(__file__).resolve().parents[1] |
| parser.add_argument("--checkpoint", type=Path, default=root / "dense_uv_parser/runs/dense_uv_parser_v61/best.pt") |
| parser.add_argument("--output_dir", type=Path, required=True) |
| parser.add_argument("--epochs", type=int, default=2) |
| parser.add_argument("--train_samples", type=int, default=0, help="0 uses all original 162000 training skins") |
| parser.add_argument("--val_samples", type=int, default=512) |
| parser.add_argument("--test_samples", type=int, default=256, |
| help="Additional original validation skins used only for the final report") |
| parser.add_argument("--batch_size", type=int, default=8) |
| parser.add_argument("--workers", type=int, default=4) |
| parser.add_argument("--lr", type=float, default=1e-5) |
| parser.add_argument("--semantic_lr", type=float, default=5e-5) |
| parser.add_argument("--skip_dropout", type=float, default=0.10) |
| parser.add_argument("--eval_every", type=int, default=500) |
| parser.add_argument("--log_every", type=int, default=25) |
| parser.add_argument("--max_clean_regression", type=float, default=0.015) |
| parser.add_argument("--evaluate_only", action="store_true") |
| opt = parser.parse_args() |
| if min(opt.epochs, opt.val_samples, opt.batch_size, opt.eval_every, opt.log_every) < 1: |
| parser.error("epochs, val_samples, batch_size, eval_every, log_every must be positive") |
| if not 0 <= opt.skip_dropout < 1 or min(opt.train_samples, opt.test_samples, opt.workers) < 0: |
| parser.error("invalid dropout or sample/worker count") |
| opt.output_dir.mkdir(parents=True, exist_ok=False) |
| device = torch.device("cuda") |
| torch.set_num_threads(8) |
| torch.backends.cudnn.benchmark = True |
| torch.set_float32_matmul_precision("high") |
| core.seed_everything(1234, reproducible=False) |
| checkpoint = torch.load(opt.checkpoint, map_location="cpu", weights_only=False) |
| if checkpoint["model_config"].get("predict_outer_uv_occupancy"): |
| raise ValueError("Expected the v61 model without an occupancy head") |
| args_dict = vars(core.build_arg_parser().parse_args(["--data_dir", str(root / "skins")])) |
| args_dict.update(checkpoint["args"]) |
| args_dict.update(data_dir=str(root / "skins"), |
| mappings_dir=str(root.parent / "github/differentiable_minecraft_renderer/mappings_256x512"), |
| output_dir=str(opt.output_dir), privileged_views="", resume=str(opt.checkpoint), |
| siglip_local_files_only=True, best_metric="semantic_robustness", |
| outer_uv_min_source_pixels=33, |
| lambda_outer_false_negative=1.25, |
| route_outer_class_weight_cap=1.0) |
| args = SimpleNamespace(**args_dict) |
| args.views = core.parse_views(args.views) |
| args.epochs = checkpoint["epoch"] + opt.epochs |
| |
| dataset = SkinUVDataset(args.data_dir, max_samples=args.max_samples) |
| original_cache = root / "dense_uv_parser" / checkpoint["args"]["siglip_cache_dir"] / "metadata.json" |
| if original_cache.is_file(): |
| original_names = json.loads(original_cache.read_text())["filenames"] |
| if [path.name for path in dataset.skin_paths] != original_names: |
| raise ValueError("Dataset ordering differs from the v61 cache; reconstruct the original split first") |
| train_set, val_set = build_split(dataset, args.seed, args.val_split) |
| if opt.train_samples: |
| train_set = Subset(dataset, train_set.indices[:opt.train_samples]) |
| test_set = Subset(dataset, val_set.indices[opt.val_samples:opt.val_samples + opt.test_samples]) |
| val_set = Subset(dataset, val_set.indices[:opt.val_samples]) |
| if not len(train_set) or not len(val_set): |
| raise ValueError("Empty training/validation split") |
| if set(train_set.indices) & set(val_set.indices): |
| raise RuntimeError("Training/validation leakage") |
| loader_args = dict(batch_size=opt.batch_size, num_workers=opt.workers, pin_memory=True) |
| if opt.workers: |
| loader_args.update(persistent_workers=True, prefetch_factor=2) |
| train_loader = DataLoader(train_set, shuffle=True, |
| generator=torch.Generator().manual_seed(args.seed + 1), **loader_args) |
| val_loader = DataLoader(val_set, shuffle=False, **loader_args) |
| test_loader = DataLoader(test_set, shuffle=False, **loader_args) if len(test_set) else None |
| model_kwargs = {key: val for key, val in checkpoint["model_config"].items() |
| if key in inspect.signature(DenseUVParserNet).parameters} |
| model = DenseUVParserNet(**model_kwargs).to(device) |
| model.load_state_dict(checkpoint["model"], strict=True) |
| |
| model.route_role_prior.requires_grad_(False) |
| attach_semantic_runtime(model, "siglip2", args.siglip_model, device, local_files_only=True) |
| renderer = DifferentiableRenderer(args.mappings_dir).to(device) |
| loss_kwargs = {key: value for key, value in args_dict.items() |
| if key in inspect.signature(DenseUVParserLoss).parameters} |
| loss_kwargs.update(use_uv=False, affine_translation_limit=model.affine_translation_limit, |
| affine_log_scale_limit=model.affine_log_scale_limit) |
| criterion = DenseUVParserLoss(**loss_kwargs).to(device) |
| masks = tuple(mask.to(device) for mask in build_part_layer_masks()) |
| semantic, visual = [], [] |
| for name, parameter in model.named_parameters(): |
| if parameter.requires_grad: |
| (semantic if name.startswith(("semantic_", "outer_presence", "outer_coverage")) |
| else visual).append(parameter) |
| optimizer = torch.optim.AdamW([ |
| {"params": visual, "lr": opt.lr, "initial_lr": opt.lr}, |
| {"params": semantic, "lr": opt.semantic_lr, "initial_lr": opt.semantic_lr}, |
| ], weight_decay=args.weight_decay) |
| handles = skip_regularizer(model, opt.skip_dropout) |
| manifest = {"checkpoint": str(opt.checkpoint.resolve()), |
| "checkpoint_sha256": hashlib.sha256(opt.checkpoint.read_bytes()).hexdigest(), |
| "source_epoch": checkpoint["epoch"], "checkout": str(root), |
| "code_files": {Path(m.__file__).name: str(Path(m.__file__).resolve()) |
| for m in (core, inspect.getmodule(DenseUVParserNet), inspect.getmodule(DenseUVParserLoss))}, |
| "options": {k: str(v) if isinstance(v, Path) else v for k, v in vars(opt).items()}, |
| "train_indices": list(train_set.indices), "val_indices": list(val_set.indices), |
| "test_indices": list(test_set.indices), |
| "dataset_filenames_sha256": hashlib.sha256( |
| "\n".join(p.name for p in dataset.skin_paths).encode()).hexdigest(), |
| "label_source": "original skin RGBA and v61 renderer surface visibility; no v94 pseudo-labels", |
| "semantic_features": "online from actual input, frozen SigLIP2", |
| "optimizer": "fresh AdamW for changed fine-tuning objective; v61 model loaded strictly", |
| "routing": "v61 thresholds held fixed for hard UV validation, source pixel minimum 33"} |
| for path in manifest["code_files"].values(): |
| if not Path(path).is_relative_to(root): |
| raise RuntimeError(f"Imported outside checkout: {path}") |
| write_json(opt.output_dir / "config.json", {"args": vars(args), "experiment": manifest}) |
| source_dir = opt.output_dir / "source" |
| source_dir.mkdir() |
| for source in (root / "dense_uv_parser").glob("*.py"): |
| shutil.copy2(source, source_dir / source.name) |
| shutil.copy2(root / "renderer.py", source_dir / "renderer.py") |
| started = time.time() |
| write_json(opt.output_dir / "status.json", {"state": "baseline", "pid": os.getpid()}) |
| print("Evaluating untouched v61 baseline", flush=True) |
| baseline = evaluate(model, renderer, val_loader, args, device) |
| write_json(opt.output_dir / "baseline.json", baseline) |
| print("baseline=" + json.dumps(baseline), flush=True) |
| if opt.evaluate_only: |
| write_json(opt.output_dir / "status.json", {"state": "evaluated", "pid": os.getpid()}) |
| return |
| best_score = robustness_score(baseline) |
| step = 0 |
| stop = False |
| rejection_streak = 0 |
| args.semantic_generalization = manifest["options"] |
| for epoch in range(1, opt.epochs + 1): |
| running = 0.0 |
| for batch_index, batch in enumerate(train_loader): |
| model.train() |
| progress = (step / max(1, opt.epochs * len(train_loader) - 1)) |
| factor = 0.2 + 0.8 * (1 + math.cos(math.pi * progress)) / 2 |
| for group in optimizer.param_groups: |
| group["lr"] = group["initial_lr"] * factor |
| uv = batch["uv"].to(device, non_blocking=True) |
| images, targets, _, ids = core.build_parser_inputs(uv, renderer, args.views, True, args) |
| flat = flat_foreground_mask(images, targets) |
| augmented = step % 2 == 1 |
| actual = appearance_augment(images, targets["foreground"]) if augmented else images |
| optimizer.zero_grad(set_to_none=True) |
| with core.autocast_context(device, "bf16"): |
| output = model(actual, view_ids=ids, semantic_foreground=targets["foreground"]) |
| losses = criterion(output, targets) |
| loss = losses["loss_total"] |
| loss = loss + 0.10 * flat_route_loss(output["layer"], targets["route_role"], flat) |
| loss = loss + 0.15 * head_outer_hard_positive_loss(output["layer"], targets) |
| attributes = build_semantic_attribute_targets(uv, *masks) |
| loss = loss + 0.25 * F.binary_cross_entropy_with_logits( |
| output["outer_presence_logits"], attributes["outer_presence"]) |
| loss = loss + 0.25 * F.smooth_l1_loss(output["outer_coverage"], attributes["outer_coverage"]) |
| if not augmented: |
| geometry = core.differentiable_geometry_losses(images, uv, output, targets, renderer, args.views) |
| for key in ("soft_uv_rgb", "soft_uv_alpha", "soft_uv_inner_recall", |
| "soft_uv_outer_recall", "render_rgb", "render_alpha"): |
| loss = loss + getattr(args, "lambda_" + key) * geometry["loss_" + key] |
| if not torch.isfinite(loss): |
| raise FloatingPointError(f"Nonfinite loss at step {step}") |
| loss.backward() |
| norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0, error_if_nonfinite=True) |
| optimizer.step() |
| step += 1 |
| running += float(loss.detach()) |
| if step % opt.log_every == 0 or step == 1: |
| status = {"state": "training", "pid": os.getpid(), "step": step, "epoch": epoch, |
| "total_steps": opt.epochs * len(train_loader), "loss": running / (batch_index + 1), |
| "gradient_norm": float(norm), "semantic_lr": optimizer.param_groups[1]["lr"], |
| "elapsed_seconds": time.time() - started} |
| write_json(opt.output_dir / "status.json", status) |
| print(json.dumps(status), flush=True) |
| if step % opt.eval_every == 0 or batch_index + 1 == len(train_loader): |
| metrics = evaluate(model, renderer, val_loader, args, device) |
| score = robustness_score(metrics) |
| accepted = eligible(metrics, baseline, opt.max_clean_regression) |
| promoted = accepted and score > best_score |
| record = {"step": step, "epoch": epoch, "score": score, "clean_guard_pass": accepted, |
| "promoted": promoted, "metrics": metrics, "elapsed_seconds": time.time() - started} |
| with (opt.output_dir / "metrics.jsonl").open("a") as stream: |
| stream.write(json.dumps(record) + "\n") |
| print("evaluation=" + json.dumps(record), flush=True) |
| if promoted: |
| best_score = score |
| args.finetune_step = step |
| saved_metrics = {"val": metrics, "semantic_robustness": score} |
| core.save_checkpoint(opt.output_dir / "latest.pt", model, optimizer, None, |
| checkpoint["epoch"] + epoch, args, saved_metrics, best_metric=-best_score) |
| if promoted: |
| core.save_checkpoint(opt.output_dir / "best.pt", model, optimizer, None, |
| checkpoint["epoch"] + epoch, args, saved_metrics, best_metric=-best_score) |
| write_json(opt.output_dir / "last_evaluation.json", record) |
| rejection_streak = 0 if accepted else rejection_streak + 1 |
| if rejection_streak >= 2: |
| stop = True |
| break |
| if stop: |
| break |
| for handle in handles: |
| handle.remove() |
| if test_loader is not None and (opt.output_dir / "best.pt").is_file(): |
| write_json(opt.output_dir / "status.json", {"state": "final_test", "pid": os.getpid(), "step": step}) |
| best = torch.load(opt.output_dir / "best.pt", map_location="cpu", weights_only=False) |
| model.load_state_dict(best["model"], strict=True) |
| candidate_test = evaluate(model, renderer, test_loader, args, device) |
| model.load_state_dict(checkpoint["model"], strict=True) |
| baseline_test = evaluate(model, renderer, test_loader, args, device) |
| write_json(opt.output_dir / "heldout_test.json", { |
| "note": "Not used for checkpoint selection; appearance stress is synthetic, not real-image ground truth", |
| "samples": len(test_set), "candidate": candidate_test, "baseline": baseline_test}) |
| write_json(opt.output_dir / "status.json", {"state": "stopped_clean_regression" if stop else "complete", |
| "pid": os.getpid(), "step": step, "best_score": best_score, |
| "baseline_score": robustness_score(baseline), "best_checkpoint_exists": (opt.output_dir / "best.pt").exists(), |
| "elapsed_seconds": time.time() - started}) |
|
|
|
|
| if __name__ == "__main__": |
| try: |
| main() |
| except Exception as error: |
| if "--output_dir" in sys.argv: |
| index = sys.argv.index("--output_dir") + 1 |
| if index < len(sys.argv): |
| failed_dir = Path(sys.argv[index]) |
| if failed_dir.is_dir() and not isinstance(error, FileExistsError): |
| write_json(failed_dir / "status.json", { |
| "state": "failed", "pid": os.getpid(), |
| "error": f"{type(error).__name__}: {error}"}) |
| raise |
|
|
