Colorizer round6: source and write verification before training
Browse files- experiments/decision-20260928-v2/PROTOCOL.md +47 -0
- experiments/decision-20260928-v2/REVIEW_HANDOFF.md +30 -0
- experiments/decision-20260928-v2/SHA256SUMS.json +22 -0
- experiments/decision-20260928-v2/source/data.py +116 -0
- experiments/decision-20260928-v2/source/decision_model.py +117 -0
- experiments/decision-20260928-v2/source/inference.py +111 -0
- experiments/decision-20260928-v2/source/legacy_eval.py +150 -0
- experiments/decision-20260928-v2/source/metrics.py +35 -0
- experiments/decision-20260928-v2/source/model.py +122 -0
- experiments/decision-20260928-v2/source/objectives.py +17 -0
- experiments/decision-20260928-v2/source/persistence.py +42 -0
- experiments/decision-20260928-v2/source/r7.py +234 -0
- experiments/decision-20260928-v2/source/semantic_model.py +94 -0
- experiments/decision-20260928-v2/source/spatial.py +34 -0
- experiments/decision-20260928-v2/source/train_suite.py +359 -0
- experiments/decision-20260928-v2/status.json +5 -0
- experiments/decision-20260928-v2/supplement/baseline_nara_resolution.json +128 -0
- experiments/decision-20260928-v2/supplement/baseline_resolution.json +23 -0
- experiments/decision-20260928-v2/supplement/cache_loader_check.json +8 -0
- experiments/decision-20260928-v2/supplement/nara_sources.json +50 -0
- experiments/decision-20260928-v2/supplement/resolution_review.py +27 -0
experiments/decision-20260928-v2/PROTOCOL.md
ADDED
|
@@ -0,0 +1,47 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Penultimate architecture and curriculum decision — 28 September 2026
|
| 2 |
+
|
| 3 |
+
## Objective and decision scope
|
| 4 |
+
Select one reproducible course for the user's next, longer training run. Prioritize coherent plausible colors, fewer within-surface patches, and retained color over recovering an unknowable original hue. Strictly fewer than 4,000,000 deployed learned parameters, including encoder and any selector. This screen cannot establish a universal architecture ceiling. No experimental model replaces root production weights during this pass.
|
| 5 |
+
|
| 6 |
+
## Evidence audit
|
| 7 |
+
The uploaded project log describes an older classification U-Net. Current production is round6 palette step9000 at base revision `1a9eb8af2754ad2329a24cfe50d388cb559441d0`, 3,994,676 parameters. The authoritative current repo revision before this pass was `26d731862cc0ea531e7198fa823cb09db25bd243`.
|
| 8 |
+
|
| 9 |
+
The semantic encoder/palette change improved substantially over the old U-Net. Round7's easier-of-original-or-teacher objective selected original targets only 5.8–9.5% of the time; it did not learn explicit alternative outputs. Critic variants (including 27M, smaller, and dense critics) failed visual release review. A region auxiliary and repeating the same decoder did not produce a clear gain. These results motivate separating color ambiguity from spatial assignment, not another saturation/critic sweep.
|
| 10 |
+
|
| 11 |
+
## Primary research and implications
|
| 12 |
+
- Zhang et al., Colorful Image Colorization: https://richzhang.github.io/colorization/ — classification/rebalancing addresses uncertainty; pixel-wise decoding still needs spatial coordination. Earlier log language that classification is universally required was an overstatement: later palette regression already outperformed that baseline.
|
| 13 |
+
- Deshpande et al., Learning Diverse Image Colorization: https://arxiv.org/abs/1612.01958 and https://arxiv.org/html/1612.01958v2 — low-dimensional color-field variables coordinate alternatives over long spatial distances. Motivates global modes instead of pixel noise.
|
| 14 |
+
- Wang et al., PalGAN: https://arxiv.org/html/2210.11204v1 — probabilistic palette prediction and assignment separate ambiguity from boundary handling. This project's small palette model is not a PalGAN reproduction.
|
| 15 |
+
- Lee et al., Stochastic Multiple Choice Learning: https://arxiv.org/abs/1606.07839 — winner-based learning supports distinct hypotheses, but oracle quality is not automatic-default quality. We explicitly measure the gap and train an image-only selector.
|
| 16 |
+
- Mehta and Rastegari, MobileViTv2: https://arxiv.org/abs/2206.02680 — efficient global attention in a mobile encoder. Test the ImageNet-pretrained 0.75 variant as a distinct backbone challenger, with all deployment parameters counted.
|
| 17 |
+
- Saharia et al., Palette: https://arxiv.org/abs/2111.05826 — diffusion supports colorization diversity. Not a justified final-run choice here without training/validating a compact denoiser; repeating the old decoder did not test diffusion. Excluded from this finite screen because establishing a from-scratch diffusion model and sampler is substantially larger work than these warm-started alternatives.
|
| 18 |
+
|
| 19 |
+
## Controlled comparisons
|
| 20 |
+
Six arms: single palette vs four coherent palettes, crossed with teacher / mixed-target / staged curriculum. Same MobileNet v3 initialization, fixed8192 training images, seeds, image-index batches, augmentation draws, AdamW, batch24, 3000 updates. Encoder LR1e-5, decoder LR1e-4, warmup150 and cosine decay to15%, weight decay0.01, frozen encoder BatchNorm, gradient norm cap1, BF16 forwards/FP32 losses. New mode tensors and selector are included in the decoder group. The multi head adds only1,156 parameters (3,995,832 total).
|
| 21 |
+
|
| 22 |
+
Multi head: shared spatial masks and encoder; four latent palette variants; image-only score selects one complete output. Training loss is95% winner reconstruction +5% mean reconstruction +0.1 selector cross-entropy against the detached winning index. Reconstruction is L1/20 +0.25 sliced color distribution distance +0.1 boundary difference. Winner is selected per entire image, never per pixel. Report mode usage, between-mode diversity, default vs oracle losses and mode-switching under input resize. No original color is available to the runtime selector. No hue rotation, saturation bonus or unvalidated generated 'plausible' labels.
|
| 23 |
+
|
| 24 |
+
Curricula:
|
| 25 |
+
1. Teacher: deterministic cached DDColor artistic targets.
|
| 26 |
+
2. Mixed: independently sample original photograph or teacher target with probability0.5 per image each update. Never average two hues and never pick the easier target. Both represent complete color proposals; teacher remains fallible.
|
| 27 |
+
3. Staged: first15% teacher, next35% ramp original probability to0.5; final third adds35% mild spectral/film/noise augmentation plus same-mode resize consistency (mean coefficient0.1, computed every fourth update). This tests a curriculum bundle, not isolated causal effects of each component. Anchored reconstruction remains active to guard against consistency-induced gray collapse.
|
| 28 |
+
|
| 29 |
+
Seventh arm: MobileViTv2-0.75 + single palette, copied compatible decoder weights and new lateral projections. 3,481,205 parameters.1000 extra teacher alignment updates, then3000 staged updates; encoder LR2e-5, decoder2e-4. Pretrained encoder pinned to the revision recorded in encoder_provenance.json. This is a practical challenger with unequal prior colorization training and an extra alignment stage, not a pure equal-initialization ablation. A negative short pilot does not prove the encoder incapable.
|
| 30 |
+
|
| 31 |
+
## Data and evaluation
|
| 32 |
+
Training subset sampled from existing filtered16,230-image Imagenette/COCO training pool; no new teacher expense. Existing target cache hashes and row IDs checked. Pin dataset revisions. New128 development and160 test images from COCO val/0001 exclude earlier round7, region, recurrent evaluation IDs and exact-byte training duplicates. Teacher/encoder upstream exposure and near duplicates remain possible. Fix splits before training. Choose each arm's candidate from evaluations every1000 updates on development only; test opened after choices frozen. Keep all evaluated checkpoints.
|
| 33 |
+
|
| 34 |
+
Screen on ordinary and film-style grayscale. Record raw and guided patch metrics, chroma retention, missed color, spill, auxiliary original-color error, resize and flip sensitivity, latency, and image-bootstrap paired intervals. Gray output can win a smoothness metric; do not rank on patch score alone. Screening guardrails allow missed color/spill at most+1.5points, coverage at least95% baseline, and error at most110% baseline. These are diagnostic eligibility conditions, not proof of plausibility. All checkpoint scans are exploratory and intervals do not correct multiple comparisons.
|
| 35 |
+
|
| 36 |
+
Visual review: 32 fixed fresh examples per group, archival photos never used as neutral targets, seven familiar probes, and aspect-preserving production-pipeline outputs. Label comparisons anonymously with saved keys; assistant review is qualitative, not an independent human preference study. Archive source URLs and checksums retained. A modest historical panel is not a representative archival benchmark.
|
| 37 |
+
|
| 38 |
+
## Targeted fine-tuning and RL
|
| 39 |
+
Late-stage robustness and consistency are tested rather than assumed helpful. No RL launched in this screen: no independently validated preference reward currently exists, and previous learned critics were exploitable. Prefer direct differentiable supervision for measurable consistency; if a later reward is proposed it must first rank coherent alternate hues above gray collapse and patchwork on held-out images. No claim that RL is intrinsically unsuitable.
|
| 40 |
+
|
| 41 |
+
## Persistence, budget, and final selection
|
| 42 |
+
One L4 job with60min hard limit and50min internal training deadline; all outputs checkpointed to this main-branch experiment folder with SHA256 download verification. Source/access/preflight checks before training. The first attempt was canceled before training after about12min because of repeated NPZ decompression. Corrected version2 materializes each cache once and logs data preparation. A separate confirmation, if warranted, is capped18min; total maximum hardware charges at the verified$0.80/h rate approximately$1.20, excluding no separate service purchases. Current account balance is not exposed by the connector. No automatic paid reruns.
|
| 43 |
+
|
| 44 |
+
Choose architecture/curriculum from test plus visual review, with production baseline retained as an eligible fallback. If a candidate clearly wins, confirm against the matching control on a second seed or targeted ablation within the remaining cap. Save final_config.json and NEXT_PASS.md with exact checkpoints, source revisions, dataset/curriculum, learning rates, retention gates and validation instructions. The next pass starts the longer run and performs Space/model-card/repository work. This pass does not claim any longer run already completed.
|
| 45 |
+
|
| 46 |
+
## Version2 corrections before training
|
| 47 |
+
NPZ cache members are decompressed once per shard, not once per selected image; a synthetic equivalence and timing check is saved. MobileViT now uses a separate post-warmup curriculum step. Six additional National Archives inputs supplement the original historical panel. Version1 was canceled before any training checkpoint; its source/preflight remains for audit.
|
experiments/decision-20260928-v2/REVIEW_HANDOFF.md
ADDED
|
@@ -0,0 +1,30 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Experimental pass handoff — 2026-09-28
|
| 2 |
+
|
| 3 |
+
## Status
|
| 4 |
+
Version1 job `6aba3fb36b030d633f69c572` was canceled before training to fix repeated NPZ decompression. The replacement runs the seven-arm suite. Source, architecture preflight, parameter counts and upload/download verification are already committed to `experiments/decision-20260928-v2` on main. At this handoff no trained candidate has been reviewed or approved. Do not describe the suite as complete or declare a winner without reading its outputs.
|
| 5 |
+
|
| 6 |
+
Root production remains v3 at `1a9eb8af2754ad2329a24cfe50d388cb559441d0`. The next pass should review this suite before launching the final longer run. Never overwrite root model files merely because an experimental score is lower.
|
| 7 |
+
|
| 8 |
+
## Exact procedure for selecting the final course
|
| 9 |
+
1. Read status.json, failure.json if present, summary.json and candidates.json. Verify every candidate's actual steps and parameter count; a partial/timed-out screen cannot establish the missing comparison. All model counts include encoder and selector.
|
| 10 |
+
2. Read each development result and selected test result. The test is newly held out from project training and listed prior evaluations; upstream encoder/teacher exposure and near duplicates cannot be excluded. Use paired intervals as exploratory diagnostics, not multiple-comparison-corrected proof.
|
| 11 |
+
3. View comparison JPEGs **before reading their key.json files**. Inspect both full frames and useful crops. Review all 32 fixed COCO examples for single/multi groups, the encoder group, and actual-aspect production_test images. Log visible findings per image: incoherent hue changes on one surface, missing colour, boundary leakage, implausible skin/sky/foliage, whole-image colour cast. An alternate plausible clothing colour is not a failure simply because it differs from the original.
|
| 12 |
+
4. View the seven probes, historical comparison images, and actual-aspect production_archive images. Read archive_sources.json to identify failed downloads; do not report attempted images as evaluated. Supplement with the six National Archives inputs whose source records are in `supplement/nara_sources.json`. These are evaluation-only. Never supervise historical grayscale inputs with zero chroma.
|
| 13 |
+
5. Inspect multi diagnostics: all mode frequencies, diversity, default-v-oracle losses and resize/flip mode switches. An oracle-only gain is insufficient for automatic deployment. An optional alternatives UI may be valuable but must not conceal a worse default. Inspect MobileViT visual outputs in light of its shorter colourisation history; do not declare a universally inferior backbone from a short warm-start test.
|
| 14 |
+
6. Prefer a candidate only if images show a consistent benefit without unacceptable loss of colour, and quantitative diagnostics do not reveal collapse. Retain v3 as the fallback. If one method looks clearly better, a second-seed matched confirmation is authorized within the **remaining** pass budget; maximum additional L4 runtime18min, not an unbounded rerun. Previous attempt used about12min; replacement capped60min at$0.80/h. Preserve all results before exit.
|
| 15 |
+
7. Write REVIEW.md and final_config.json with the actual checkpoint path, immutable commit, architecture, complete parameter count, recipe, seed, dataset revisions, teacher-cache revision and validation evidence. If nothing wins, state that plainly and use the established palette architecture. Do not promise that longer continuation necessarily improves it.
|
| 16 |
+
|
| 17 |
+
## Specific implementation caveats to carry forward
|
| 18 |
+
- Single/multi arms: same initial encoder and decoder, index batches, learning rates and update counts. Staged augmentation changes random-number consumption late in training; its flip/target-choice draws are not identical to teacher/mix. This is a curriculum-bundle comparison.
|
| 19 |
+
- MobileViT receives1000 teacher alignment steps plus3000 curriculum steps. Version2 uses the post-warmup curriculum step for both augmentation/target ramps and consistency, correcting the canceled first attempt.
|
| 20 |
+
- Each arm's candidate can come from an earlier development checkpoint. optimizer.pt belongs to the **last** step, not necessarily that selected checkpoint. Start a new optimizer from selected candidate weights; never pair candidate weights with a mismatched optimizer.
|
| 21 |
+
- Existing r7 loss uses teacher/original targets for supervision; neither original colour error nor patch score alone measures plausibility.
|
| 22 |
+
- Do not add RL without an independently validated reward. Previous critic failures are project evidence for caution, not proof RL can never help.
|
| 23 |
+
- Keep original-image luminance and aspect ratio through the established inference pipeline. Single/multi DecisionColorizer needs its own loader; root load_semantic cannot silently load a new architecture.
|
| 24 |
+
- If deploying MobileViT, include its upstream Apple license/notice and appropriate pretrained-model attribution.
|
| 25 |
+
|
| 26 |
+
## Inference compute experiment already inspected
|
| 27 |
+
Seven real historical photographs were evaluated with unchanged v3 weights at256/384/512 input size and the existing radius8 filter. On the six new National Archives images, larger size notably improves the children-in-a-field image (more coherent green foliage and blue sky), and slightly clarifies some people. It does not fix the near-monochrome casts on machinery, street and interior scenes. The familiar powerhouse pipe retains uneven colouring and changes hue with size. Therefore higher resolution is useful in some cases but is not a general cure. At fixed radius8, its physical smoothing footprint also changes: do not call this an isolated network-resolution ablation. Record measured CPU times as local indicative timings, not a deployment latency guarantee; they include a second diagnostic forward.
|
| 28 |
+
|
| 29 |
+
## Final-pass scope
|
| 30 |
+
After evidence review, start one longer run on the selected course, with checkpoint retention, bounded runtime, early exit for non-finite losses, durable metrics and uploads. Final production promotion requires comparing its saved checkpoints visually to v3 and the selected pilot. Develop the Space/ZeroGPU app, requirements and model card, and clean the repository without removing experiments or the stable branch needed by the user. Never state that the model is mathematically optimal or that no improvement is possible under4M parameters.
|
experiments/decision-20260928-v2/SHA256SUMS.json
ADDED
|
@@ -0,0 +1,22 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"PROTOCOL.md": "7d73f476df74bf01b6ea93346fd08b98dde0456af62eda5b297b6a3044e33869",
|
| 3 |
+
"REVIEW_HANDOFF.md": "feae011a9a04136caf796224ae92e913eda5bc9e47c94a7dfea481e071bb58b2",
|
| 4 |
+
"source/data.py": "eb43b886fe45435841439c36646985f914386a48faa712fb6806accac2ad48e0",
|
| 5 |
+
"source/decision_model.py": "c7c4370d93bc52fdd363d6d4fcc684c317806d9b393bd969f9ee670a01487f27",
|
| 6 |
+
"source/inference.py": "0ac1f382205b42cef0afaa9173017fa660e6859194a65a61b9b9052c437cd476",
|
| 7 |
+
"source/legacy_eval.py": "dd7e8bd5f961dfb6f1c2ce98fac41824c5009188e306000b195e922e4ed03e58",
|
| 8 |
+
"source/metrics.py": "af0ac996f6be0a351a55e372fab637a0d06740afe4a744cb80d0eae56ec97f66",
|
| 9 |
+
"source/model.py": "4cc57f82ffd6378bdf23a088ce7a1ed56c09e5673de3ee5232b8ee1cacc9be0a",
|
| 10 |
+
"source/objectives.py": "133e85b6f2914c0e3803cc26a26fa29b779a49036219dd245c2ecfe0ba411e20",
|
| 11 |
+
"source/persistence.py": "8b26dd8acdc135ecbfcdd30b7f7723feb6008ac820c083403da1e820bdc15d33",
|
| 12 |
+
"source/r7.py": "019ea992052f0233954b86e06e31c85c74ca313e67a55ff7fb03248f8874aaf5",
|
| 13 |
+
"source/semantic_model.py": "b019358cbcaa214743cbd244d20eb6e59eeff0ab5264093ec0927be8529b25a2",
|
| 14 |
+
"source/spatial.py": "d13abe399ef049e21a6459a7003461afaae0c00e0c560262c6cf375de4c9884a",
|
| 15 |
+
"source/train_suite.py": "fe09629930b44a2ce4d630339b15f6cb3c09a4504de890df41016113163efb6a",
|
| 16 |
+
"status.json": "820e46a29d75ee397d8d2814ad0429cb3d300952f32030d83e06a4dc0b56f803",
|
| 17 |
+
"supplement/baseline_nara_resolution.json": "8920921e71ef95ed442495171a0ce229b5338170c5ccdf6d60e241fdd61cb570",
|
| 18 |
+
"supplement/baseline_resolution.json": "e770a0dff51460de711b05a2dcb107e128f8635f05284bc075816da8a44f51e2",
|
| 19 |
+
"supplement/cache_loader_check.json": "b7f282c99f9371b40f1a0b403bd31296ed1c93496d99ed60e144e088397f22f9",
|
| 20 |
+
"supplement/nara_sources.json": "d4551fb5857161ecd72be4f315dc9e643e33700827e7fcf51e90049a6ef0f6ba",
|
| 21 |
+
"supplement/resolution_review.py": "d942406942a5197b54db053ea696fc59ff8b5d0093d4e0b5817431db55e9d115"
|
| 22 |
+
}
|
experiments/decision-20260928-v2/source/data.py
ADDED
|
@@ -0,0 +1,116 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Fixed color-bin semantics and reproducible dataset membership."""
|
| 2 |
+
import hashlib
|
| 3 |
+
import io
|
| 4 |
+
import json
|
| 5 |
+
from pathlib import Path
|
| 6 |
+
import numpy as np
|
| 7 |
+
from PIL import Image
|
| 8 |
+
import pyarrow.parquet as pq
|
| 9 |
+
from scipy.spatial import cKDTree
|
| 10 |
+
from skimage.color import rgb2lab
|
| 11 |
+
import torch
|
| 12 |
+
from torch.utils.data import Dataset
|
| 13 |
+
|
| 14 |
+
class ColorBins:
|
| 15 |
+
def __init__(self, centers, weights=None):
|
| 16 |
+
self.centers = np.asarray(centers, np.float32)
|
| 17 |
+
if self.centers.ndim != 2 or self.centers.shape[1] != 2:
|
| 18 |
+
raise ValueError("Expected Q by 2 color centers")
|
| 19 |
+
self.tree = cKDTree(self.centers)
|
| 20 |
+
self.weights = np.ones(len(self.centers), np.float32) if weights is None else np.asarray(weights,np.float32)
|
| 21 |
+
|
| 22 |
+
def encode(self, ab, k=5, sigma=5):
|
| 23 |
+
k = min(k, len(self.centers))
|
| 24 |
+
if sigma <= 0 or k < 1:
|
| 25 |
+
raise ValueError("Invalid encoding parameters")
|
| 26 |
+
dist, idx = self.tree.query(ab.reshape(-1,2), k=k)
|
| 27 |
+
dist, idx = dist.reshape(-1,k), idx.reshape(-1,k)
|
| 28 |
+
logw = -dist**2/(2*sigma**2)
|
| 29 |
+
logw -= logw.max(1, keepdims=True)
|
| 30 |
+
weight = np.exp(logw); weight /= weight.sum(1,keepdims=True)
|
| 31 |
+
shape = (*ab.shape[:-1], k)
|
| 32 |
+
return idx.reshape(shape).astype(np.int64), weight.reshape(shape).astype(np.float32)
|
| 33 |
+
|
| 34 |
+
def load_table(path):
|
| 35 |
+
table = pq.read_table(path)
|
| 36 |
+
return table
|
| 37 |
+
|
| 38 |
+
def validate_manifest(path, table, manifest):
|
| 39 |
+
with open(path,'rb') as f:
|
| 40 |
+
digest=hashlib.file_digest(f,'sha256').hexdigest()
|
| 41 |
+
if digest!=manifest['dataset_sha256'] or len(table)!=manifest['rows']:
|
| 42 |
+
raise ValueError('Dataset differs from pinned split manifest')
|
| 43 |
+
groups=[list(map(int,manifest[k])) for k in ['train','validation','test']]
|
| 44 |
+
if any(not g or len(g)!=len(set(g)) or min(g)<0 or max(g)>=len(table) for g in groups):
|
| 45 |
+
raise ValueError('Invalid or repeated indices in manifest')
|
| 46 |
+
if any(set(groups[i])&set(groups[j]) for i,j in [(0,1),(0,2),(1,2)]):
|
| 47 |
+
raise ValueError('Train/validation/test overlap')
|
| 48 |
+
|
| 49 |
+
def decode_image(table, index):
|
| 50 |
+
obj = table['image'][int(index)].as_py()
|
| 51 |
+
return Image.open(io.BytesIO(obj['bytes'])).convert('RGB')
|
| 52 |
+
|
| 53 |
+
def legacy_split(n, seed=0):
|
| 54 |
+
# Exactly reproduce datasets.Dataset.train_test_split(test_size=.05, seed=0).
|
| 55 |
+
order = np.random.default_rng(seed).permutation(n)
|
| 56 |
+
nval = int(np.ceil(.05*n))
|
| 57 |
+
return order[nval:], order[:nval]
|
| 58 |
+
|
| 59 |
+
def stratified_subset(indices, labels, per_class, seed):
|
| 60 |
+
rng = np.random.default_rng(seed)
|
| 61 |
+
selected = []
|
| 62 |
+
for label in sorted(set(labels)):
|
| 63 |
+
group = np.asarray([i for i in indices if labels[i] == label])
|
| 64 |
+
selected.extend(rng.choice(group, min(per_class,len(group)), replace=False).tolist())
|
| 65 |
+
return np.asarray(selected, dtype=np.int64)
|
| 66 |
+
|
| 67 |
+
class PhotoDataset(Dataset):
|
| 68 |
+
def __init__(self, table, indices, bins, size=256, augment=False):
|
| 69 |
+
self.table, self.indices, self.bins = table, list(map(int,indices)), bins
|
| 70 |
+
self.size, self.augment = size, augment
|
| 71 |
+
|
| 72 |
+
def __len__(self):
|
| 73 |
+
return len(self.indices)
|
| 74 |
+
|
| 75 |
+
def __getitem__(self, item):
|
| 76 |
+
# Keep square preprocessing for controlled comparison with historical runs.
|
| 77 |
+
# Production inference preserves aspect ratio separately.
|
| 78 |
+
image = decode_image(self.table,self.indices[item]).resize((self.size,self.size),Image.Resampling.BILINEAR)
|
| 79 |
+
arr = np.asarray(image,dtype=np.float32)/255
|
| 80 |
+
if self.augment and torch.rand(()) < .5:
|
| 81 |
+
arr = arr[:,::-1].copy()
|
| 82 |
+
lab = rgb2lab(arr).astype(np.float32)
|
| 83 |
+
idx, weight = self.bins.encode(lab[...,1:])
|
| 84 |
+
return (torch.from_numpy((lab[...,:1]/50-1).transpose(2,0,1).copy()),
|
| 85 |
+
torch.from_numpy(lab[...,1:].transpose(2,0,1).copy()),
|
| 86 |
+
torch.from_numpy(idx),torch.from_numpy(weight))
|
| 87 |
+
|
| 88 |
+
def estimate_weights(table, indices, bins, mix=.7, limit=1000, seed=0):
|
| 89 |
+
"""Re-estimate the prior ON checkpoint bins; never replace the bin vocabulary."""
|
| 90 |
+
if not 0 <= mix <= 1:
|
| 91 |
+
raise ValueError("Rebalance mixture must be in [0,1]")
|
| 92 |
+
rng=np.random.default_rng(seed)
|
| 93 |
+
chosen=rng.choice(indices,min(limit,len(indices)),replace=False)
|
| 94 |
+
counts=np.ones(len(bins.centers),np.float64)*1e-3
|
| 95 |
+
for index in chosen:
|
| 96 |
+
image=decode_image(table,index).resize((32,32))
|
| 97 |
+
ab=rgb2lab(np.asarray(image,dtype=np.float32)/255)[...,1:]
|
| 98 |
+
# Low chroma is a heuristic, not proof that an image was originally B&W.
|
| 99 |
+
if np.linalg.norm(ab,axis=-1).mean()<3:
|
| 100 |
+
continue
|
| 101 |
+
idx,weight=bins.encode(ab)
|
| 102 |
+
np.add.at(counts,idx.ravel(),weight.ravel())
|
| 103 |
+
prior=counts/counts.sum()
|
| 104 |
+
weights=1/((1-mix)*prior+mix/len(prior))
|
| 105 |
+
weights/=np.sum(prior*weights)
|
| 106 |
+
return weights.astype(np.float32),prior.astype(np.float32)
|
| 107 |
+
|
| 108 |
+
def write_manifest(path, table, train, validation, test, dataset_sha256):
|
| 109 |
+
groups=[set(map(int,g)) for g in [train,validation,test]]
|
| 110 |
+
assert not groups[0]&groups[1] and not groups[0]&groups[2] and not groups[1]&groups[2]
|
| 111 |
+
content={'dataset':'johnowhitaker/imagenette2-320','dataset_sha256':dataset_sha256,
|
| 112 |
+
'rows':len(table),'split_seed':0,'legacy_holdout_fraction':.05,
|
| 113 |
+
'train':list(map(int,train)),'validation':list(map(int,validation)), 'test':list(map(int,test)),
|
| 114 |
+
'caveat':'Matches the recent seed-0 holdout; older upstream training exposure is not established.'}
|
| 115 |
+
Path(path).write_text(json.dumps(content,indent=2))
|
| 116 |
+
return content
|
experiments/decision-20260928-v2/source/decision_model.py
ADDED
|
@@ -0,0 +1,117 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Single / multiple coherent palettes, optional compact global-attention encoder.
|
| 2 |
+
|
| 3 |
+
Original experiment inspired by diverse colorization and multiple-choice learning;
|
| 4 |
+
not a reproduction. All deployed parameters, including the selector, are counted.
|
| 5 |
+
"""
|
| 6 |
+
import json
|
| 7 |
+
from pathlib import Path
|
| 8 |
+
import torch
|
| 9 |
+
from torch import nn
|
| 10 |
+
import torch.nn.functional as F
|
| 11 |
+
from safetensors.torch import save_file, load_file
|
| 12 |
+
from semantic_model import SemanticColorizer
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
class DecisionColorizer(SemanticColorizer):
|
| 16 |
+
def __init__(self, modes=1, backbone='mobilenet', pretrained=False,
|
| 17 |
+
encoder_revision=None, **kwargs):
|
| 18 |
+
super().__init__(pretrained=False, **kwargs)
|
| 19 |
+
self.modes = modes
|
| 20 |
+
self.backbone = backbone
|
| 21 |
+
d = self.config['width']; q = self.config['queries']
|
| 22 |
+
if backbone == 'mobilevit':
|
| 23 |
+
import timm
|
| 24 |
+
from huggingface_hub import hf_hub_download
|
| 25 |
+
name = 'mobilevitv2_075.cvnets_in1k'
|
| 26 |
+
overlay = None
|
| 27 |
+
if pretrained:
|
| 28 |
+
assert encoder_revision, 'Pin pretrained encoder revision'
|
| 29 |
+
filename = hf_hub_download('timm/'+name, 'model.safetensors', revision=encoder_revision)
|
| 30 |
+
overlay = {'file': filename, 'hf_hub_id': None}
|
| 31 |
+
self.encoder = timm.create_model(name, pretrained=pretrained,
|
| 32 |
+
pretrained_cfg_overlay=overlay, features_only=True, out_indices=(1,2,3,4))
|
| 33 |
+
self.lateral = nn.ModuleList([nn.Conv2d(c,d,1) for c in self.encoder.feature_info.channels()])
|
| 34 |
+
cfg = self.encoder.pretrained_cfg
|
| 35 |
+
self.rgb_mean.copy_(torch.tensor(cfg['mean']).view(1,3,1,1))
|
| 36 |
+
self.rgb_std.copy_(torch.tensor(cfg['std']).view(1,3,1,1))
|
| 37 |
+
if modes > 1:
|
| 38 |
+
self.mode_embeddings = nn.Parameter(torch.randn(modes,d)*.03)
|
| 39 |
+
self.mode_offsets = nn.Parameter(torch.zeros(modes,q,2))
|
| 40 |
+
# Whole-image palette alternatives; no independently sampled pixels.
|
| 41 |
+
with torch.no_grad():
|
| 42 |
+
a=torch.arange(modes)*2*torch.pi/modes
|
| 43 |
+
self.mode_offsets[:,:,0]=.025*a.cos()[:,None]
|
| 44 |
+
self.mode_offsets[:,:,1]=.025*a.sin()[:,None]
|
| 45 |
+
self.mode_score = nn.Linear(d,modes)
|
| 46 |
+
nn.init.zeros_(self.mode_score.weight);nn.init.zeros_(self.mode_score.bias)
|
| 47 |
+
self.config.update(architecture='DecisionColorizer', format_version=3,
|
| 48 |
+
modes=modes, backbone=backbone, encoder_revision=encoder_revision)
|
| 49 |
+
count = sum(p.numel() for p in self.parameters())
|
| 50 |
+
assert count < 4_000_000, count
|
| 51 |
+
|
| 52 |
+
def features(self,L):
|
| 53 |
+
h,w=L.shape[-2:]
|
| 54 |
+
x=F.pad(self.neutral_rgb(L),(0,(-w)%32,0,(-h)%32),mode='replicate')
|
| 55 |
+
x=(x-self.rgb_mean)/self.rgb_std
|
| 56 |
+
if self.backbone=='mobilevit':
|
| 57 |
+
fs=self.encoder(x)
|
| 58 |
+
else:
|
| 59 |
+
fs=[]
|
| 60 |
+
for i,layer in enumerate(self.encoder):
|
| 61 |
+
x=layer(x)
|
| 62 |
+
if i in [3,6,12,16]:fs.append(x)
|
| 63 |
+
projected=[layer(f) for layer,f in zip(self.lateral,fs)]
|
| 64 |
+
x=projected[-1]
|
| 65 |
+
for i in range(2,-1,-1):
|
| 66 |
+
x=self.refine[2-i](F.interpolate(x,size=projected[i].shape[-2:],mode='bilinear',align_corners=False)+projected[i])
|
| 67 |
+
memory=self.memory_norm(torch.cat([F.adaptive_avg_pool2d(f,(8,8)).flatten(2).transpose(1,2) for f in projected[1:]],1))
|
| 68 |
+
q=self.queries[None].expand(L.shape[0],-1,-1)
|
| 69 |
+
for block in self.query_blocks:q=block(q,memory)
|
| 70 |
+
q=self.query_norm(q)
|
| 71 |
+
masks=torch.einsum('bqd,bdhw->bqhw',q,self.pixel(x))/(self.config['width']**.5)
|
| 72 |
+
weights=F.softmax(masks.float(),dim=1)
|
| 73 |
+
residual=2*torch.tanh(self.residual(x).float())
|
| 74 |
+
return q,weights,residual,memory.mean(1)
|
| 75 |
+
|
| 76 |
+
def forward(self,L,return_details=False):
|
| 77 |
+
q,weights,residual,context=self.features(L)
|
| 78 |
+
if self.modes==1:
|
| 79 |
+
palette=80*torch.tanh(self.palette(q))[:,None]
|
| 80 |
+
scores=q.new_zeros((len(L),1))
|
| 81 |
+
else:
|
| 82 |
+
mq=q[:,None]+self.mode_embeddings[None,:,None]
|
| 83 |
+
palette=80*torch.tanh(self.palette(mq)+self.mode_offsets[None])
|
| 84 |
+
scores=self.mode_score(context)
|
| 85 |
+
outputs=torch.einsum('bqhw,bkqc->bkchw',weights,palette.float())+residual[:,None]
|
| 86 |
+
b,k,c,h,w=outputs.shape
|
| 87 |
+
full=F.interpolate(outputs.reshape(b*k,c,h,w),size=(h*4,w*4),mode='bilinear',align_corners=False)
|
| 88 |
+
full=full[...,:L.shape[-2],:L.shape[-1]].reshape(b,k,c,*L.shape[-2:])
|
| 89 |
+
default=full[torch.arange(b,device=L.device),scores.argmax(1)]
|
| 90 |
+
if return_details:
|
| 91 |
+
return {'output':default,'all':full,'scores':scores.float(),
|
| 92 |
+
'weights':weights,'palette':palette.float()}
|
| 93 |
+
return default
|
| 94 |
+
|
| 95 |
+
|
| 96 |
+
def initialize_from_v3(model,baseline):
|
| 97 |
+
old=baseline.state_dict();state=model.state_dict()
|
| 98 |
+
copied=[]
|
| 99 |
+
for k,v in old.items():
|
| 100 |
+
if model.backbone=='mobilevit' and (k.startswith(('encoder.','lateral.')) or k in ['rgb_mean','rgb_std']):continue
|
| 101 |
+
if k in state and v.shape==state[k].shape:
|
| 102 |
+
state[k]=v;copied.append(k)
|
| 103 |
+
model.load_state_dict(state,strict=True)
|
| 104 |
+
return copied
|
| 105 |
+
|
| 106 |
+
|
| 107 |
+
def save_decision(m,path):
|
| 108 |
+
p=Path(path);p.mkdir(parents=True,exist_ok=True)
|
| 109 |
+
save_file({k:v.detach().cpu().contiguous() for k,v in m.state_dict().items()},str(p/'model.safetensors'))
|
| 110 |
+
(p/'config.json').write_text(json.dumps(m.config,indent=2))
|
| 111 |
+
|
| 112 |
+
|
| 113 |
+
def load_decision(path,device='cpu'):
|
| 114 |
+
p=Path(path);cfg=json.loads((p/'config.json').read_text())
|
| 115 |
+
m=DecisionColorizer(**{k:cfg[k] for k in ['modes','backbone','width','queries','head','encoder_revision']})
|
| 116 |
+
m.load_state_dict(load_file(str(p/'model.safetensors')),strict=True)
|
| 117 |
+
return m.to(device).eval()
|
experiments/decision-20260928-v2/source/inference.py
ADDED
|
@@ -0,0 +1,111 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Aspect-preserving colourisation with a compact semantic model.
|
| 2 |
+
|
| 3 |
+
Input and output are PIL images. No teacher, critic or second learned model is
|
| 4 |
+
loaded. Original lightness and alpha are retained; out-of-gamut chroma is reduced.
|
| 5 |
+
"""
|
| 6 |
+
import json
|
| 7 |
+
import math
|
| 8 |
+
from pathlib import Path
|
| 9 |
+
import numpy as np
|
| 10 |
+
import torch
|
| 11 |
+
import torch.nn.functional as F
|
| 12 |
+
from PIL import Image, ImageOps
|
| 13 |
+
from skimage.color import rgb2lab
|
| 14 |
+
from semantic_model import load_semantic
|
| 15 |
+
|
| 16 |
+
MAX_PIXELS = 12_000_000
|
| 17 |
+
|
| 18 |
+
def load_colorizer(path, device='cpu'):
|
| 19 |
+
model = load_semantic(path, device)
|
| 20 |
+
count = sum(p.numel() for p in model.parameters())
|
| 21 |
+
if count >= 4_000_000:
|
| 22 |
+
raise ValueError('Model exceeds the four-million-parameter limit')
|
| 23 |
+
return model
|
| 24 |
+
|
| 25 |
+
def box_mean(x, r):
|
| 26 |
+
return F.avg_pool2d(x, 2*r+1, 1, r, count_include_pad=False)
|
| 27 |
+
|
| 28 |
+
def prepare(image, size=256):
|
| 29 |
+
if not isinstance(image, Image.Image):
|
| 30 |
+
raise TypeError('Expected a PIL image')
|
| 31 |
+
if not 128 <= int(size) <= 512:
|
| 32 |
+
raise ValueError('Input size must be between 128 and 512')
|
| 33 |
+
image = ImageOps.exif_transpose(image)
|
| 34 |
+
if image.width * image.height > MAX_PIXELS:
|
| 35 |
+
raise ValueError('Please resize the image to at most 12 megapixels')
|
| 36 |
+
alpha = np.asarray(image.getchannel('A')).copy() if 'A' in image.getbands() else None
|
| 37 |
+
rgb = np.asarray(image.convert('RGB'), dtype=np.float32) / 255
|
| 38 |
+
light = rgb2lab(rgb)[..., 0].astype(np.float32)
|
| 39 |
+
scale = min(int(size)/max(image.size), 1.)
|
| 40 |
+
shape = (max(8, round(image.height*scale)), max(8, round(image.width*scale)))
|
| 41 |
+
x = torch.from_numpy(light)[None,None]/50-1
|
| 42 |
+
small = F.interpolate(x, size=shape, mode='bilinear', align_corners=False, antialias=True)
|
| 43 |
+
return light, alpha, small
|
| 44 |
+
|
| 45 |
+
@torch.inference_mode()
|
| 46 |
+
def chroma_coefficients(model, small, radius=8):
|
| 47 |
+
if radius not in (0,4,8,12,16):
|
| 48 |
+
raise ValueError('Unsupported smoothing radius')
|
| 49 |
+
device = next(model.parameters()).device
|
| 50 |
+
L = small.to(device)
|
| 51 |
+
ab = model(L).float()
|
| 52 |
+
if not torch.isfinite(ab).all():
|
| 53 |
+
raise RuntimeError('Model returned non-finite colours')
|
| 54 |
+
if radius == 0:
|
| 55 |
+
return torch.zeros_like(ab).cpu(), ab.cpu()
|
| 56 |
+
guide = (L.float()+1)/2
|
| 57 |
+
mi, mp = box_mean(guide,radius), box_mean(ab,radius)
|
| 58 |
+
var = (box_mean(guide*guide,radius)-mi*mi).clamp_min(0)
|
| 59 |
+
cov = box_mean(guide*ab,radius)-mi*mp
|
| 60 |
+
a = cov/(var+.001)
|
| 61 |
+
b = mp-a*mi
|
| 62 |
+
# Coefficients are upsampled, then evaluated against original-resolution L.
|
| 63 |
+
return box_mean(a,radius).cpu(), box_mean(b,radius).cpu()
|
| 64 |
+
|
| 65 |
+
def _linear_rgb(light, ab):
|
| 66 |
+
fy = (light+16)/116
|
| 67 |
+
f = np.stack([fy+ab[...,0]/500,fy,fy-ab[...,1]/200],axis=-1)
|
| 68 |
+
xyz = np.where(f>6/29,f**3,(f-4/29)*(3*(6/29)**2))
|
| 69 |
+
xyz *= np.array([.95047,1.,1.08883],np.float32)
|
| 70 |
+
matrix = np.array([[3.24048134,-1.53715152,-.49853633],[-.96925495,1.87599,.04155593],[.05564664,-.20404134,1.05731107]],np.float32)
|
| 71 |
+
return xyz @ matrix.T
|
| 72 |
+
|
| 73 |
+
def render(light, alpha, coefficients, saturation=1.):
|
| 74 |
+
saturation = float(saturation)
|
| 75 |
+
if not math.isfinite(saturation) or not 0 <= saturation <= 1.5:
|
| 76 |
+
raise ValueError('Colour strength must be between 0 and 1.5')
|
| 77 |
+
a,b = [F.interpolate(v.float(),size=light.shape,mode='bilinear',align_corners=False)[0].permute(1,2,0).numpy() for v in coefficients]
|
| 78 |
+
ab = (a*(light[...,None]/100)+b)*saturation
|
| 79 |
+
# Binary-search chroma compression retains Lab hue and lightness.
|
| 80 |
+
linear = _linear_rgb(light,ab)
|
| 81 |
+
invalid = ((linear < -1e-5)|(linear > 1+1e-5)).any(-1)
|
| 82 |
+
if invalid.any():
|
| 83 |
+
L = light[invalid]; colors=ab[invalid];lo=np.zeros(len(L),np.float32);hi=np.ones(len(L),np.float32)
|
| 84 |
+
for _ in range(9):
|
| 85 |
+
mid=(lo+hi)/2; candidate=_linear_rgb(L,colors*mid[:,None])
|
| 86 |
+
valid=((candidate>=-1e-5)&(candidate<=1+1e-5)).all(-1)
|
| 87 |
+
lo=np.where(valid,mid,lo);hi=np.where(valid,hi,mid)
|
| 88 |
+
ab[invalid]=colors*lo[:,None]
|
| 89 |
+
linear[invalid]=_linear_rgb(L,ab[invalid])
|
| 90 |
+
linear=np.clip(linear,0,1)
|
| 91 |
+
rgb=np.where(linear<=.0031308,12.92*linear,1.055*np.power(linear,1/2.4)-.055)
|
| 92 |
+
pixels=np.uint8(np.clip(np.rint(rgb*255),0,255))
|
| 93 |
+
if alpha is not None: pixels=np.concatenate([pixels,alpha[...,None]],axis=-1)
|
| 94 |
+
return Image.fromarray(pixels)
|
| 95 |
+
|
| 96 |
+
def colorize(model, image, size=256, radius=8, saturation=1.):
|
| 97 |
+
light,alpha,small=prepare(image,size)
|
| 98 |
+
return render(light,alpha,chroma_coefficients(model,small,radius),saturation)
|
| 99 |
+
|
| 100 |
+
def main():
|
| 101 |
+
import argparse
|
| 102 |
+
parser=argparse.ArgumentParser(description='Compact photo colouriser')
|
| 103 |
+
parser.add_argument('input');parser.add_argument('output')
|
| 104 |
+
parser.add_argument('--model',default='.');parser.add_argument('--device',default='cpu')
|
| 105 |
+
parser.add_argument('--size',type=int,default=256);parser.add_argument('--saturation',type=float,default=1.)
|
| 106 |
+
args=parser.parse_args()
|
| 107 |
+
model=load_colorizer(args.model,args.device)
|
| 108 |
+
with Image.open(args.input) as image:
|
| 109 |
+
colorize(model,image,args.size,saturation=args.saturation).save(args.output)
|
| 110 |
+
|
| 111 |
+
if __name__=='__main__':main()
|
experiments/decision-20260928-v2/source/legacy_eval.py
ADDED
|
@@ -0,0 +1,150 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Round 4: matched multiscale structure and DDColor distillation experiments."""
|
| 2 |
+
import os,sys,json,time,random,hashlib,traceback,shutil
|
| 3 |
+
from pathlib import Path
|
| 4 |
+
import numpy as np
|
| 5 |
+
import cv2
|
| 6 |
+
from PIL import Image,ImageDraw
|
| 7 |
+
from skimage import data as sample_data
|
| 8 |
+
from skimage.color import rgb2lab,lab2rgb
|
| 9 |
+
import pyarrow as pa
|
| 10 |
+
import pyarrow.parquet as pq
|
| 11 |
+
import torch
|
| 12 |
+
import torch.nn.functional as F
|
| 13 |
+
from torch.utils.data import Dataset,DataLoader
|
| 14 |
+
from huggingface_hub import HfApi,hf_hub_download,PyTorchModelHubMixin
|
| 15 |
+
from model import load_model,save_model
|
| 16 |
+
from data import decode_image,ColorBins
|
| 17 |
+
from metrics import soft_ce,per_image_metrics
|
| 18 |
+
from spatial import guided_chroma
|
| 19 |
+
|
| 20 |
+
OUT=Path('round6');OUT.mkdir(exist_ok=True)
|
| 21 |
+
DURABLE=None
|
| 22 |
+
START=time.monotonic();DEADLINE=START+3.5*3600
|
| 23 |
+
DEVICE='cuda';SEED=409;STEPS=3000
|
| 24 |
+
BASE_REV='704fa80d792c3d759db91daa00b2dcfe6f0f6412'
|
| 25 |
+
BASE_HASH='0e4c417375684a044860f8af3ac3a2fb44e1a5729254ca33abda758ea71aea6e'
|
| 26 |
+
TEACHER_ID='piddnad/ddcolor_artistic'
|
| 27 |
+
torch.set_num_threads(6)
|
| 28 |
+
|
| 29 |
+
def deadline():
|
| 30 |
+
if time.monotonic()>DEADLINE:raise TimeoutError('Graceful export before remote hard timeout')
|
| 31 |
+
|
| 32 |
+
def gray_arrays(im,style='gray'):
|
| 33 |
+
im=im.resize((256,256),Image.Resampling.BILINEAR)
|
| 34 |
+
rgb=np.asarray(im,np.float32)/255
|
| 35 |
+
lab=rgb2lab(rgb).astype(np.float32)
|
| 36 |
+
if style=='film':
|
| 37 |
+
g=(rgb*np.array([.42,.45,.13],np.float32)).sum(-1)
|
| 38 |
+
g=np.clip((g**1.15-.5)*.8+.5,0,1)
|
| 39 |
+
else:g=np.asarray(im.convert('L'),np.float32)/255
|
| 40 |
+
# Feed the same neutral sRGB input to teacher and student.
|
| 41 |
+
gray=np.repeat(g[...,None],3,-1).astype(np.float32)
|
| 42 |
+
L=rgb2lab(gray)[...,0].astype(np.float32)
|
| 43 |
+
return (L[None]/50-1).copy(),lab[...,1:].transpose(2,0,1).copy(),gray
|
| 44 |
+
|
| 45 |
+
class Photos(Dataset):
|
| 46 |
+
def __init__(self,table,ids,bins=None,cache=None,augment=False,style='gray'):
|
| 47 |
+
self.table,self.ids,self.bins,self.cache,self.augment,self.style=table,ids,bins,cache,augment,style
|
| 48 |
+
def __len__(self):return len(self.ids)
|
| 49 |
+
def __getitem__(self,k):
|
| 50 |
+
L,ab,gray=gray_arrays(decode_image(self.table,self.ids[k]),self.style)
|
| 51 |
+
teach=np.array(self.cache[k],np.float32) if self.cache is not None else np.zeros((2,64,64),np.float32)
|
| 52 |
+
if self.augment:
|
| 53 |
+
if random.random()<.5:L=L[...,::-1].copy();ab=ab[...,::-1].copy();teach=teach[...,::-1].copy()
|
| 54 |
+
if random.random()<.3:L=np.clip(L*random.uniform(.85,1.15)+random.uniform(-.08,.08),-1,1)
|
| 55 |
+
if self.bins is not None:
|
| 56 |
+
idx,w=self.bins.encode(ab.transpose(1,2,0));return torch.from_numpy(L),torch.from_numpy(ab),torch.from_numpy(idx),torch.from_numpy(w),torch.from_numpy(teach)
|
| 57 |
+
return torch.from_numpy(L),torch.from_numpy(ab),torch.from_numpy(gray.transpose(2,0,1)),self.ids[k]
|
| 58 |
+
|
| 59 |
+
def patch_excess(p,t):
|
| 60 |
+
"""Excess color differences across 4–64px regions, in target-flat areas."""
|
| 61 |
+
values=[]
|
| 62 |
+
for scale in [4,16]:
|
| 63 |
+
a=F.avg_pool2d(p,scale);b=F.avg_pool2d(t,scale)
|
| 64 |
+
for offset in [1,4]:
|
| 65 |
+
for dim in [-1,-2]:
|
| 66 |
+
x=a.narrow(dim,offset,a.shape[dim]-offset)-a.narrow(dim,0,a.shape[dim]-offset)
|
| 67 |
+
y=b.narrow(dim,offset,b.shape[dim]-offset)-b.narrow(dim,0,b.shape[dim]-offset)
|
| 68 |
+
dx=x.norm(dim=1);dy=y.norm(dim=1);mask=(dy<3).float()
|
| 69 |
+
values.append(((dx-dy-1).relu()*mask).sum((1,2))/mask.sum((1,2)).clamp_min(1))
|
| 70 |
+
return sum(values)/len(values)
|
| 71 |
+
|
| 72 |
+
def structural_loss(p,t):
|
| 73 |
+
terms=[]
|
| 74 |
+
for scale in [4,16]:
|
| 75 |
+
a=F.avg_pool2d(p,scale);b=F.avg_pool2d(t,scale)
|
| 76 |
+
for offset in [1,4]:
|
| 77 |
+
for dim in [-1,-2]:
|
| 78 |
+
x=a.narrow(dim,offset,a.shape[dim]-offset)-a.narrow(dim,0,a.shape[dim]-offset)
|
| 79 |
+
y=b.narrow(dim,offset,b.shape[dim]-offset)-b.narrow(dim,0,b.shape[dim]-offset)
|
| 80 |
+
terms.append(F.smooth_l1_loss(x/10,y/10,beta=.3))
|
| 81 |
+
return sum(terms)/len(terms)
|
| 82 |
+
|
| 83 |
+
def additional_losses(p,t,teacher):
|
| 84 |
+
small=F.avg_pool2d(p,4);target=F.avg_pool2d(t,4)
|
| 85 |
+
# Teacher is fallible; reduce guidance when it conflicts strongly with ground truth.
|
| 86 |
+
confidence=torch.exp(-(teacher-target).norm(dim=1,keepdim=True)/30).detach()
|
| 87 |
+
kd=(F.smooth_l1_loss(small/20,teacher/20,beta=.5,reduction='none')*confidence).mean()
|
| 88 |
+
neutral=(t.norm(dim=1)<3).float()
|
| 89 |
+
neutral_loss=(p.norm(dim=1)*neutral).sum()/neutral.sum().clamp_min(1)/20
|
| 90 |
+
return structural_loss(p,t),kd,neutral_loss
|
| 91 |
+
|
| 92 |
+
def measure(L,p,t,raw=None):
|
| 93 |
+
r=per_image_metrics(L,p,t);C=p.norm(dim=1);T=t.norm(dim=1)
|
| 94 |
+
for name,mask,event in [('missed_color',T>12,C<5),('neutral_spill',T<3,C>10)]:
|
| 95 |
+
r[name]=(mask&event).sum((1,2))/mask.sum((1,2)).clamp_min(1)
|
| 96 |
+
r['color_coverage']=(C>10).float().mean((1,2));r['patch_excess']=patch_excess(p,t)
|
| 97 |
+
r['raw_patch_excess']=patch_excess(p if raw is None else raw,t)
|
| 98 |
+
return r
|
| 99 |
+
|
| 100 |
+
@torch.inference_mode()
|
| 101 |
+
def teacher_predict(teacher,gray):
|
| 102 |
+
gray=F.interpolate(gray.to(DEVICE),size=(512,512),mode='bilinear',align_corners=False)
|
| 103 |
+
# Full precision is deliberate: spectral normalization / attention are not assumed BF16-safe.
|
| 104 |
+
out=teacher(gray).float()
|
| 105 |
+
assert out.shape[1]==2 and torch.isfinite(out).all()
|
| 106 |
+
return F.interpolate(out,size=(256,256),mode='bilinear',align_corners=False)
|
| 107 |
+
|
| 108 |
+
@torch.inference_mode()
|
| 109 |
+
def evaluate(model,table,ids,is_teacher=False):
|
| 110 |
+
result={};model.eval()
|
| 111 |
+
for style in ['gray','film']:
|
| 112 |
+
rows=[]
|
| 113 |
+
for L,t,gray,indices in DataLoader(Photos(table,ids,style=style),batch_size=2 if is_teacher else 8,num_workers=4):
|
| 114 |
+
deadline();L,t=L.to(DEVICE),t.to(DEVICE)
|
| 115 |
+
raw=teacher_predict(model,gray) if is_teacher else model.decode(model(L),.38)
|
| 116 |
+
pred=raw if is_teacher else guided_chroma(L,raw,8)
|
| 117 |
+
metrics=measure(L,pred,t,raw)
|
| 118 |
+
for j,index in enumerate(indices):rows.append({'index':int(index)}|{k:float(v[j]) for k,v in metrics.items()})
|
| 119 |
+
keep=[x for x in rows if x['target_chroma']>=5]
|
| 120 |
+
result[style]={'n_total':len(rows),'n_color':len(keep),'summary':{k:float(np.mean([x[k] for x in keep])) for k in keep[0] if k!='index'},'per_image':rows}
|
| 121 |
+
return result
|
| 122 |
+
|
| 123 |
+
def selection(result,baseline):
|
| 124 |
+
scores=[];eligible=True
|
| 125 |
+
for style in ['gray','film']:
|
| 126 |
+
s=result[style]['summary'];b=baseline[style]['summary']
|
| 127 |
+
eligible &= (s['ab_error']<=b['ab_error']*1.05 and s['neutral_spill']<=b['neutral_spill']+.01
|
| 128 |
+
and s['missed_color']<=b['missed_color']+.01 and s['color_coverage']>=.95*b['color_coverage']
|
| 129 |
+
and s['patch_excess']<.85*b['patch_excess'] and s['raw_patch_excess']<.9*b['raw_patch_excess'])
|
| 130 |
+
scores.append(s['patch_excess']/max(b['patch_excess'],1e-6)+.3*s['neutral_spill']+.2*s['ab_error']/b['ab_error'])
|
| 131 |
+
return float(np.mean(scores)),bool(eligible)
|
| 132 |
+
|
| 133 |
+
@torch.inference_mode()
|
| 134 |
+
def probes(model,tag,teacher=False):
|
| 135 |
+
names=['astronaut','coffee','chelsea','rocket','camera','coins','moon'];images=[]
|
| 136 |
+
for name in names:
|
| 137 |
+
original=Image.fromarray(getattr(sample_data,name)()).convert('RGB')
|
| 138 |
+
L,_,gray=gray_arrays(original);x=torch.from_numpy(L)[None].to(DEVICE)
|
| 139 |
+
ab=teacher_predict(model,torch.from_numpy(gray.transpose(2,0,1))[None]) if teacher else guided_chroma(x,model.decode(model(x),.38),8)
|
| 140 |
+
lab=np.concatenate([(L[0]*50+50)[...,None],ab[0].cpu().numpy().transpose(1,2,0)],-1)
|
| 141 |
+
rgb=np.uint8(np.clip(lab2rgb(lab),0,1)*255)
|
| 142 |
+
images.append((name,Image.fromarray(np.uint8(gray*255)),Image.fromarray(rgb)))
|
| 143 |
+
canvas=Image.new('RGB',(512,len(names)*280),'white');d=ImageDraw.Draw(canvas)
|
| 144 |
+
for i,(name,gray,col) in enumerate(images):
|
| 145 |
+
canvas.paste(gray,(0,i*280+24));canvas.paste(col,(256,i*280+24));d.text((4,i*280+4),name+' input | '+tag,fill='black')
|
| 146 |
+
canvas.save(OUT/(tag+'.jpg'),quality=90)
|
| 147 |
+
|
| 148 |
+
def write(name,obj):
|
| 149 |
+
(OUT/name).write_text(json.dumps(obj,indent=2));return obj
|
| 150 |
+
|
experiments/decision-20260928-v2/source/metrics.py
ADDED
|
@@ -0,0 +1,35 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Metrics are fidelity/artifact proxies; none certifies plausible color by itself."""
|
| 2 |
+
import numpy as np
|
| 3 |
+
import torch
|
| 4 |
+
|
| 5 |
+
def soft_ce(logits, idx, weights, class_weights=None):
|
| 6 |
+
# FP32 reduction under autocast; no full target distribution materialized.
|
| 7 |
+
logits=logits.float()
|
| 8 |
+
idx=idx.permute(0,3,1,2)
|
| 9 |
+
weights=weights.permute(0,3,1,2)
|
| 10 |
+
nll=-((logits.gather(1,idx)-logits.logsumexp(1,keepdim=True))*weights).sum(1)
|
| 11 |
+
if class_weights is not None:
|
| 12 |
+
nll=nll*class_weights[idx[:,0]]
|
| 13 |
+
return nll.mean()
|
| 14 |
+
|
| 15 |
+
def per_image_metrics(L, pred, target):
|
| 16 |
+
error=(pred-target).square().sum(1).sqrt().mean((1,2))
|
| 17 |
+
chroma=pred.square().sum(1).sqrt().mean((1,2))
|
| 18 |
+
true_chroma=target.square().sum(1).sqrt().mean((1,2))
|
| 19 |
+
seams=[]
|
| 20 |
+
for dim in [-1,-2]:
|
| 21 |
+
dp=pred.diff(dim=dim).square().sum(1).sqrt()
|
| 22 |
+
dt=target.diff(dim=dim).square().sum(1).sqrt()
|
| 23 |
+
dl=L.diff(dim=dim).abs()[:,0]*50
|
| 24 |
+
# Penalize excess chroma discontinuity where luminance is flat.
|
| 25 |
+
mask=(dl<2).float()
|
| 26 |
+
seams.append(((dp-dt).relu()*mask).sum((1,2))/mask.sum((1,2)).clamp_min(1))
|
| 27 |
+
return {'ab_error':error,'chroma':chroma,'target_chroma':true_chroma,
|
| 28 |
+
'excess_chroma_edge':sum(seams)/2}
|
| 29 |
+
|
| 30 |
+
def summarize(rows):
|
| 31 |
+
keys=['ce','ab_error','chroma','target_chroma','excess_chroma_edge']
|
| 32 |
+
result={k:float(np.mean([r[k] for r in rows])) for k in keys}
|
| 33 |
+
result['n']=len(rows)
|
| 34 |
+
result['chroma_ratio']=result['chroma']/max(result['target_chroma'],1e-9)
|
| 35 |
+
return result
|
experiments/decision-20260928-v2/source/model.py
ADDED
|
@@ -0,0 +1,122 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Shared, checkpoint-compatible Mini U-Net. Parameters remain under 4M."""
|
| 2 |
+
import json
|
| 3 |
+
import math
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
import numpy as np
|
| 6 |
+
import torch
|
| 7 |
+
from torch import nn
|
| 8 |
+
import torch.nn.functional as F
|
| 9 |
+
from huggingface_hub import PyTorchModelHubMixin, snapshot_download
|
| 10 |
+
from safetensors.torch import load_file, save_file
|
| 11 |
+
def double_conv(in_ch, out_ch):
|
| 12 |
+
return nn.Sequential(
|
| 13 |
+
nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False),
|
| 14 |
+
nn.BatchNorm2d(out_ch),
|
| 15 |
+
nn.ReLU(inplace=True),
|
| 16 |
+
nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False),
|
| 17 |
+
nn.BatchNorm2d(out_ch),
|
| 18 |
+
nn.ReLU(inplace=True),
|
| 19 |
+
)
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
class DilatedContextBlock(nn.Module):
|
| 23 |
+
def __init__(self, channels, mid_ch=96, dilations=(2, 4, 8)):
|
| 24 |
+
super().__init__()
|
| 25 |
+
self.proj_in = nn.Sequential(
|
| 26 |
+
nn.Conv2d(channels, mid_ch, 1, bias=False),
|
| 27 |
+
nn.BatchNorm2d(mid_ch),
|
| 28 |
+
nn.ReLU(inplace=True),
|
| 29 |
+
)
|
| 30 |
+
layers = []
|
| 31 |
+
for d in dilations:
|
| 32 |
+
layers += [
|
| 33 |
+
nn.Conv2d(mid_ch, mid_ch, 3, padding=d, dilation=d, bias=False),
|
| 34 |
+
nn.BatchNorm2d(mid_ch),
|
| 35 |
+
nn.ReLU(inplace=True),
|
| 36 |
+
]
|
| 37 |
+
self.dilated = nn.Sequential(*layers)
|
| 38 |
+
self.proj_out = nn.Sequential(
|
| 39 |
+
nn.Conv2d(mid_ch, channels, 1, bias=False),
|
| 40 |
+
nn.BatchNorm2d(channels),
|
| 41 |
+
)
|
| 42 |
+
nn.init.zeros_(self.proj_out[-1].weight)
|
| 43 |
+
self.relu = nn.ReLU(inplace=True)
|
| 44 |
+
|
| 45 |
+
def forward(self, x):
|
| 46 |
+
y = self.proj_in(x)
|
| 47 |
+
y = self.dilated(y)
|
| 48 |
+
y = self.proj_out(y)
|
| 49 |
+
return self.relu(x + y)
|
| 50 |
+
|
| 51 |
+
|
| 52 |
+
class SmallUNetColorizer(
|
| 53 |
+
nn.Module,
|
| 54 |
+
PyTorchModelHubMixin,
|
| 55 |
+
pipeline_tag="image-to-image",
|
| 56 |
+
license="apache-2.0",
|
| 57 |
+
tags=["colorization", "unet", "image-to-image", "classification"],
|
| 58 |
+
):
|
| 59 |
+
def __init__(self, bin_centers, in_ch: int = 1, base: int = 44,
|
| 60 |
+
context_mid_ch: int = 96, context_dilations=(2, 4, 8)):
|
| 61 |
+
super().__init__()
|
| 62 |
+
self.in_ch, self.base = in_ch, base
|
| 63 |
+
num_bins = len(bin_centers)
|
| 64 |
+
self.num_bins = num_bins
|
| 65 |
+
self.register_buffer("bin_centers", torch.tensor(bin_centers, dtype=torch.float32))
|
| 66 |
+
|
| 67 |
+
self.enc1 = double_conv(in_ch, base)
|
| 68 |
+
self.enc2 = double_conv(base, base * 2)
|
| 69 |
+
self.enc3 = double_conv(base * 2, base * 4)
|
| 70 |
+
self.enc4 = double_conv(base * 4, base * 8)
|
| 71 |
+
self.pool = nn.MaxPool2d(2)
|
| 72 |
+
self.context = DilatedContextBlock(base * 8, mid_ch=context_mid_ch,
|
| 73 |
+
dilations=tuple(context_dilations))
|
| 74 |
+
self.up3 = nn.ConvTranspose2d(base * 8, base * 4, 2, stride=2)
|
| 75 |
+
self.dec3 = double_conv(base * 8, base * 4)
|
| 76 |
+
self.up2 = nn.ConvTranspose2d(base * 4, base * 2, 2, stride=2)
|
| 77 |
+
self.dec2 = double_conv(base * 4, base * 2)
|
| 78 |
+
self.up1 = nn.ConvTranspose2d(base * 2, base, 2, stride=2)
|
| 79 |
+
self.dec1 = double_conv(base * 2, base)
|
| 80 |
+
self.out_conv = nn.Conv2d(base, num_bins, 1)
|
| 81 |
+
|
| 82 |
+
def forward(self, x):
|
| 83 |
+
h, w = x.shape[-2:]
|
| 84 |
+
x = F.pad(x, (0, (-w) % 8, 0, (-h) % 8), mode="replicate")
|
| 85 |
+
e1 = self.enc1(x)
|
| 86 |
+
e2 = self.enc2(self.pool(e1))
|
| 87 |
+
e3 = self.enc3(self.pool(e2))
|
| 88 |
+
e4 = self.context(self.enc4(self.pool(e3)))
|
| 89 |
+
d3 = self.dec3(torch.cat([self.up3(e4), e3], dim=1))
|
| 90 |
+
d2 = self.dec2(torch.cat([self.up2(d3), e2], dim=1))
|
| 91 |
+
d1 = self.dec1(torch.cat([self.up1(d2), e1], dim=1))
|
| 92 |
+
return self.out_conv(d1)[..., :h, :w]
|
| 93 |
+
|
| 94 |
+
def decode(self, logits, temperature: float = 0.38):
|
| 95 |
+
if not math.isfinite(temperature) or temperature <= 0:
|
| 96 |
+
raise ValueError("temperature must be finite and positive")
|
| 97 |
+
probs_t = F.softmax(logits.float() / temperature, dim=1)
|
| 98 |
+
return torch.einsum("bqhw,qc->bchw", probs_t, self.bin_centers)
|
| 99 |
+
|
| 100 |
+
|
| 101 |
+
def load_model(source, revision=None, device="cpu"):
|
| 102 |
+
path = Path(source)
|
| 103 |
+
if not path.is_dir():
|
| 104 |
+
path = Path(snapshot_download(source, revision=revision,
|
| 105 |
+
allow_patterns=["config.json", "model.safetensors"]))
|
| 106 |
+
config = json.loads((path / "config.json").read_text())
|
| 107 |
+
state = load_file(str(path / "model.safetensors"))
|
| 108 |
+
centers = torch.tensor(config["bin_centers"], dtype=torch.float32)
|
| 109 |
+
if not torch.equal(centers, state["bin_centers"]):
|
| 110 |
+
raise ValueError("Checkpoint config and state color bins differ; refusing ambiguous decode")
|
| 111 |
+
model = SmallUNetColorizer(**config)
|
| 112 |
+
model.load_state_dict(state, strict=True)
|
| 113 |
+
model.to(device).eval()
|
| 114 |
+
return model
|
| 115 |
+
|
| 116 |
+
def save_model(model, path):
|
| 117 |
+
path = Path(path); path.mkdir(parents=True, exist_ok=True)
|
| 118 |
+
model.save_pretrained(path)
|
| 119 |
+
# Mixin config can retain constructor bins; use the actual authoritative buffer.
|
| 120 |
+
cfg = json.loads((path / "config.json").read_text())
|
| 121 |
+
cfg["bin_centers"] = model.bin_centers.detach().cpu().tolist()
|
| 122 |
+
(path / "config.json").write_text(json.dumps(cfg, indent=2))
|
experiments/decision-20260928-v2/source/objectives.py
ADDED
|
@@ -0,0 +1,17 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import math
|
| 2 |
+
import torch
|
| 3 |
+
import torch.nn.functional as F
|
| 4 |
+
|
| 5 |
+
def color_loss(pred,target):
|
| 6 |
+
"""Fit one coherent target; distribution matching discourages gray averages."""
|
| 7 |
+
pred=F.adaptive_avg_pool2d(pred,(64,64)).float();target=target.float()
|
| 8 |
+
pixel=F.l1_loss(pred/20,target/20)
|
| 9 |
+
p=F.avg_pool2d(pred,2).flatten(2);t=F.avg_pool2d(target,2).flatten(2)
|
| 10 |
+
angle=torch.arange(8,device=p.device,dtype=p.dtype)*math.pi/8
|
| 11 |
+
directions=torch.stack([angle.cos(),angle.sin()],1)
|
| 12 |
+
a=torch.einsum('kc,bcn->bkn',directions,p).sort(-1).values
|
| 13 |
+
b=torch.einsum('kc,bcn->bkn',directions,t).sort(-1).values
|
| 14 |
+
distribution=F.l1_loss(a/20,b/20)
|
| 15 |
+
# Match boundaries in the teacher target, never minimize gradients toward zero.
|
| 16 |
+
edge=(F.l1_loss((pred[:,:,:,1:]-pred[:,:,:,:-1])/20,(target[:,:,:,1:]-target[:,:,:,:-1])/20)+F.l1_loss((pred[:,:,1:]-pred[:,:,:-1])/20,(target[:,:,1:]-target[:,:,:-1])/20))/2
|
| 17 |
+
return pixel+.25*distribution+.1*edge,{'pixel':pixel,'distribution':distribution,'edge':edge}
|
experiments/decision-20260928-v2/source/persistence.py
ADDED
|
@@ -0,0 +1,42 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Commit run artifacts atomically and verify every changed file by SHA256."""
|
| 2 |
+
import hashlib, json, os, time
|
| 3 |
+
from pathlib import Path
|
| 4 |
+
from huggingface_hub import HfApi, CommitOperationAdd, hf_hub_download
|
| 5 |
+
|
| 6 |
+
REPO = 'User-2468/mini-unet-colorizer'
|
| 7 |
+
PREFIX = 'experiments/round6-20260928'
|
| 8 |
+
|
| 9 |
+
class DurableRun:
|
| 10 |
+
def __init__(self, root):
|
| 11 |
+
if not os.environ.get('HF_TOKEN'):
|
| 12 |
+
raise RuntimeError('A write token is required before starting')
|
| 13 |
+
self.api = HfApi(token=os.environ['HF_TOKEN'])
|
| 14 |
+
if self.api.whoami()['name'] != 'User-2468':
|
| 15 |
+
raise RuntimeError('Unexpected account')
|
| 16 |
+
self.root = Path(root)
|
| 17 |
+
self.saved = {}
|
| 18 |
+
|
| 19 |
+
def sync(self, reason):
|
| 20 |
+
files = [p for p in sorted(self.root.rglob('*')) if p.is_file()]
|
| 21 |
+
hashes = {p.relative_to(self.root).as_posix(): hashlib.sha256(p.read_bytes()).hexdigest() for p in files}
|
| 22 |
+
changed = [p for p in files if self.saved.get(p.relative_to(self.root).as_posix()) != hashes[p.relative_to(self.root).as_posix()]]
|
| 23 |
+
if not changed:
|
| 24 |
+
return
|
| 25 |
+
for attempt in range(3):
|
| 26 |
+
try:
|
| 27 |
+
head = self.api.model_info(REPO, revision='main').sha
|
| 28 |
+
operations = [CommitOperationAdd(path_in_repo=PREFIX+'/'+p.relative_to(self.root).as_posix(), path_or_fileobj=str(p)) for p in changed]
|
| 29 |
+
operations.append(CommitOperationAdd(path_in_repo=PREFIX+'/SHA256SUMS.json', path_or_fileobj=json.dumps(hashes,sort_keys=True,indent=2).encode()))
|
| 30 |
+
commit = self.api.create_commit(repo_id=REPO, revision='main', parent_commit=head, operations=operations, commit_message='Colorizer round6: '+reason)
|
| 31 |
+
for p in changed:
|
| 32 |
+
name = p.relative_to(self.root).as_posix()
|
| 33 |
+
downloaded = hf_hub_download(REPO, PREFIX+'/'+name, revision=commit.oid, force_download=True, token=os.environ['HF_TOKEN'])
|
| 34 |
+
if hashlib.sha256(Path(downloaded).read_bytes()).hexdigest() != hashes[name]:
|
| 35 |
+
raise RuntimeError('Remote checksum mismatch: '+name)
|
| 36 |
+
self.saved = hashes
|
| 37 |
+
print('PERSISTED',reason,commit.oid,'verified_files',len(changed),flush=True)
|
| 38 |
+
return
|
| 39 |
+
except Exception:
|
| 40 |
+
if attempt == 2:
|
| 41 |
+
raise
|
| 42 |
+
time.sleep(2**attempt)
|
experiments/decision-20260928-v2/source/r7.py
ADDED
|
@@ -0,0 +1,234 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import os,sys,json,random,time,hashlib,shutil,math,traceback,io
|
| 2 |
+
from pathlib import Path
|
| 3 |
+
import numpy as np
|
| 4 |
+
import torch
|
| 5 |
+
import torch.nn.functional as F
|
| 6 |
+
from torch.utils.data import Dataset,DataLoader
|
| 7 |
+
from huggingface_hub import hf_hub_download
|
| 8 |
+
import pyarrow as pa
|
| 9 |
+
import pyarrow.parquet as pq
|
| 10 |
+
from PIL import Image,ImageDraw
|
| 11 |
+
from skimage.color import rgb2lab,lab2rgb
|
| 12 |
+
from skimage.segmentation import slic
|
| 13 |
+
import cv2
|
| 14 |
+
import persistence
|
| 15 |
+
from semantic_model import load_semantic,save_semantic
|
| 16 |
+
from objectives import color_loss
|
| 17 |
+
from data import decode_image
|
| 18 |
+
import legacy_eval as ev
|
| 19 |
+
REPO='User-2468/mini-unet-colorizer'
|
| 20 |
+
REV='7b10485ee9d205ce2f7b660dfc6a90012d78b1cb'
|
| 21 |
+
OLD='experiments/round6-20260928'
|
| 22 |
+
OUT=Path('/work/results');OUT.mkdir(parents=True,exist_ok=True)
|
| 23 |
+
persistence.PREFIX='experiments/round7-20260928'
|
| 24 |
+
ev.OUT=OUT
|
| 25 |
+
START=time.monotonic()
|
| 26 |
+
SEED=709
|
| 27 |
+
STEPS=2000
|
| 28 |
+
DEVICE='cpu' if os.environ.get('SMOKE')=='1' else 'cuda'
|
| 29 |
+
ev.DEVICE=DEVICE
|
| 30 |
+
DURABLE=None
|
| 31 |
+
def write(name,obj):
|
| 32 |
+
p=OUT/name;p.parent.mkdir(parents=True,exist_ok=True);p.write_text(json.dumps(obj,indent=2));return obj
|
| 33 |
+
def deadline():
|
| 34 |
+
if time.monotonic()-START>9000:raise TimeoutError('Graceful deadline')
|
| 35 |
+
ev.deadline=deadline
|
| 36 |
+
def losses(p,t):
|
| 37 |
+
p=F.adaptive_avg_pool2d(p.float(),(64,64));t=t.float()
|
| 38 |
+
pixel=(p-t).abs().mean((1,2,3))/20
|
| 39 |
+
a=F.avg_pool2d(p,2).flatten(2);b=F.avg_pool2d(t,2).flatten(2)
|
| 40 |
+
angles=torch.arange(8,device=p.device)*math.pi/8
|
| 41 |
+
directions=torch.stack([angles.cos(),angles.sin()],1)
|
| 42 |
+
ap=torch.einsum('kc,bcn->bkn',directions,a).sort(-1).values
|
| 43 |
+
bp=torch.einsum('kc,bcn->bkn',directions,b).sort(-1).values
|
| 44 |
+
distribution=(ap-bp).abs().mean((1,2))/20
|
| 45 |
+
edge=((torch.diff(p,dim=-1)-torch.diff(t,dim=-1)).abs().mean((1,2,3))+(torch.diff(p,dim=-2)-torch.diff(t,dim=-2)).abs().mean((1,2,3)))/40
|
| 46 |
+
return pixel+.25*distribution+.1*edge
|
| 47 |
+
def select_loss(p,teacher,original,mode):
|
| 48 |
+
a=losses(p,teacher)
|
| 49 |
+
if mode=='teacher':return a.mean(),torch.zeros_like(a,dtype=torch.bool)
|
| 50 |
+
b=losses(p,original);choose=b<a
|
| 51 |
+
return torch.where(choose,b,a).mean(),choose
|
| 52 |
+
def smoke():
|
| 53 |
+
torch.set_num_threads(4);torch.manual_seed(SEED)
|
| 54 |
+
p=torch.randn(2,2,64,64,requires_grad=True);t=torch.randn_like(p)
|
| 55 |
+
assert torch.allclose(losses(p,t).mean(),color_loss(p,t)[0],atol=1e-6)
|
| 56 |
+
# Whole-image choice: exact red or blue is better than an average or a patchwork.
|
| 57 |
+
red=torch.zeros(1,2,64,64);red[:,0]=30
|
| 58 |
+
blue=-red;mix=red.clone();mix[:,:,:,32:]=blue[:,:,:,32:]
|
| 59 |
+
assert select_loss(red,red,blue,'set')[0]==0
|
| 60 |
+
assert select_loss(torch.zeros_like(red),red,blue,'set')[0]>0
|
| 61 |
+
assert select_loss(mix,red,blue,'set')[0]>0
|
| 62 |
+
m=load_semantic('/work/initial',DEVICE)
|
| 63 |
+
assert sum(p.numel() for p in m.parameters())==3994676
|
| 64 |
+
x=torch.rand(2,1,96,128,device=DEVICE)*2-1
|
| 65 |
+
pred=m(x);loss,_=select_loss(pred,t.to(DEVICE),-t.to(DEVICE),'set')
|
| 66 |
+
loss.backward()
|
| 67 |
+
assert torch.isfinite(loss) and all(torch.isfinite(p.grad).all() for p in m.parameters() if p.grad is not None)
|
| 68 |
+
save_semantic(m,OUT/'initial')
|
| 69 |
+
re=load_semantic(OUT/'initial',DEVICE)
|
| 70 |
+
with torch.no_grad():assert torch.allclose(m.eval()(x),re(x),atol=1e-5)
|
| 71 |
+
write('smoke.json',{'passed':True,'parameters':3994676,'loss_equivalence':True,'whole_image_choice':True,'finite_backward':True,'reload':True})
|
| 72 |
+
del m,re
|
| 73 |
+
class Photos(Dataset):
|
| 74 |
+
def __init__(self,table,ids,cache,archive):
|
| 75 |
+
self.table,self.ids,self.cache,self.archive=table,ids,cache,archive;self.epoch=0
|
| 76 |
+
def __len__(self):return len(self.ids)
|
| 77 |
+
def __getitem__(self,k):
|
| 78 |
+
im=decode_image(self.table,self.ids[k]);L,ab,gray=ev.gray_arrays(im)
|
| 79 |
+
# Independent per-example RNG; augmentation differences never alter image ordering/flips.
|
| 80 |
+
rng=np.random.default_rng(np.random.SeedSequence([SEED,self.epoch,k]))
|
| 81 |
+
target=np.array(self.cache[k],np.float32)
|
| 82 |
+
if self.archive and rng.random()<.5:
|
| 83 |
+
rgb=np.asarray(im.resize((256,256),Image.Resampling.BILINEAR),np.float32)/255
|
| 84 |
+
weights=rng.dirichlet(np.array([3.,5.,2.]))
|
| 85 |
+
g=(rgb*weights).sum(-1)
|
| 86 |
+
g=np.clip((g**rng.uniform(.7,1.5)-.5)*rng.uniform(.6,1.1)+.5+rng.uniform(-.08,.08),0,1)
|
| 87 |
+
if rng.random()<.5:g=cv2.GaussianBlur(g,(5,5),float(rng.uniform(.3,1.2)))
|
| 88 |
+
if rng.random()<.6:g=np.clip(g+rng.normal(0,rng.uniform(.003,.025),g.shape),0,1)
|
| 89 |
+
L=(rgb2lab(np.repeat(g[...,None],3,-1).astype(np.float32))[...,0][None]/50-1).astype(np.float32)
|
| 90 |
+
else:
|
| 91 |
+
if rng.random()<.5:L=np.clip(L*rng.uniform(.7,1.15)+rng.uniform(-.12,.12),-1,1)
|
| 92 |
+
flip=np.random.default_rng(np.random.SeedSequence([SEED+1,self.epoch,k])).random()<.5
|
| 93 |
+
if flip:L=L[...,::-1].copy();ab=ab[...,::-1].copy();target=target[...,::-1].copy()
|
| 94 |
+
original=cv2.resize(ab.transpose(1,2,0),(64,64),interpolation=cv2.INTER_AREA).transpose(2,0,1).copy()
|
| 95 |
+
return torch.from_numpy(L),torch.from_numpy(target),torch.from_numpy(original)
|
| 96 |
+
@torch.inference_mode()
|
| 97 |
+
def predict(m,L):
|
| 98 |
+
x=torch.from_numpy(L)[None].to(DEVICE);raw=m(x)
|
| 99 |
+
guided=ev.guided_chroma(x,raw,8)
|
| 100 |
+
return raw,guided
|
| 101 |
+
def toimage(L,ab):
|
| 102 |
+
rgb=lab2rgb(np.concatenate([(L[0]*50+50)[...,None],ab],-1))
|
| 103 |
+
return Image.fromarray(np.uint8(np.clip(rgb,0,1)*255))
|
| 104 |
+
@torch.inference_mode()
|
| 105 |
+
def review(m,table,ids,name):
|
| 106 |
+
m.eval();rows=[]
|
| 107 |
+
# Reference-derived surface regions for diagnostics only, never inference.
|
| 108 |
+
for i in ids:
|
| 109 |
+
L,target,_=ev.gray_arrays(decode_image(table,i))
|
| 110 |
+
raw,pred=predict(m,L);x=torch.from_numpy(L)[None].to(DEVICE)
|
| 111 |
+
resize=F.interpolate(x,(224,224),mode='bilinear',align_corners=False)
|
| 112 |
+
resized=F.interpolate(m(resize),(256,256),mode='bilinear',align_corners=False)
|
| 113 |
+
contrast=m((x*.95+.02).clamp(-1,1))
|
| 114 |
+
pab=pred[0].cpu().numpy().transpose(1,2,0)
|
| 115 |
+
rgb=np.asarray(decode_image(table,i).resize((256,256)),np.float32)/255
|
| 116 |
+
seg=slic(rgb,n_segments=100,compactness=10,start_label=0)
|
| 117 |
+
tab=target.transpose(1,2,0);variations=[]
|
| 118 |
+
for label in np.unique(seg):
|
| 119 |
+
mask=seg==label
|
| 120 |
+
if mask.sum()<64:continue
|
| 121 |
+
ground=tab[mask];center=np.median(ground,axis=0)
|
| 122 |
+
if np.percentile(np.linalg.norm(ground-center,axis=1),90)>5:continue
|
| 123 |
+
col=pab[mask];variations.append(float(np.mean(np.linalg.norm(col-np.median(col,axis=0),axis=1))))
|
| 124 |
+
rows.append({'id':int(i),'resize_ab_change':float((raw-resized).abs().mean()),'contrast_ab_change':float((raw-contrast).abs().mean()),'surface_chroma_variation':float(np.mean(variations)) if variations else None,'chroma_coverage':float((pred.norm(dim=1)>10).float().mean()),'mean_chroma':float(pred.norm(dim=1).mean())})
|
| 125 |
+
write(name+'_diagnostics.json',{'rows':rows,'summary':{k:float(np.mean([r[k] for r in rows if r[k] is not None])) for k in rows[0] if k!='id'},'warning':'Surface uniformity alone rewards grey collapse; never use alone to select model. Color differences from original are diagnostic only.'})
|
| 126 |
+
return rows
|
| 127 |
+
@torch.inference_mode()
|
| 128 |
+
def grid(models,table,ids,label):
|
| 129 |
+
order=np.random.default_rng(SEED+2).permutation(len(models))
|
| 130 |
+
names=list(models);key={chr(65+j):names[int(i)] for j,i in enumerate(order)}
|
| 131 |
+
write(label+'_key.json',key)
|
| 132 |
+
for page in range(0,min(len(ids),32),8):
|
| 133 |
+
chosen=ids[page:page+8];canvas=Image.new('RGB',(256*(len(models)+1),274*len(chosen)),'white');draw=ImageDraw.Draw(canvas)
|
| 134 |
+
for row,index in enumerate(chosen):
|
| 135 |
+
L,_,gray=ev.gray_arrays(decode_image(table,index));draw.text((2,row*274),f'{index}: gray | '+' | '.join(key),fill='black')
|
| 136 |
+
canvas.paste(Image.fromarray(np.uint8(gray*255)),(0,row*274+18))
|
| 137 |
+
for j,i in enumerate(order):
|
| 138 |
+
_,p=predict(models[names[int(i)]],L)
|
| 139 |
+
canvas.paste(toimage(L,p[0].cpu().numpy().transpose(1,2,0)),((j+1)*256,row*274+18))
|
| 140 |
+
canvas.save(OUT/f'{label}_{page//8}.jpg',quality=92)
|
| 141 |
+
def main():
|
| 142 |
+
global DURABLE
|
| 143 |
+
torch.set_num_threads(6);torch.manual_seed(SEED);random.seed(SEED);np.random.seed(SEED)
|
| 144 |
+
DURABLE=persistence.DurableRun(OUT)
|
| 145 |
+
shutil.copy2('/work/runner.py',OUT/'runner.py');shutil.copy2('/work/PROTOCOL.md',OUT/'PROTOCOL.md')
|
| 146 |
+
shutil.copytree('/work/source',OUT/'source',ignore=shutil.ignore_patterns('__pycache__'),dirs_exist_ok=True)
|
| 147 |
+
smoke();DURABLE.sync('round7 initial checkpoint and smoke verified')
|
| 148 |
+
if DEVICE=='cpu':return
|
| 149 |
+
torch.backends.cudnn.benchmark=True
|
| 150 |
+
tables=[]
|
| 151 |
+
for repo,rev,files in [('johnowhitaker/imagenette2-320','771c1310a2487e8076ede6b7d6307244aa8400af',['default/train/0000.parquet']),('detection-datasets/coco','26ddc382fe75dfc2a0655b5977e296ea10efebce',['default/train/0000.parquet','default/train/0001.parquet'])]:
|
| 152 |
+
for f in files:tables.append(pq.read_table(hf_hub_download(repo,f,repo_type='dataset',revision=rev),columns=['image']))
|
| 153 |
+
table=pa.concat_tables(tables);manifest=json.loads(Path('/work/source/previous_manifest.json').read_text());ids=manifest['train']
|
| 154 |
+
cache=np.empty((len(ids),2,64,64),np.float16)
|
| 155 |
+
sums=json.loads(Path(hf_hub_download(REPO,OLD+'/SHA256SUMS.json',revision=REV)).read_text())
|
| 156 |
+
for start in range(0,len(ids),2048):
|
| 157 |
+
f=f'teacher_cache/{start:05d}.npz';path=Path(hf_hub_download(REPO,OLD+'/'+f,revision=REV))
|
| 158 |
+
assert hashlib.sha256(path.read_bytes()).hexdigest()==sums[f]
|
| 159 |
+
with np.load(path) as z:
|
| 160 |
+
assert z['indices'].tolist()==ids[start:start+len(z['indices'])]
|
| 161 |
+
cache[start:start+len(z['target'])]=z['target']
|
| 162 |
+
fresh=pq.read_table(hf_hub_download('detection-datasets/coco','default/val/0001.parquet',repo_type='dataset',revision='26ddc382fe75dfc2a0655b5977e296ea10efebce'),columns=['image'])
|
| 163 |
+
hashes={hashlib.sha256(table['image'][int(i)].as_py()['bytes']).digest() for i in ids}
|
| 164 |
+
freshids=[]
|
| 165 |
+
for i in np.random.default_rng(SEED).permutation(len(fresh)):
|
| 166 |
+
if hashlib.sha256(fresh['image'][int(i)].as_py()['bytes']).digest() not in hashes:freshids.append(int(i))
|
| 167 |
+
if len(freshids)==160:break
|
| 168 |
+
write('manifest.json',{'train':ids,'development':freshids[:80],'heldout':freshids[80:],'fresh_file':'default/val/0001.parquet','revision':'26ddc382fe75dfc2a0655b5977e296ea10efebce','prior_model_revision':REV,'upstream_pretraining_overlap_unknown':True,'near_duplicates_not_excluded':True})
|
| 169 |
+
DURABLE.sync('data cache verified and split fixed')
|
| 170 |
+
m=load_semantic('/work/initial',DEVICE)
|
| 171 |
+
write('initial_validation.json',ev.evaluate(m,table,manifest['validation']))
|
| 172 |
+
review(m,fresh,freshids[:80],'initial_development')
|
| 173 |
+
review(m,fresh,freshids[80:],'initial_heldout')
|
| 174 |
+
write('initial_heldout_reference_metrics.json',ev.evaluate(m,fresh,freshids[80:]))
|
| 175 |
+
del m
|
| 176 |
+
arms=[('teacher_clean','teacher',False),('set_clean','set',False),('teacher_archive','teacher',True),('set_archive','set',True)]
|
| 177 |
+
for name,mode,archive in arms:
|
| 178 |
+
torch.manual_seed(SEED);np.random.seed(SEED);random.seed(SEED)
|
| 179 |
+
m=load_semantic('/work/initial',DEVICE)
|
| 180 |
+
opt=torch.optim.AdamW([{'params':m.encoder.parameters(),'lr':5e-6},{'params':[p for n,p in m.named_parameters() if not n.startswith('encoder.')],'lr':5e-5}],weight_decay=.01)
|
| 181 |
+
sched=torch.optim.lr_scheduler.LambdaLR(opt,lambda step:min((step+1)/100,1)*(.2+.8*(1+math.cos(math.pi*min(step,STEPS)/STEPS))/2))
|
| 182 |
+
ds=Photos(table,ids,cache,archive)
|
| 183 |
+
loader=DataLoader(ds,batch_size=32,shuffle=True,num_workers=4,pin_memory=True,generator=torch.Generator().manual_seed(SEED))
|
| 184 |
+
it=iter(loader);history=[];choices=0;seen=0
|
| 185 |
+
for step in range(1,STEPS+1):
|
| 186 |
+
deadline();m.train()
|
| 187 |
+
for layer in m.encoder.modules():
|
| 188 |
+
if isinstance(layer,torch.nn.BatchNorm2d):layer.eval()
|
| 189 |
+
try:L,t,o=next(it)
|
| 190 |
+
except StopIteration:ds.epoch+=1;it=iter(loader);L,t,o=next(it)
|
| 191 |
+
L,t,o=L.to(DEVICE),t.to(DEVICE),o.to(DEVICE);opt.zero_grad(set_to_none=True)
|
| 192 |
+
with torch.autocast('cuda',dtype=torch.bfloat16):p=m(L)
|
| 193 |
+
loss,choice=select_loss(p,t,o,mode)
|
| 194 |
+
assert torch.isfinite(loss)
|
| 195 |
+
loss.backward();torch.nn.utils.clip_grad_norm_(m.parameters(),1,error_if_nonfinite=True);opt.step();sched.step()
|
| 196 |
+
choices+=int(choice.sum());seen+=len(L)
|
| 197 |
+
if step%100==0:
|
| 198 |
+
history.append({'step':step,'loss':float(loss),'original_choice_fraction':choices/seen})
|
| 199 |
+
print('STEP',name,step,round(float(loss),4),flush=True)
|
| 200 |
+
if step%500==0:
|
| 201 |
+
path=OUT/name/f'step_{step:05d}';save_semantic(m,path)
|
| 202 |
+
torch.save({'optimizer':opt.state_dict(),'scheduler':sched.state_dict(),'step':step,'epoch':ds.epoch,'torch_rng':torch.get_rng_state(),'cuda_rng':torch.cuda.get_rng_state_all(),'numpy_rng':np.random.get_state(),'python_rng':random.getstate()},path/'training_state.pt')
|
| 203 |
+
write(name+'/history.json',history);write('status.json',{'status':'training','arm':name,'step':step,'elapsed_seconds':time.monotonic()-START})
|
| 204 |
+
DURABLE.sync(name+' checkpoint '+str(step))
|
| 205 |
+
write(name+'_validation.json',ev.evaluate(m,table,manifest['validation']))
|
| 206 |
+
review(m,fresh,freshids[:80],name+'_development')
|
| 207 |
+
review(m,fresh,freshids[80:],name+'_heldout')
|
| 208 |
+
write(name+'_heldout_reference_metrics.json',ev.evaluate(m,fresh,freshids[80:]))
|
| 209 |
+
ev.probes(m,name)
|
| 210 |
+
DURABLE.sync(name+' complete diagnostics')
|
| 211 |
+
del m,opt,sched,it,loader;torch.cuda.empty_cache()
|
| 212 |
+
models={'round6':load_semantic('/work/initial',DEVICE)}
|
| 213 |
+
for name,_,_ in arms:models[name]=load_semantic(OUT/name/f'step_{STEPS:05d}',DEVICE)
|
| 214 |
+
grid(models,fresh,freshids[:80],'blinded_development')
|
| 215 |
+
grid(models,fresh,freshids[80:],'blinded_heldout')
|
| 216 |
+
# Familiar historical development probes; never claim these are a held-out archival benchmark.
|
| 217 |
+
import requests
|
| 218 |
+
sources={'migrant_mother':'https://cdn.loc.gov/service/pnp/ppmsca/50200/50236v.jpg','power_house_mechanic':'https://www.archives.gov/exhibits/picturing_the_century/images/port_hine_022_v36.jpg'}
|
| 219 |
+
for name,url in sources.items():
|
| 220 |
+
try:
|
| 221 |
+
response=requests.get(url,timeout=45);response.raise_for_status()
|
| 222 |
+
im=Image.open(io.BytesIO(response.content)).convert('RGB')
|
| 223 |
+
tb=pa.table({'image':[{'bytes':response.content,'path':None}]})
|
| 224 |
+
grid(models,tb,[0],name)
|
| 225 |
+
write(name+'_source.json',{'url':url,'sha256':hashlib.sha256(response.content).hexdigest(),'use':'previously seen archival development probe; original colors unknown'})
|
| 226 |
+
except Exception as e:write(name+'_error.json',{'type':type(e).__name__,'message':str(e)})
|
| 227 |
+
write('status.json',{'status':'completed','elapsed_seconds':time.monotonic()-START,'production_approved':False,'next':'Review blinded grids for coherent plausible color and gray/sepia collapse. This pilot does not test a stochastic conditioned decoder.'})
|
| 228 |
+
DURABLE.sync('round7 four-arm experiment completed')
|
| 229 |
+
if __name__=='__main__':
|
| 230 |
+
try:main()
|
| 231 |
+
except BaseException as e:
|
| 232 |
+
write('failure.json',{'type':type(e).__name__,'traceback':traceback.format_exc().replace(os.environ.get('HF_TOKEN','__EMPTY__'),'[REDACTED]')})
|
| 233 |
+
if DURABLE:DURABLE.sync('round7 failure report and latest saved checkpoints')
|
| 234 |
+
raise
|
experiments/decision-20260928-v2/source/semantic_model.py
ADDED
|
@@ -0,0 +1,94 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Compact pretrained semantic colorizer; dense and shared-palette variants."""
|
| 2 |
+
import json
|
| 3 |
+
from pathlib import Path
|
| 4 |
+
import torch
|
| 5 |
+
from torch import nn
|
| 6 |
+
import torch.nn.functional as F
|
| 7 |
+
from torchvision.models import mobilenet_v3_large, MobileNet_V3_Large_Weights
|
| 8 |
+
from safetensors.torch import save_file, load_file
|
| 9 |
+
|
| 10 |
+
class QueryBlock(nn.Module):
|
| 11 |
+
def __init__(self,d=96):
|
| 12 |
+
super().__init__()
|
| 13 |
+
self.self_attn=nn.MultiheadAttention(d,4,batch_first=True,dropout=0)
|
| 14 |
+
self.cross_attn=nn.MultiheadAttention(d,4,batch_first=True,dropout=0)
|
| 15 |
+
self.norms=nn.ModuleList([nn.LayerNorm(d) for _ in range(3)])
|
| 16 |
+
self.ff=nn.Sequential(nn.Linear(d,2*d),nn.GELU(),nn.Linear(2*d,d))
|
| 17 |
+
def forward(self,q,memory):
|
| 18 |
+
x=self.norms[0](q);q=q+self.self_attn(x,x,x,need_weights=False)[0]
|
| 19 |
+
x=self.norms[1](q);q=q+self.cross_attn(x,memory,memory,need_weights=False)[0]
|
| 20 |
+
return q+self.ff(self.norms[2](q))
|
| 21 |
+
|
| 22 |
+
def refine(d):
|
| 23 |
+
return nn.Sequential(nn.Conv2d(d,d,3,padding=1,bias=False),nn.GroupNorm(8,d),nn.SiLU())
|
| 24 |
+
|
| 25 |
+
class SemanticColorizer(nn.Module):
|
| 26 |
+
def __init__(self,head='palette',pretrained=False,width=128,queries=16):
|
| 27 |
+
super().__init__()
|
| 28 |
+
if head not in ['palette','dense']:raise ValueError(head)
|
| 29 |
+
self.config={'architecture':'SemanticColorizer','head':head,'width':width,'queries':queries,'format_version':1}
|
| 30 |
+
self.encoder=mobilenet_v3_large(weights=MobileNet_V3_Large_Weights.IMAGENET1K_V2 if pretrained else None,progress=False).features
|
| 31 |
+
self.lateral=nn.ModuleList([nn.Conv2d(c,width,1) for c in [24,40,112,960]])
|
| 32 |
+
self.refine=nn.ModuleList([refine(width) for _ in range(3)])
|
| 33 |
+
self.register_buffer('rgb_mean',torch.tensor([.485,.456,.406]).view(1,3,1,1))
|
| 34 |
+
self.register_buffer('rgb_std',torch.tensor([.229,.224,.225]).view(1,3,1,1))
|
| 35 |
+
if head=='palette':
|
| 36 |
+
self.queries=nn.Parameter(torch.randn(queries,width)*.2)
|
| 37 |
+
self.query_blocks=nn.ModuleList([QueryBlock(width) for _ in range(2)])
|
| 38 |
+
self.memory_norm=nn.LayerNorm(width)
|
| 39 |
+
self.query_norm=nn.LayerNorm(width)
|
| 40 |
+
self.pixel=nn.Conv2d(width,width,1)
|
| 41 |
+
self.palette=nn.Sequential(nn.Linear(width,width),nn.GELU(),nn.Linear(width,2))
|
| 42 |
+
self.residual=nn.Conv2d(width,2,1)
|
| 43 |
+
nn.init.normal_(self.palette[-1].weight,std=.01);nn.init.zeros_(self.palette[-1].bias)
|
| 44 |
+
nn.init.zeros_(self.residual.weight);nn.init.zeros_(self.residual.bias)
|
| 45 |
+
else:
|
| 46 |
+
self.dense=nn.Sequential(refine(width),nn.Conv2d(width,2,1))
|
| 47 |
+
nn.init.normal_(self.dense[-1].weight,std=.01);nn.init.zeros_(self.dense[-1].bias)
|
| 48 |
+
count=sum(p.numel() for p in self.parameters())
|
| 49 |
+
if count>=4_000_000:raise ValueError(f'Parameter budget exceeded: {count}')
|
| 50 |
+
|
| 51 |
+
@staticmethod
|
| 52 |
+
def neutral_rgb(L):
|
| 53 |
+
light=(L.float()*50+50).clamp(0,100)
|
| 54 |
+
y=torch.where(light>8,((light+16)/116)**3,light/903.296296)
|
| 55 |
+
g=torch.where(y<=.0031308,12.92*y,1.055*y.clamp_min(1e-8).pow(1/2.4)-.055)
|
| 56 |
+
return g.expand(-1,3,-1,-1)
|
| 57 |
+
|
| 58 |
+
def forward(self,L):
|
| 59 |
+
h,w=L.shape[-2:]
|
| 60 |
+
x=F.pad(self.neutral_rgb(L),(0,(-w)%32,0,(-h)%32),mode='replicate')
|
| 61 |
+
x=(x-self.rgb_mean)/self.rgb_std
|
| 62 |
+
features=[]
|
| 63 |
+
for i,layer in enumerate(self.encoder):
|
| 64 |
+
x=layer(x)
|
| 65 |
+
if i in [3,6,12,16]:features.append(x)
|
| 66 |
+
projected=[layer(f) for layer,f in zip(self.lateral,features)]
|
| 67 |
+
x=projected[-1]
|
| 68 |
+
for i in range(2,-1,-1):
|
| 69 |
+
x=self.refine[2-i](F.interpolate(x,size=projected[i].shape[-2:],mode='bilinear',align_corners=False)+projected[i])
|
| 70 |
+
if self.config['head']=='palette':
|
| 71 |
+
memory=torch.cat([F.adaptive_avg_pool2d(f,(8,8)).flatten(2).transpose(1,2) for f in projected[1:]],1)
|
| 72 |
+
memory=self.memory_norm(memory)
|
| 73 |
+
q=self.queries[None].expand(L.shape[0],-1,-1)
|
| 74 |
+
for block in self.query_blocks:q=block(q,memory)
|
| 75 |
+
q=self.query_norm(q)
|
| 76 |
+
palette=80*torch.tanh(self.palette(q))
|
| 77 |
+
masks=torch.einsum('bqd,bdhw->bqhw',q,self.pixel(x))/(self.config['width']**.5)
|
| 78 |
+
weights=F.softmax(masks.float(),dim=1)
|
| 79 |
+
ab=torch.einsum('bqhw,bqc->bchw',weights,palette.float())+2*torch.tanh(self.residual(x).float())
|
| 80 |
+
else:ab=80*torch.tanh(self.dense(x).float())
|
| 81 |
+
return F.interpolate(ab,size=(x.shape[-2]*4,x.shape[-1]*4),mode='bilinear',align_corners=False)[...,:h,:w]
|
| 82 |
+
def decode(self,z,temperature=.38):return z.float()
|
| 83 |
+
|
| 84 |
+
def save_semantic(model,path):
|
| 85 |
+
path=Path(path);path.mkdir(parents=True,exist_ok=True)
|
| 86 |
+
save_file({k:v.detach().cpu().contiguous() for k,v in model.state_dict().items()},str(path/'model.safetensors'))
|
| 87 |
+
(path/'config.json').write_text(json.dumps(model.config,indent=2))
|
| 88 |
+
(path/'README.md').write_text('# Experimental semantic colorizer\n\nHead: '+model.config['head']+'. Parameters: '+str(sum(p.numel() for p in model.parameters()))+'.\n\nNot approved for production. Requires semantic_model.py; incompatible with the old U-Net loader. Input is Lab lightness normalized to [-1,1]; output is Lab ab. See the run protocol, provenance, selection and visual comparisons. Predictions are plausible colors, not recovered historical truth.\n')
|
| 89 |
+
|
| 90 |
+
def load_semantic(path,device='cpu'):
|
| 91 |
+
path=Path(path);cfg=json.loads((path/'config.json').read_text())
|
| 92 |
+
model=SemanticColorizer(**{k:cfg[k] for k in ['head','width','queries']})
|
| 93 |
+
model.load_state_dict(load_file(str(path/'model.safetensors')),strict=True)
|
| 94 |
+
return model.to(device).eval()
|
experiments/decision-20260928-v2/source/spatial.py
ADDED
|
@@ -0,0 +1,34 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Controlled spatial decoding alternatives; no learned parameters."""
|
| 2 |
+
import math
|
| 3 |
+
import torch
|
| 4 |
+
import torch.nn.functional as F
|
| 5 |
+
|
| 6 |
+
def box_mean(x,radius):
|
| 7 |
+
# Border-normalized local windows, also work on images smaller than kernel.
|
| 8 |
+
return F.avg_pool2d(x,2*radius+1,stride=1,padding=radius,count_include_pad=False)
|
| 9 |
+
|
| 10 |
+
def guided_chroma(L,ab,radius=4,epsilon=.001):
|
| 11 |
+
"""Scalar luminance-guided local linear filter (He et al., ECCV 2010).
|
| 12 |
+
|
| 13 |
+
L is the network's [-1,1] luminance. Epsilon is in [0,1] luminance units.
|
| 14 |
+
Uses luminance only; reference colors never enter inference.
|
| 15 |
+
"""
|
| 16 |
+
if not isinstance(radius,int) or radius<0 or not math.isfinite(epsilon) or epsilon<=0:
|
| 17 |
+
raise ValueError('Invalid guided-filter radius or epsilon')
|
| 18 |
+
if radius==0:return ab
|
| 19 |
+
I=(L.float()+1)/2;p=ab.float()
|
| 20 |
+
mi=box_mean(I,radius);mp=box_mean(p,radius)
|
| 21 |
+
var=(box_mean(I*I,radius)-mi*mi).clamp_min(0)
|
| 22 |
+
cov=box_mean(I*p,radius)-mi*mp
|
| 23 |
+
a=cov/(var+epsilon);b=mp-a*mi
|
| 24 |
+
return box_mean(a,radius)*I+box_mean(b,radius)
|
| 25 |
+
|
| 26 |
+
def spatial_decode(model,logits,L,temperature=.38,pool=1,radius=0,epsilon=.001):
|
| 27 |
+
if pool<1 or not isinstance(pool,int):raise ValueError('pool must be a positive integer')
|
| 28 |
+
if pool>1:
|
| 29 |
+
# Average evidence before annealing, then upsample chroma, not RGB.
|
| 30 |
+
small=F.avg_pool2d(logits,pool,ceil_mode=True,count_include_pad=False)
|
| 31 |
+
ab=model.decode(small,temperature)
|
| 32 |
+
ab=F.interpolate(ab,size=L.shape[-2:],mode='bilinear',align_corners=False)
|
| 33 |
+
else:ab=model.decode(logits,temperature)
|
| 34 |
+
return guided_chroma(L,ab,radius,epsilon)
|
experiments/decision-20260928-v2/source/train_suite.py
ADDED
|
@@ -0,0 +1,359 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Budget-bounded architecture/curriculum screening. Sources pinned by bootstrap."""
|
| 2 |
+
import os,sys,json,time,math,random,hashlib,shutil,traceback,copy,io
|
| 3 |
+
from pathlib import Path
|
| 4 |
+
import numpy as np
|
| 5 |
+
import torch
|
| 6 |
+
from torch import nn
|
| 7 |
+
import torch.nn.functional as F
|
| 8 |
+
from PIL import Image,ImageDraw
|
| 9 |
+
from skimage.color import rgb2lab
|
| 10 |
+
import cv2,pyarrow as pa,pyarrow.parquet as pq
|
| 11 |
+
from huggingface_hub import HfApi,hf_hub_download
|
| 12 |
+
from safetensors.torch import load_file
|
| 13 |
+
from semantic_model import load_semantic
|
| 14 |
+
from decision_model import DecisionColorizer,initialize_from_v3,save_decision,load_decision
|
| 15 |
+
import persistence
|
| 16 |
+
import r7 as r
|
| 17 |
+
from data import decode_image
|
| 18 |
+
|
| 19 |
+
REPO='User-2468/mini-unet-colorizer'
|
| 20 |
+
BASE='1a9eb8af2754ad2329a24cfe50d388cb559441d0'
|
| 21 |
+
ARTIFACTS='26d731862cc0ea531e7198fa823cb09db25bd243'
|
| 22 |
+
CACHE_REV='7b10485ee9d205ce2f7b660dfc6a90012d78b1cb'
|
| 23 |
+
COCO_REV='26ddc382fe75dfc2a0655b5977e296ea10efebce'
|
| 24 |
+
PREFIX=os.environ.get('EXPERIMENT_PREFIX','experiments/decision-20260928')
|
| 25 |
+
SEED=int(os.environ.get('SEED','92893'))
|
| 26 |
+
STEPS=int(os.environ.get('STEPS','3000'));BATCH=24
|
| 27 |
+
START=time.monotonic();TRAIN_LIMIT=float(os.environ.get('TRAIN_LIMIT_SECONDS','3000'))
|
| 28 |
+
OUT=Path('/work/results');OUT.mkdir(exist_ok=True,parents=True)
|
| 29 |
+
persistence.PREFIX=PREFIX;r.OUT=OUT;r.ev.OUT=OUT;r.DEVICE='cuda';r.ev.DEVICE='cuda';r.ev.deadline=lambda:None
|
| 30 |
+
api=HfApi();DURABLE=None
|
| 31 |
+
def write(n,obj):
|
| 32 |
+
p=OUT/n;p.parent.mkdir(parents=True,exist_ok=True);p.write_text(json.dumps(obj,indent=2));return obj
|
| 33 |
+
def dl(f,rev=ARTIFACTS):return Path(hf_hub_download(REPO,f,revision=rev))
|
| 34 |
+
def seed(s):torch.manual_seed(s);np.random.seed(s);random.seed(s)
|
| 35 |
+
|
| 36 |
+
def objective(details,target,step):
|
| 37 |
+
allout=details['all'];b,k=allout.shape[:2]
|
| 38 |
+
losses=r.losses(allout.flatten(0,1),target[:,None].expand(-1,k,-1,-1,-1).flatten(0,1)).reshape(b,k)
|
| 39 |
+
if k==1:return losses.mean(),losses.argmin(1),losses
|
| 40 |
+
# One winning complete image, never a per-pixel best-of-K patchwork.
|
| 41 |
+
winner=losses.detach().argmin(1)
|
| 42 |
+
selected=losses.gather(1,winner[:,None]).mean()
|
| 43 |
+
# Small nonwinner term prevents completely untrained branches; no saturation reward.
|
| 44 |
+
loss=.95*selected+.05*losses.mean()+.1*F.cross_entropy(details['scores'],winner)
|
| 45 |
+
return loss,winner,losses
|
| 46 |
+
|
| 47 |
+
class Batches:
|
| 48 |
+
def __init__(self,rgb,teacher):
|
| 49 |
+
self.rgb=torch.from_numpy(rgb).cuda().permute(0,3,1,2)
|
| 50 |
+
self.teacher=torch.from_numpy(teacher).cuda().float()
|
| 51 |
+
self.n=len(rgb)
|
| 52 |
+
@staticmethod
|
| 53 |
+
def light(g):
|
| 54 |
+
y=torch.where(g<=.04045,g/12.92,((g+.055)/1.055).pow(2.4))
|
| 55 |
+
l=torch.where(y>.008856,116*y.clamp_min(1e-9).pow(1/3)-16,903.296296*y)
|
| 56 |
+
return l/50-1
|
| 57 |
+
def batch(self,step,recipe,original,curriculum_step=None):
|
| 58 |
+
phase=step if curriculum_step is None else curriculum_step
|
| 59 |
+
gen=torch.Generator(device='cuda').manual_seed(SEED*100000+step)
|
| 60 |
+
ids=torch.randint(self.n,(BATCH,),generator=gen,device='cuda')
|
| 61 |
+
rgb=self.rgb[ids].float()/255;t=self.teacher[ids];o=original[ids]
|
| 62 |
+
g=(rgb*rgb.new_tensor([.299,.587,.114])[None,:,None,None]).sum(1,keepdim=True)
|
| 63 |
+
# Emulate standard 8-bit PIL grayscale conversion before Lab conversion.
|
| 64 |
+
g=(g*255).round()/255
|
| 65 |
+
L=self.light(g)
|
| 66 |
+
if recipe=='staged' and phase>STEPS*2/3:
|
| 67 |
+
take=torch.rand(BATCH,1,1,1,generator=gen,device='cuda')<.35
|
| 68 |
+
weights=torch.rand(BATCH,3,1,1,generator=gen,device='cuda')+rgb.new_tensor([.8,1.2,.2])[None,:,None,None]
|
| 69 |
+
weights/=weights.sum(1,keepdim=True)
|
| 70 |
+
film=(rgb*weights).sum(1,keepdim=True)
|
| 71 |
+
gamma=.85+.4*torch.rand(BATCH,1,1,1,generator=gen,device='cuda')
|
| 72 |
+
film=(film.pow(gamma)-.5)*(.8+.3*torch.rand(BATCH,1,1,1,generator=gen,device='cuda'))+.5
|
| 73 |
+
film=film+.005*torch.randn(film.shape,generator=gen,device='cuda')
|
| 74 |
+
L=torch.where(take,self.light(film.clamp(0,1)),L)
|
| 75 |
+
else:
|
| 76 |
+
gain=.9+.2*torch.rand(BATCH,1,1,1,generator=gen,device='cuda')
|
| 77 |
+
L=(L*gain).clamp(-1,1)
|
| 78 |
+
flips=torch.rand(BATCH,1,1,1,generator=gen,device='cuda')<.5
|
| 79 |
+
L=torch.where(flips,L.flip(-1),L);t=torch.where(flips,t.flip(-1),t);o=torch.where(flips,o.flip(-1),o)
|
| 80 |
+
prob=0 if recipe=='teacher' else .5
|
| 81 |
+
if recipe=='staged':prob*=min(max((phase/STEPS-.15)/.35,0),1)
|
| 82 |
+
pick=torch.rand(BATCH,1,1,1,generator=gen,device='cuda')<prob
|
| 83 |
+
# Stochastic full targets: never choose the easier one, never average hues.
|
| 84 |
+
target=torch.where(pick,o,t)
|
| 85 |
+
return L,target,ids,float(pick.float().mean())
|
| 86 |
+
|
| 87 |
+
def summary(result):return {s:result[s]['summary'] for s in ['gray','film']}
|
| 88 |
+
def score(result,baseline):
|
| 89 |
+
val=0;okay=True
|
| 90 |
+
for style in ['gray','film']:
|
| 91 |
+
s=result[style]['summary'];b=baseline[style]['summary']
|
| 92 |
+
okay &= s['missed_color']<=b['missed_color']+.015 and s['neutral_spill']<=b['neutral_spill']+.015 and s['color_coverage']>=.95*b['color_coverage'] and s['ab_error']<=b['ab_error']*1.10
|
| 93 |
+
val+=s['patch_excess']/max(b['patch_excess'],1e-6)+.2*s['ab_error']/b['ab_error']+2*s['missed_color']+s['neutral_spill']
|
| 94 |
+
return val/2,bool(okay)
|
| 95 |
+
|
| 96 |
+
@torch.inference_mode()
|
| 97 |
+
def diagnostics(model,table,ids):
|
| 98 |
+
model.eval();rows=[]
|
| 99 |
+
for idx in ids:
|
| 100 |
+
L,t,_=r.ev.gray_arrays(decode_image(table,idx));x=torch.from_numpy(L)[None].cuda()
|
| 101 |
+
v=model(x,True);a=v['all'];sel=int(v['scores'].argmax(1))
|
| 102 |
+
resized=model(F.interpolate(x,(224,224),mode='bilinear',align_corners=False),True)
|
| 103 |
+
p=F.interpolate(resized['output'],(256,256),mode='bilinear',align_corners=False)
|
| 104 |
+
flip=model(x.flip(-1),True)
|
| 105 |
+
colors=a[0].mean((-1,-2));k=model.modes
|
| 106 |
+
diff=(a[:,:,None]-a[:,None]).abs().mean((0,3,4,5))
|
| 107 |
+
target=torch.from_numpy(t)[None].cuda();target64=F.avg_pool2d(target,4)
|
| 108 |
+
losses=r.losses(a.flatten(0,1),target64.expand(k,-1,-1,-1))
|
| 109 |
+
mask=v['weights'].clamp_min(1e-8)
|
| 110 |
+
rows.append({'index':int(idx),'selected_mode':sel,'resized_mode':int(resized['scores'].argmax(1)),
|
| 111 |
+
'flip_mode':int(flip['scores'].argmax(1)),'resize_ab_change':float((p-v['output']).abs().mean()),
|
| 112 |
+
'flip_ab_change':float((flip['output'].flip(-1)-v['output']).abs().mean()),
|
| 113 |
+
'mode_diversity':float(diff.sum()/max(k*(k-1),1)),
|
| 114 |
+
'default_target_loss':float(losses[sel]),'oracle_target_loss':float(losses.min()),
|
| 115 |
+
'mask_entropy':float(-(mask*mask.log()).sum(1).mean()),
|
| 116 |
+
'palette_chroma':float(v['palette'].norm(dim=-1).mean()),'output_chroma':float(v['output'].norm(dim=1).mean())})
|
| 117 |
+
keys=['resize_ab_change','flip_ab_change','mode_diversity','default_target_loss','oracle_target_loss','mask_entropy','palette_chroma','output_chroma']
|
| 118 |
+
return {'rows':rows,'summary':{k:float(np.mean([x[k] for x in rows])) for k in keys},
|
| 119 |
+
'mode_frequency':np.bincount([x['selected_mode'] for x in rows],minlength=model.modes).tolist(),
|
| 120 |
+
'resize_switch_fraction':float(np.mean([x['selected_mode']!=x['resized_mode'] for x in rows])),
|
| 121 |
+
'note':'Oracle uses unknown original only for diagnosis. Default mode is selected by the learned image-only score, never by ground truth.'}
|
| 122 |
+
|
| 123 |
+
def paired_bootstrap(result,base):
|
| 124 |
+
rng=np.random.default_rng(937);out={}
|
| 125 |
+
for style in ['gray','film']:
|
| 126 |
+
br={z['index']:z for z in base[style]['per_image']};rows=[z for z in result[style]['per_image'] if z['target_chroma']>=5]
|
| 127 |
+
out[style]={}
|
| 128 |
+
for key in ['patch_excess','raw_patch_excess','missed_color','neutral_spill','color_coverage','ab_error']:
|
| 129 |
+
d=np.array([z[key]-br[z['index']][key] for z in rows]);boot=d[rng.integers(len(d),size=(3000,len(d)))].mean(1)
|
| 130 |
+
out[style][key]={'delta':float(d.mean()),'ci95':np.quantile(boot,[.025,.975]).tolist()}
|
| 131 |
+
return out
|
| 132 |
+
|
| 133 |
+
@torch.inference_mode()
|
| 134 |
+
def pipeline_grid(models,table,ids,label):
|
| 135 |
+
from inference import colorize
|
| 136 |
+
from PIL import ImageOps
|
| 137 |
+
names=list(models);order=np.random.default_rng(9184).permutation(len(names))
|
| 138 |
+
key={chr(65+j):names[int(i)] for j,i in enumerate(order)};write(label+'_key.json',key)
|
| 139 |
+
for start in range(0,min(32,len(ids)),4):
|
| 140 |
+
canvas=Image.new('RGB',(320*(len(names)+1),344*len(ids[start:start+4])),'white');d=ImageDraw.Draw(canvas)
|
| 141 |
+
for row,idx in enumerate(ids[start:start+4]):
|
| 142 |
+
original=decode_image(table,idx).convert('L').convert('RGB')
|
| 143 |
+
original.thumbnail((800,800))
|
| 144 |
+
d.text((2,row*344),f'{idx}: grayscale | '+' | '.join(key),fill='black')
|
| 145 |
+
canvas.paste(ImageOps.pad(original,(320,320),color='white'),(0,row*344+24))
|
| 146 |
+
for col,j in enumerate(order):
|
| 147 |
+
output=colorize(models[names[int(j)]].eval(),original,size=256,radius=8)
|
| 148 |
+
canvas.paste(ImageOps.pad(output,(320,320),color='white'),((col+1)*320,row*344+24))
|
| 149 |
+
canvas.save(OUT/f'{label}_{start//4}.jpg',quality=93)
|
| 150 |
+
|
| 151 |
+
@torch.inference_mode()
|
| 152 |
+
def resolution_grid(model,table,records):
|
| 153 |
+
from inference import colorize
|
| 154 |
+
from PIL import ImageOps
|
| 155 |
+
for start in range(0,len(table),3):
|
| 156 |
+
count=min(3,len(table)-start)
|
| 157 |
+
canvas=Image.new('RGB',(1280,344*count),'white');d=ImageDraw.Draw(canvas)
|
| 158 |
+
for row,idx in enumerate(range(start,start+count)):
|
| 159 |
+
im=decode_image(table,idx).convert('L').convert('RGB');im.thumbnail((1000,1000))
|
| 160 |
+
d.text((2,row*344),records[idx]['name']+' | input | 256 | 384 | 512',fill='black')
|
| 161 |
+
canvas.paste(ImageOps.pad(im,(320,320),color='white'),(0,row*344+24))
|
| 162 |
+
for col,size in enumerate([256,384,512],1):
|
| 163 |
+
pred=colorize(model,im,size=size,radius=8)
|
| 164 |
+
canvas.paste(ImageOps.pad(pred,(320,320),color='white'),(col*320,row*344+24))
|
| 165 |
+
canvas.save(OUT/f'v3_resolution_{start//3}.jpg',quality=94)
|
| 166 |
+
|
| 167 |
+
def main():
|
| 168 |
+
global DURABLE
|
| 169 |
+
assert api.whoami()['name']=='User-2468'
|
| 170 |
+
DURABLE=persistence.DurableRun(OUT)
|
| 171 |
+
shutil.copytree('/work/source',OUT/'source',dirs_exist_ok=True,ignore=shutil.ignore_patterns('__pycache__'))
|
| 172 |
+
shutil.copy2('/work/PROTOCOL.md',OUT/'PROTOCOL.md')
|
| 173 |
+
if Path('/work/REVIEW_HANDOFF.md').exists():shutil.copy2('/work/REVIEW_HANDOFF.md',OUT/'REVIEW_HANDOFF.md')
|
| 174 |
+
if Path('/work/supplement').exists():shutil.copytree('/work/supplement',OUT/'supplement',dirs_exist_ok=True)
|
| 175 |
+
write('status.json',{'status':'preflight','seed':SEED,'production_approved':False})
|
| 176 |
+
DURABLE.sync('source and write verification before training')
|
| 177 |
+
torch.set_num_threads(6);seed(SEED)
|
| 178 |
+
torch.backends.cuda.matmul.allow_tf32=False;torch.backends.cudnn.allow_tf32=False
|
| 179 |
+
initial=Path('/work/initial');initial.mkdir(exist_ok=True)
|
| 180 |
+
for f in ['model.safetensors','config.json']:shutil.copy2(dl(f,BASE),initial/f)
|
| 181 |
+
baseline=load_semantic(initial,'cuda').requires_grad_(False)
|
| 182 |
+
counts={};smokes={}
|
| 183 |
+
encoder_rev=api.model_info('timm/mobilevitv2_075.cvnets_in1k').sha
|
| 184 |
+
write('encoder_provenance.json',{'repo':'timm/mobilevitv2_075.cvnets_in1k','revision':encoder_rev,'license':'Apple sample code license; include upstream LICENSE if deployed','timm':'1.0.19'})
|
| 185 |
+
constructors=[('single',1,'mobilenet'),('multi',4,'mobilenet'),('mobilevit',1,'mobilevit')]
|
| 186 |
+
for name,k,backbone in constructors:
|
| 187 |
+
seed(SEED);m=DecisionColorizer(modes=k,backbone=backbone,pretrained=backbone=='mobilevit',encoder_revision=encoder_rev if backbone=='mobilevit' else None).cuda()
|
| 188 |
+
copied=initialize_from_v3(m,baseline);m.eval();count=sum(p.numel() for p in m.parameters());assert count<4_000_000
|
| 189 |
+
counts[name]=count
|
| 190 |
+
x=torch.rand(2,1,97,113,device='cuda')*2-1
|
| 191 |
+
with torch.no_grad():
|
| 192 |
+
v=m(x,True);assert v['all'].shape==(2,k,2,97,113) and torch.isfinite(v['all']).all()
|
| 193 |
+
parity=float((v['output']-baseline(x)).abs().max()) if name=='single' else None
|
| 194 |
+
if parity is not None:assert parity<1e-3,parity
|
| 195 |
+
loss,_,_=objective(m(x,True),torch.randn(2,2,64,64,device='cuda'),1)
|
| 196 |
+
loss.backward();assert torch.isfinite(loss) and all(torch.isfinite(p.grad).all() for p in m.parameters() if p.grad is not None)
|
| 197 |
+
m.zero_grad(set_to_none=True);save_decision(m,OUT/'initial'/name)
|
| 198 |
+
re=load_decision(OUT/'initial'/name,'cuda')
|
| 199 |
+
with torch.no_grad():assert torch.allclose(m(x),re(x),atol=1e-4)
|
| 200 |
+
smokes[name]={'parameters':count,'copied_tensors':len(copied),'odd_shape':True,'finite_backward':True,'reload':True,'baseline_max_abs':parity}
|
| 201 |
+
del m,re
|
| 202 |
+
write('smoke.json',smokes);DURABLE.sync('all three architectures pass preflight')
|
| 203 |
+
torch.backends.cuda.matmul.allow_tf32=True;torch.backends.cudnn.allow_tf32=True;torch.backends.cudnn.benchmark=True
|
| 204 |
+
tables=[]
|
| 205 |
+
for repo,rev,files in [('johnowhitaker/imagenette2-320','771c1310a2487e8076ede6b7d6307244aa8400af',['default/train/0000.parquet']),('detection-datasets/coco',COCO_REV,['default/train/0000.parquet','default/train/0001.parquet'])]:
|
| 206 |
+
for f in files:
|
| 207 |
+
print('PREP downloading',repo,f,flush=True)
|
| 208 |
+
tables.append(pq.read_table(hf_hub_download(repo,f,repo_type='dataset',revision=rev),columns=['image']))
|
| 209 |
+
print('PREP loaded',repo,f,len(tables[-1]),flush=True)
|
| 210 |
+
table=pa.concat_tables(tables)
|
| 211 |
+
full_manifest=json.loads(dl('experiments/region-coherence-20260928/source/previous_manifest.json').read_text())
|
| 212 |
+
full_ids=full_manifest['train']
|
| 213 |
+
# Same fixed training subset for every arm, independent of new dev/test.
|
| 214 |
+
ids=np.random.default_rng(92893).choice(full_ids,8192,replace=False).tolist();pos={i:j for j,i in enumerate(ids)}
|
| 215 |
+
teacher=np.empty((len(ids),2,64,64),np.float16)
|
| 216 |
+
sums=json.loads(dl('experiments/round6-20260928/SHA256SUMS.json',CACHE_REV).read_text())
|
| 217 |
+
for start in range(0,len(full_ids),2048):
|
| 218 |
+
f=f'teacher_cache/{start:05d}.npz';p=dl('experiments/round6-20260928/'+f,CACHE_REV)
|
| 219 |
+
assert hashlib.sha256(p.read_bytes()).hexdigest()==sums[f]
|
| 220 |
+
with np.load(p) as z:
|
| 221 |
+
# NpzFile indexing decompresses a member on every access. Materialize once.
|
| 222 |
+
cache_ids=z['indices'];cache_targets=z['target']
|
| 223 |
+
assert cache_ids.tolist()==full_ids[start:start+len(cache_ids)]
|
| 224 |
+
src=[j for j,i in enumerate(cache_ids) if int(i) in pos]
|
| 225 |
+
dst=[pos[int(cache_ids[j])] for j in src]
|
| 226 |
+
teacher[dst]=cache_targets[src]
|
| 227 |
+
print('PREP cache',start,'selected',len(src),flush=True)
|
| 228 |
+
print('PREP decoding training images',len(ids),flush=True)
|
| 229 |
+
rgb=np.stack([np.asarray(decode_image(table,i).resize((256,256),Image.Resampling.BILINEAR)) for i in ids])
|
| 230 |
+
original=np.stack([cv2.resize(rgb2lab(im.astype(np.float32)/255)[...,1:],(64,64),interpolation=cv2.INTER_AREA).transpose(2,0,1) for im in rgb]).astype(np.float32)
|
| 231 |
+
batcher=Batches(rgb,teacher);original=torch.from_numpy(original).cuda();del rgb,teacher
|
| 232 |
+
print('PREP training tensors ready',flush=True)
|
| 233 |
+
fresh=pq.read_table(hf_hub_download('detection-datasets/coco','default/val/0001.parquet',repo_type='dataset',revision=COCO_REV),columns=['image'])
|
| 234 |
+
exclude=set()
|
| 235 |
+
for path,keys in [('round7-20260928',['development','heldout']),('region-coherence-20260928',['development_regression','heldout_regression','new_holdout']),('recurrent-refinement-20260928',['fresh_holdout','development_regression'])]:
|
| 236 |
+
manifest=json.loads(dl('experiments/'+path+'/manifest.json').read_text())
|
| 237 |
+
for key in keys:exclude.update(manifest[key])
|
| 238 |
+
hashes={hashlib.sha256(table['image'][i].as_py()['bytes']).digest() for i in full_ids}
|
| 239 |
+
freshids=[int(i) for i in np.random.default_rng(92893).permutation(len(fresh)) if int(i) not in exclude and hashlib.sha256(fresh['image'][int(i)].as_py()['bytes']).digest() not in hashes][:288]
|
| 240 |
+
assert len(freshids)==288
|
| 241 |
+
dev,test=freshids[:128],freshids[128:]
|
| 242 |
+
write('manifest.json',{'train':ids,'development':dev,'test':test,'prior_eval_excluded':sorted(exclude),'coco_revision':COCO_REV,'fresh_file':'default/val/0001.parquet','teacher_revision':CACHE_REV,'base_revision':BASE,'seed':SEED,'near_duplicates_and_upstream_overlap_not_excluded':True})
|
| 243 |
+
base_dev=r.ev.evaluate(baseline,fresh,dev);write('baseline_development.json',base_dev)
|
| 244 |
+
DURABLE.sync('dataset checks and unseen development/test split frozen')
|
| 245 |
+
arms=[(h,c) for h in ['single','multi'] for c in ['teacher','mix','staged']]+[('mobilevit','staged')]
|
| 246 |
+
if os.environ.get('ARMS'):arms=[tuple(x.split(':')) for x in os.environ['ARMS'].split(',')]
|
| 247 |
+
selected={};histories={};dev_results={};begin_train=time.monotonic()
|
| 248 |
+
for head,recipe in arms:
|
| 249 |
+
name=head+'_'+recipe;seed(SEED)
|
| 250 |
+
m=load_decision(OUT/'initial'/head,'cuda')
|
| 251 |
+
# New encoder gets a separate alignment warmup, explicitly outside matched six arms.
|
| 252 |
+
warmup=1000 if head=='mobilevit' else 0
|
| 253 |
+
groups=[{'params':[p for n,p in m.named_parameters() if n.startswith('encoder.')],'lr':2e-5 if head=='mobilevit' else 1e-5},
|
| 254 |
+
{'params':[p for n,p in m.named_parameters() if not n.startswith('encoder.')],'lr':2e-4 if head=='mobilevit' else 1e-4}]
|
| 255 |
+
opt=torch.optim.AdamW(groups,weight_decay=.01)
|
| 256 |
+
for g in opt.param_groups:g['base_lr']=g['lr']
|
| 257 |
+
history=[];best=1e9;beststep=0;last=0;arm_start=time.monotonic();mode_counts=np.zeros(m.modes,np.int64)
|
| 258 |
+
for s in range(1,STEPS+warmup+1):
|
| 259 |
+
if time.monotonic()-START>TRAIN_LIMIT:break
|
| 260 |
+
step=max(1,s-warmup);active_recipe='teacher' if s<=warmup else recipe
|
| 261 |
+
L,target,_,fraction=batcher.batch(s,active_recipe,original,curriculum_step=step)
|
| 262 |
+
m.train()
|
| 263 |
+
for layer in m.encoder.modules():
|
| 264 |
+
if isinstance(layer,nn.BatchNorm2d):layer.eval()
|
| 265 |
+
opt.zero_grad(set_to_none=True)
|
| 266 |
+
factor=min(s/150,1)*(.15+.85*(1+math.cos(math.pi*max(0,s-warmup)/STEPS))/2)
|
| 267 |
+
for g in opt.param_groups:g['lr']=g['base_lr']*factor
|
| 268 |
+
with torch.autocast('cuda',dtype=torch.bfloat16):
|
| 269 |
+
details=m(L,True)
|
| 270 |
+
loss,winners,terms=objective(details,target,step)
|
| 271 |
+
consistency=loss.new_zeros(())
|
| 272 |
+
if active_recipe=='staged' and step>STEPS*2/3 and s%4==0:
|
| 273 |
+
small=m(F.interpolate(L,(192,192),mode='bilinear',align_corners=False),True)
|
| 274 |
+
a=F.adaptive_avg_pool2d(small['all'].flatten(0,1).float(),(64,64))
|
| 275 |
+
b=F.adaptive_avg_pool2d(details['all'].flatten(0,1).detach().float(),(64,64))
|
| 276 |
+
consistency=F.l1_loss(a/20,b/20)
|
| 277 |
+
if m.modes>1:consistency+=.1*F.kl_div(small['scores'].log_softmax(1),details['scores'].detach().softmax(1),reduction='batchmean')
|
| 278 |
+
loss=loss+.4*consistency # every fourth step -> mean weight .1
|
| 279 |
+
assert torch.isfinite(loss),name
|
| 280 |
+
loss.backward();torch.nn.utils.clip_grad_norm_(m.parameters(),1,error_if_nonfinite=True);opt.step();last=s
|
| 281 |
+
mode_counts+=np.bincount(winners.detach().cpu().numpy(),minlength=m.modes)
|
| 282 |
+
if s%100==0:
|
| 283 |
+
row={'step':s,'loss':float(loss.detach()),'reconstruction':float(terms.mean().detach()),'consistency':float(consistency.detach()),'photo_fraction_this_batch':fraction,'winner_counts':mode_counts.tolist(),'seconds':time.monotonic()-arm_start}
|
| 284 |
+
history.append(row);print('STEP',name,s,round(row['loss'],4),flush=True)
|
| 285 |
+
if s%1000==0 or s==STEPS+warmup:
|
| 286 |
+
m.eval();res=r.ev.evaluate(m,fresh,dev);val,okay=score(res,base_dev)
|
| 287 |
+
write(name+f'/development_{s:05d}.json',res)
|
| 288 |
+
write(name+'/history.json',history)
|
| 289 |
+
save_decision(m,OUT/name/f'step_{s:05d}')
|
| 290 |
+
# Ineligible models still retained for diagnosis, ranked after eligible.
|
| 291 |
+
ranked=val+(0 if okay else 100)
|
| 292 |
+
if ranked<best:
|
| 293 |
+
best=ranked;beststep=s;save_decision(m,OUT/name/'candidate');dev_results[name]=res
|
| 294 |
+
write('status.json',{'status':'training','arm':name,'step':s,'elapsed_seconds':time.monotonic()-START,'production_approved':False})
|
| 295 |
+
DURABLE.sync(name+' checkpoint '+str(s))
|
| 296 |
+
if not beststep:raise RuntimeError('Budget exhausted before first candidate for '+name)
|
| 297 |
+
write(name+'/history.json',history)
|
| 298 |
+
torch.save({'optimizer':opt.state_dict(),'step':last,'seed':SEED,'torch_rng':torch.get_rng_state(),'cuda_rng':torch.cuda.get_rng_state_all(),'exact_resume':False},OUT/name/'optimizer.pt')
|
| 299 |
+
selected[name]={'step':beststep,'rank_score':best,'parameters':counts[head],'steps_completed':last,'warmup_steps':warmup,'seconds':time.monotonic()-arm_start}
|
| 300 |
+
write('candidates.json',selected);DURABLE.sync(name+' selected checkpoint on development only')
|
| 301 |
+
del m,opt;torch.cuda.empty_cache()
|
| 302 |
+
if last<STEPS+warmup:break
|
| 303 |
+
# Test is opened only after all checkpoint choices have been frozen.
|
| 304 |
+
write('candidates.json',selected);DURABLE.sync('freeze candidates before opening final test')
|
| 305 |
+
models={'v3':baseline};test_results={}
|
| 306 |
+
base_test=r.ev.evaluate(baseline,fresh,test);write('v3_test.json',base_test);test_results['v3']=base_test
|
| 307 |
+
for name in selected:
|
| 308 |
+
m=load_decision(OUT/name/'candidate','cuda');models[name]=m
|
| 309 |
+
res=r.ev.evaluate(m,fresh,test);test_results[name]=res
|
| 310 |
+
write(name+'_test.json',res);write(name+'_test_paired.json',paired_bootstrap(res,base_test))
|
| 311 |
+
write(name+'_diagnostics.json',diagnostics(m,fresh,test[:48]))
|
| 312 |
+
r.ev.probes(m,name+'_probes')
|
| 313 |
+
x=torch.zeros(1,1,256,256,device='cuda')
|
| 314 |
+
with torch.inference_mode():
|
| 315 |
+
for _ in range(5):m(x)
|
| 316 |
+
torch.cuda.synchronize();t0=time.monotonic()
|
| 317 |
+
for _ in range(30):m(x)
|
| 318 |
+
torch.cuda.synchronize()
|
| 319 |
+
write(name+'_latency.json',{'milliseconds':1000*(time.monotonic()-t0)/30,'hardware':torch.cuda.get_device_name(),'size':256,'precision':'FP32'})
|
| 320 |
+
DURABLE.sync(name+' final test and diagnostic outputs')
|
| 321 |
+
r.ev.probes(baseline,'v3_probes')
|
| 322 |
+
# Smaller comparison groups preserve enough image detail to inspect.
|
| 323 |
+
for head in ['single','multi']:
|
| 324 |
+
group={n:m for n,m in models.items() if n=='v3' or n.startswith(head+'_')}
|
| 325 |
+
r.grid(group,fresh,test,head+'_test')
|
| 326 |
+
if 'mobilevit_staged' in models:r.grid({'v3':baseline,'single_staged':models['single_staged'],'mobilevit_staged':models['mobilevit_staged']},fresh,test,'encoder_test')
|
| 327 |
+
# Historical originals deliberately never serve as neutral-color supervision.
|
| 328 |
+
archive=json.loads(Path('/work/archive_sources.json').read_text())
|
| 329 |
+
import requests
|
| 330 |
+
records=[];images=[]
|
| 331 |
+
for entry in archive:
|
| 332 |
+
try:
|
| 333 |
+
response=requests.get(entry['image_url'],timeout=20);response.raise_for_status()
|
| 334 |
+
im=Image.open(io.BytesIO(response.content)).convert('RGB');im.load()
|
| 335 |
+
images.append({'bytes':response.content,'path':None});records.append(entry|{'sha256':hashlib.sha256(response.content).hexdigest()})
|
| 336 |
+
except Exception as e:records.append(entry|{'failed':type(e).__name__})
|
| 337 |
+
if images:
|
| 338 |
+
ar=pa.table({'image':images})
|
| 339 |
+
resolution_grid(baseline,ar,[x for x in records if 'failed' not in x])
|
| 340 |
+
for head in ['single','multi']:
|
| 341 |
+
r.grid({n:m for n,m in models.items() if n=='v3' or n.startswith(head+'_')},ar,list(range(len(images))),head+'_archive')
|
| 342 |
+
if 'mobilevit_staged' in models:r.grid({'v3':baseline,'mobilevit_staged':models['mobilevit_staged']},ar,list(range(len(images))),'encoder_archive')
|
| 343 |
+
pipeline_grid({'v3':baseline,**{n:m for n,m in models.items() if n.endswith('staged')}},ar,list(range(len(images))),'production_archive')
|
| 344 |
+
pipeline_grid({'v3':baseline,**{n:m for n,m in models.items() if n.endswith('staged')}},fresh,test[:16],'production_test')
|
| 345 |
+
write('archive_sources.json',records)
|
| 346 |
+
write('summary.json',{'parameters':counts,'candidates':selected,'test':{n:summary(v) for n,v in test_results.items()},'all_arms_completed':len(selected)==len(arms),'elapsed_seconds':time.monotonic()-START,'estimated_compute_usd':(time.monotonic()-START)/3600*.8,'warning':'No metric proves realism. All outputs need visual review. MobileViT received extra alignment warmup and different encoder LR; not a pure head ablation.'})
|
| 347 |
+
write('status.json',{'status':'completed','all_arms_completed':len(selected)==len(arms),'elapsed_seconds':time.monotonic()-START,'production_approved':False})
|
| 348 |
+
DURABLE.sync('architecture and curriculum suite completed')
|
| 349 |
+
|
| 350 |
+
if __name__=='__main__':
|
| 351 |
+
try:main()
|
| 352 |
+
except BaseException:
|
| 353 |
+
msg=traceback.format_exc();token=os.environ.get('HF_TOKEN')
|
| 354 |
+
if token:msg=msg.replace(token,'[REDACTED]')
|
| 355 |
+
write('failure.json',{'traceback':msg});write('status.json',{'status':'failed','elapsed_seconds':time.monotonic()-START,'production_approved':False})
|
| 356 |
+
if DURABLE:
|
| 357 |
+
try:DURABLE.sync('failure report and saved checkpoints')
|
| 358 |
+
except Exception:pass
|
| 359 |
+
print(msg,flush=True);raise SystemExit(1)
|
experiments/decision-20260928-v2/status.json
ADDED
|
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"status": "preflight",
|
| 3 |
+
"seed": 92893,
|
| 4 |
+
"production_approved": false
|
| 5 |
+
}
|
experiments/decision-20260928-v2/supplement/baseline_nara_resolution.json
ADDED
|
@@ -0,0 +1,128 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
[
|
| 2 |
+
{
|
| 3 |
+
"name": "nara_42.jpg",
|
| 4 |
+
"size": 256,
|
| 5 |
+
"seconds_cpu": 0.14755021000019042,
|
| 6 |
+
"mask_entropy": 2.0555808544158936,
|
| 7 |
+
"output_chroma": 12.437393188476562
|
| 8 |
+
},
|
| 9 |
+
{
|
| 10 |
+
"name": "nara_42.jpg",
|
| 11 |
+
"size": 384,
|
| 12 |
+
"seconds_cpu": 0.23218034499996065,
|
| 13 |
+
"mask_entropy": 2.090886116027832,
|
| 14 |
+
"output_chroma": 14.703699111938477
|
| 15 |
+
},
|
| 16 |
+
{
|
| 17 |
+
"name": "nara_42.jpg",
|
| 18 |
+
"size": 512,
|
| 19 |
+
"seconds_cpu": 0.3470577670000239,
|
| 20 |
+
"mask_entropy": 2.164964199066162,
|
| 21 |
+
"output_chroma": 13.910662651062012
|
| 22 |
+
},
|
| 23 |
+
{
|
| 24 |
+
"name": "nara_43.jpg",
|
| 25 |
+
"size": 256,
|
| 26 |
+
"seconds_cpu": 0.11568649200034997,
|
| 27 |
+
"mask_entropy": 2.221789598464966,
|
| 28 |
+
"output_chroma": 22.151323318481445
|
| 29 |
+
},
|
| 30 |
+
{
|
| 31 |
+
"name": "nara_43.jpg",
|
| 32 |
+
"size": 384,
|
| 33 |
+
"seconds_cpu": 0.17595398799994655,
|
| 34 |
+
"mask_entropy": 2.151303291320801,
|
| 35 |
+
"output_chroma": 23.234233856201172
|
| 36 |
+
},
|
| 37 |
+
{
|
| 38 |
+
"name": "nara_43.jpg",
|
| 39 |
+
"size": 512,
|
| 40 |
+
"seconds_cpu": 0.2790908069996476,
|
| 41 |
+
"mask_entropy": 2.1258065700531006,
|
| 42 |
+
"output_chroma": 20.545372009277344
|
| 43 |
+
},
|
| 44 |
+
{
|
| 45 |
+
"name": "nara_44.jpg",
|
| 46 |
+
"size": 256,
|
| 47 |
+
"seconds_cpu": 0.09606875300005413,
|
| 48 |
+
"mask_entropy": 1.8140740394592285,
|
| 49 |
+
"output_chroma": 7.230224609375
|
| 50 |
+
},
|
| 51 |
+
{
|
| 52 |
+
"name": "nara_44.jpg",
|
| 53 |
+
"size": 384,
|
| 54 |
+
"seconds_cpu": 0.1647310340003969,
|
| 55 |
+
"mask_entropy": 1.9336862564086914,
|
| 56 |
+
"output_chroma": 8.660061836242676
|
| 57 |
+
},
|
| 58 |
+
{
|
| 59 |
+
"name": "nara_44.jpg",
|
| 60 |
+
"size": 512,
|
| 61 |
+
"seconds_cpu": 0.2678544229997897,
|
| 62 |
+
"mask_entropy": 2.003618001937866,
|
| 63 |
+
"output_chroma": 10.750614166259766
|
| 64 |
+
},
|
| 65 |
+
{
|
| 66 |
+
"name": "nara_45.jpg",
|
| 67 |
+
"size": 256,
|
| 68 |
+
"seconds_cpu": 0.12318589300002714,
|
| 69 |
+
"mask_entropy": 1.9286333322525024,
|
| 70 |
+
"output_chroma": 10.617255210876465
|
| 71 |
+
},
|
| 72 |
+
{
|
| 73 |
+
"name": "nara_45.jpg",
|
| 74 |
+
"size": 384,
|
| 75 |
+
"seconds_cpu": 0.2194160090002697,
|
| 76 |
+
"mask_entropy": 1.786165475845337,
|
| 77 |
+
"output_chroma": 13.136454582214355
|
| 78 |
+
},
|
| 79 |
+
{
|
| 80 |
+
"name": "nara_45.jpg",
|
| 81 |
+
"size": 512,
|
| 82 |
+
"seconds_cpu": 0.3068678010004078,
|
| 83 |
+
"mask_entropy": 1.711348056793213,
|
| 84 |
+
"output_chroma": 15.462327003479004
|
| 85 |
+
},
|
| 86 |
+
{
|
| 87 |
+
"name": "nara_46.jpg",
|
| 88 |
+
"size": 256,
|
| 89 |
+
"seconds_cpu": 0.14160436099973595,
|
| 90 |
+
"mask_entropy": 1.984742522239685,
|
| 91 |
+
"output_chroma": 6.706954002380371
|
| 92 |
+
},
|
| 93 |
+
{
|
| 94 |
+
"name": "nara_46.jpg",
|
| 95 |
+
"size": 384,
|
| 96 |
+
"seconds_cpu": 0.20283023299998604,
|
| 97 |
+
"mask_entropy": 2.0790905952453613,
|
| 98 |
+
"output_chroma": 8.468958854675293
|
| 99 |
+
},
|
| 100 |
+
{
|
| 101 |
+
"name": "nara_46.jpg",
|
| 102 |
+
"size": 512,
|
| 103 |
+
"seconds_cpu": 0.3308731010001793,
|
| 104 |
+
"mask_entropy": 2.1869609355926514,
|
| 105 |
+
"output_chroma": 9.939047813415527
|
| 106 |
+
},
|
| 107 |
+
{
|
| 108 |
+
"name": "nara_49.jpg",
|
| 109 |
+
"size": 256,
|
| 110 |
+
"seconds_cpu": 0.14086192399963693,
|
| 111 |
+
"mask_entropy": 2.110356330871582,
|
| 112 |
+
"output_chroma": 15.332934379577637
|
| 113 |
+
},
|
| 114 |
+
{
|
| 115 |
+
"name": "nara_49.jpg",
|
| 116 |
+
"size": 384,
|
| 117 |
+
"seconds_cpu": 0.2311230479999722,
|
| 118 |
+
"mask_entropy": 2.158229112625122,
|
| 119 |
+
"output_chroma": 19.6121768951416
|
| 120 |
+
},
|
| 121 |
+
{
|
| 122 |
+
"name": "nara_49.jpg",
|
| 123 |
+
"size": 512,
|
| 124 |
+
"seconds_cpu": 0.2948662769999828,
|
| 125 |
+
"mask_entropy": 2.1546900272369385,
|
| 126 |
+
"output_chroma": 16.430278778076172
|
| 127 |
+
}
|
| 128 |
+
]
|
experiments/decision-20260928-v2/supplement/baseline_resolution.json
ADDED
|
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
[
|
| 2 |
+
{
|
| 3 |
+
"image": "07.jpg",
|
| 4 |
+
"size": 256,
|
| 5 |
+
"mask_entropy": 1.855709433555603,
|
| 6 |
+
"palette_chroma": 30.90544319152832,
|
| 7 |
+
"output_chroma": 19.500049591064453
|
| 8 |
+
},
|
| 9 |
+
{
|
| 10 |
+
"image": "07.jpg",
|
| 11 |
+
"size": 384,
|
| 12 |
+
"mask_entropy": 1.8696938753128052,
|
| 13 |
+
"palette_chroma": 20.328649520874023,
|
| 14 |
+
"output_chroma": 11.265676498413086
|
| 15 |
+
},
|
| 16 |
+
{
|
| 17 |
+
"image": "07.jpg",
|
| 18 |
+
"size": 512,
|
| 19 |
+
"mask_entropy": 1.9025498628616333,
|
| 20 |
+
"palette_chroma": 20.42593765258789,
|
| 21 |
+
"output_chroma": 12.631791114807129
|
| 22 |
+
}
|
| 23 |
+
]
|
experiments/decision-20260928-v2/supplement/cache_loader_check.json
ADDED
|
@@ -0,0 +1,8 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"old_seconds": 0.6009589720001713,
|
| 3 |
+
"new_seconds": 0.010625808999975561,
|
| 4 |
+
"speedup": 56.556538142324364,
|
| 5 |
+
"exact_same_targets": true,
|
| 6 |
+
"synthetic_rows": 128,
|
| 7 |
+
"selected_rows": 64
|
| 8 |
+
}
|
experiments/decision-20260928-v2/supplement/nara_sources.json
ADDED
|
@@ -0,0 +1,50 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
[
|
| 2 |
+
{
|
| 3 |
+
"name": "nara_42.jpg",
|
| 4 |
+
"title": "Power station",
|
| 5 |
+
"source": "https://www.archives.gov/exhibits/picturing_the_century/greatdep/greatdep_img42.html",
|
| 6 |
+
"image_url": "https://www.archives.gov/exhibits/picturing_the_century/images/greatdep_042_v71.jpg",
|
| 7 |
+
"sha256": "a73186b25853d95a8e1785d1497923c1052747e8b045c5eac4c05dade60b2e87",
|
| 8 |
+
"bytes": 47847
|
| 9 |
+
},
|
| 10 |
+
{
|
| 11 |
+
"name": "nara_43.jpg",
|
| 12 |
+
"title": "Boxing team",
|
| 13 |
+
"source": "https://www.archives.gov/exhibits/picturing_the_century/greatdep/greatdep_img43.html",
|
| 14 |
+
"image_url": "https://www.archives.gov/exhibits/picturing_the_century/images/greatdep_043_v69.jpg",
|
| 15 |
+
"sha256": "f22dc07c4bdd0f3dac4b65ff0548395dc9aa88d8169a714fc40524249233b977",
|
| 16 |
+
"bytes": 40864
|
| 17 |
+
},
|
| 18 |
+
{
|
| 19 |
+
"name": "nara_44.jpg",
|
| 20 |
+
"title": "Street corner",
|
| 21 |
+
"source": "https://www.archives.gov/exhibits/picturing_the_century/greatdep/greatdep_img44.html",
|
| 22 |
+
"image_url": "https://www.archives.gov/exhibits/picturing_the_century/images/greatdep_044_v73.jpg",
|
| 23 |
+
"sha256": "c5e9a3fa57a6a0789e4b1529910e0294d98f245b24d205a1caccd08a5a3069b2",
|
| 24 |
+
"bytes": 68857
|
| 25 |
+
},
|
| 26 |
+
{
|
| 27 |
+
"name": "nara_45.jpg",
|
| 28 |
+
"title": "Children and sugar beets",
|
| 29 |
+
"source": "https://www.archives.gov/exhibits/picturing_the_century/greatdep/greatdep_img45.html",
|
| 30 |
+
"image_url": "https://www.archives.gov/exhibits/picturing_the_century/images/greatdep_045_v75.jpg",
|
| 31 |
+
"sha256": "3a3e3900cdb1384ca028a0a8843dcbcf499ffdd34832bc6b2e3f477d520403ae",
|
| 32 |
+
"bytes": 66706
|
| 33 |
+
},
|
| 34 |
+
{
|
| 35 |
+
"name": "nara_46.jpg",
|
| 36 |
+
"title": "Abandoned house",
|
| 37 |
+
"source": "https://www.archives.gov/exhibits/picturing_the_century/greatdep/greatdep_img46.html",
|
| 38 |
+
"image_url": "https://www.archives.gov/exhibits/picturing_the_century/images/greatdep_046_v76.jpg",
|
| 39 |
+
"sha256": "e7f2881c420c51c1855eb64219dc8c31e9a112976a3c81bfe0bcacb690d6b6e8",
|
| 40 |
+
"bytes": 74972
|
| 41 |
+
},
|
| 42 |
+
{
|
| 43 |
+
"name": "nara_49.jpg",
|
| 44 |
+
"title": "Man in his home",
|
| 45 |
+
"source": "https://www.archives.gov/exhibits/picturing_the_century/greatdep/greatdep_img49.html",
|
| 46 |
+
"image_url": "https://www.archives.gov/exhibits/picturing_the_century/images/greatdep_049_v72.jpg",
|
| 47 |
+
"sha256": "6d46bcc386aabf2dc090b2a05baf2e3ead04bc1274c38d87a60e1806ecfb4bab",
|
| 48 |
+
"bytes": 54941
|
| 49 |
+
}
|
| 50 |
+
]
|
experiments/decision-20260928-v2/supplement/resolution_review.py
ADDED
|
@@ -0,0 +1,27 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Resolution-only review on fixed historical inputs; no training or pixel editing."""
|
| 2 |
+
import json,time
|
| 3 |
+
from pathlib import Path
|
| 4 |
+
import torch
|
| 5 |
+
from PIL import Image,ImageDraw,ImageOps
|
| 6 |
+
from decision_model import load_decision
|
| 7 |
+
from inference import colorize,prepare
|
| 8 |
+
root=Path(__file__).resolve().parent
|
| 9 |
+
torch.set_num_threads(4)
|
| 10 |
+
m=load_decision(root/'smoke/mobilenet1')
|
| 11 |
+
records=json.loads((root/'archive_inputs/nara_sources.json').read_text())
|
| 12 |
+
rows=[]
|
| 13 |
+
for start in range(0,len(records),3):
|
| 14 |
+
canvas=Image.new('RGB',(1280,344*len(records[start:start+3])),'white');draw=ImageDraw.Draw(canvas)
|
| 15 |
+
for j,item in enumerate(records[start:start+3]):
|
| 16 |
+
im=Image.open(root/'archive_inputs'/item['name']).convert('L').convert('RGB')
|
| 17 |
+
draw.text((3,j*344+3),item['title']+' | Input | 256 | 384 | 512',fill='black')
|
| 18 |
+
canvas.paste(ImageOps.pad(im,(320,320),color='white'),(0,j*344+24))
|
| 19 |
+
for col,size in enumerate([256,384,512],1):
|
| 20 |
+
t=time.monotonic();pred=colorize(m,im,size=size,radius=8)
|
| 21 |
+
canvas.paste(ImageOps.pad(pred,(320,320),color='white'),(col*320,j*344+24))
|
| 22 |
+
with torch.no_grad():
|
| 23 |
+
details=m(prepare(im,size)[2],True);w=details['weights'].clamp_min(1e-8)
|
| 24 |
+
rows.append({'name':item['name'],'size':size,'seconds_cpu':time.monotonic()-t,'mask_entropy':float(-(w*w.log()).sum(1).mean()),'output_chroma':float(details['output'].norm(dim=1).mean())})
|
| 25 |
+
canvas.save(root/f'baseline_nara_resolution_{start//3}.jpg',quality=94)
|
| 26 |
+
(root/'baseline_nara_resolution.json').write_text(json.dumps(rows,indent=2))
|
| 27 |
+
print('Rendered',len(records),'historical photographs at three resolutions.')
|