| |
| """Easy4K SFDP Siamese NET AI pipeline (compact).""" |
| import argparse, csv, hashlib, http.cookiejar, json, math, os, random, re, subprocess, sys, time |
| import urllib.error, urllib.parse, urllib.request |
| from functools import lru_cache |
|
|
| import cv2, numpy as np, torch, torch.nn as nn |
| from torch.utils.data import DataLoader, Dataset, Sampler |
|
|
| class _Tee: |
| """把 stdout/stderr 同时写到一份可由外部读取的实时日志。""" |
| def __init__(self, *streams): |
| self.streams = streams |
|
|
| def write(self, s): |
| for st in self.streams: |
| try: |
| st.write(s) |
| except Exception: |
| pass |
| return len(s) |
|
|
| def flush(self): |
| for st in self.streams: |
| try: |
| st.flush() |
| except Exception: |
| pass |
|
|
|
|
| _live = os.environ.get("SFDP_LIVE_LOG") |
| if _live: |
| try: |
| _lf = open(_live, "a", encoding="utf-8", buffering=1) |
| sys.stdout = _Tee(sys.stdout, _lf) |
| sys.stderr = _Tee(sys.stderr, _lf) |
| except Exception: |
| pass |
|
|
|
|
| DATA = r"D:\sfdp-ai\siamese\data" |
| TILES = os.environ.get("SFDP_TILES") or os.path.join(DATA, "tiles") |
| LABELS = os.path.join(DATA, "labels") |
| TEACHER = r"D:\sfdp-ai\bin\SfdpTeacher.exe" |
| RAW = os.path.join(DATA, "raw") |
| BILI_RAW = os.path.join(DATA, "raw_bili") |
| TEST_TILES = os.path.join(DATA, "tiles_test") |
| TEST_LABELS = os.path.join(DATA, "labels_test") |
|
|
|
|
| class DSConv(nn.Module): |
| def __init__(self, cin, cout, stride=1): |
| super().__init__() |
| self.dw = nn.Conv2d(cin, cin, 3, stride, 1, groups=cin, bias=False) |
| self.bn1 = nn.BatchNorm2d(cin) |
| self.pw = nn.Conv2d(cin, cout, 1, bias=False) |
| self.bn2 = nn.BatchNorm2d(cout) |
| self.act = nn.ReLU6(inplace=True) |
|
|
| def forward(self, x): |
| x = self.act(self.bn1(self.dw(x))) |
| return self.act(self.bn2(self.pw(x))) |
|
|
|
|
| class Encoder(nn.Module): |
| def __init__(self, emb=512): |
| super().__init__() |
| self.stem = nn.Sequential(nn.Conv2d(1, 32, 3, 2, 1, bias=False), |
| nn.BatchNorm2d(32), nn.ReLU6(inplace=True)) |
| self.blocks = nn.Sequential( |
| DSConv(32, 48, 2), DSConv(48, 48), DSConv(48, 96, 2), DSConv(96, 96), |
| DSConv(96, 160, 2), DSConv(160, 160), DSConv(160, 224, 2), DSConv(224, 224), |
| DSConv(224, 320, 2), DSConv(320, 320)) |
| self.proj = nn.Conv2d(320, emb, 1, bias=False) |
| self.bn = nn.BatchNorm2d(emb) |
| self.act = nn.ReLU6(inplace=True) |
|
|
| def forward(self, x): |
| x = self.act(self.bn(self.proj(self.blocks(self.stem(x))))) |
| return x.mean(dim=(2, 3)) |
|
|
|
|
| class SiameseSFDP(nn.Module): |
| def __init__(self, emb=512, hidden=1024): |
| super().__init__() |
| self.encoder = Encoder(emb) |
| self.head = nn.Sequential(nn.Linear(emb * 4, hidden), nn.ReLU(inplace=True), |
| nn.Dropout(0.1), nn.Linear(hidden, 1), nn.Sigmoid()) |
|
|
| def forward(self, a, b): |
| ea = self.encoder(a); eb = self.encoder(b) |
| return self.head(torch.cat([ea, eb, (ea - eb).abs(), ea * eb], dim=1)) |
|
|
|
|
| def count_params(m): |
| return sum(p.numel() for p in m.parameters()) |
|
|
|
|
| @lru_cache(maxsize=4096) |
| def load_gray(p): |
| img = cv2.imread(p, cv2.IMREAD_GRAYSCALE) |
| if img is None: |
| raise FileNotFoundError(p) |
| return img |
|
|
|
|
| _GRAY_CACHE = {} |
|
|
|
|
| def load_gray_cached(path): |
| """同一切片会被同片段的多个样本反复读取：按 worker 进程做一个小 LRU， |
| 省掉重复 PNG 解码（320x320 约 100KB，512 条 ≈ 50MB/worker）。""" |
| a = _GRAY_CACHE.get(path) |
| if a is not None: |
| return a |
| try: |
| a = load_gray(path) |
| except FileNotFoundError: |
| |
| a = None |
| for _ in range(3): |
| time.sleep(0.2) |
| try: |
| a = load_gray(path) |
| break |
| except FileNotFoundError: |
| continue |
| if a is None: |
| a = np.zeros((320, 320), np.uint8) |
| if len(_GRAY_CACHE) >= 512: |
| _GRAY_CACHE.clear() |
| _GRAY_CACHE[path] = a |
| return a |
|
|
|
|
| FALLBACK_TILES = os.path.join(DATA, "tiles") |
|
|
|
|
| def _resolve_tile(root, rel): |
| """Hot tiles (SSD) may still be copying: fall back to the original tiles dir.""" |
| p = os.path.join(root, rel) |
| if os.path.exists(p): |
| return p |
| fb = os.path.join(FALLBACK_TILES, rel) |
| return fb if os.path.exists(fb) else p |
|
|
|
|
| class PairDataset(Dataset): |
| def __init__(self, manifest, tiles_root, crop=128, augment=False): |
| self.items = [json.loads(l) for l in open(manifest, "r", encoding="utf-8") if l.strip()] |
| self.tiles_root = tiles_root; self.crop = crop; self.augment = augment |
|
|
| def __len__(self): |
| return len(self.items) |
|
|
| def __getitem__(self, i): |
| it = self.items[i] |
| a = load_gray_cached(_resolve_tile(self.tiles_root, it["a"])) |
| b = load_gray_cached(_resolve_tile(self.tiles_root, it["b"])) |
| c = self.crop; h, w = a.shape |
| x = min(max(0, it["x"]), w - c); y = min(max(0, it["y"]), h - c) |
| pa, pb = a[y:y + c, x:x + c], b[y:y + c, x:x + c] |
| if self.augment: |
| if np.random.rand() < 0.5: |
| pa, pb = pa[:, ::-1], pb[:, ::-1] |
| if np.random.rand() < 0.3: |
| g = 0.8 + np.random.rand() * 0.45 |
| pa = np.clip((pa / 255.0) ** g * 255.0, 0, 255).astype(np.uint8) |
| pb = np.clip((pb / 255.0) ** g * 255.0, 0, 255).astype(np.uint8) |
| ta = torch.from_numpy(np.ascontiguousarray(pa)).float().div_(255.0).unsqueeze(0) |
| tb = torch.from_numpy(np.ascontiguousarray(pb)).float().div_(255.0).unsqueeze(0) |
| return (ta, tb, torch.tensor([it["soft"]], dtype=torch.float32), |
| torch.tensor([it["hard"]], dtype=torch.float32)) |
|
|
|
|
| |
| BUA = ("Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 " |
| "(KHTML, like Gecko) Chrome/124.0 Safari/537.36") |
| BILI_META = os.path.join(BILI_RAW, "sources_bili.json") |
| _jar = http.cookiejar.CookieJar() |
| _op = urllib.request.build_opener(urllib.request.HTTPCookieProcessor(_jar)) |
| _op.addheaders = [("User-Agent", BUA), ("Referer", "https://www.bilibili.com/")] |
|
|
|
|
| def bjson(url, params=None): |
| full = url + ("?" + urllib.parse.urlencode(params) if params else "") |
| last = None |
| for i in range(3): |
| try: |
| with _op.open(full, timeout=25) as r: |
| return json.loads(r.read().decode("utf-8")) |
| except urllib.error.HTTPError as e: |
| if e.code == 412: |
| raise |
| last = e; time.sleep(2 + i * 2) |
| except Exception as e: |
| last = e; time.sleep(2 + i * 2) |
| raise last |
|
|
|
|
| def bview(bvid): |
| d = bjson("https://api.bilibili.com/x/web-interface/view", {"bvid": bvid}) |
| return (d.get("data") or None) if d.get("code") == 0 else None |
|
|
|
|
| def bplay(bvid, cid, qn): |
| d = bjson("https://api.bilibili.com/x/player/playurl", |
| {"bvid": bvid, "cid": cid, "qn": qn, "fnval": 0, "fourk": 0}) |
| if d.get("code") != 0: |
| return None |
| dd = d.get("data") or {}; durl = dd.get("durl") or [] |
| return {"quality": dd.get("quality"), "url": durl[0].get("url")} if durl else None |
|
|
|
|
| BILI_MAX_BPS = 6_000_000 |
|
|
|
|
| def bdownload(url, dst): |
| tmp = dst + ".part" |
| req = urllib.request.Request(url, headers={"User-Agent": BUA, "Referer": "https://www.bilibili.com/"}) |
| for _ in range(3): |
| try: |
| with urllib.request.urlopen(req, timeout=30) as r, open(tmp, "wb") as f: |
| t0 = time.time() |
| got = 0 |
| while True: |
| ch = r.read(262144) |
| if not ch: |
| break |
| f.write(ch) |
| got += len(ch) |
| |
| expected = got / BILI_MAX_BPS |
| elapsed = time.time() - t0 |
| if expected > elapsed: |
| time.sleep(expected - elapsed) |
| os.replace(tmp, dst); return True |
| except Exception: |
| if os.path.exists(tmp): |
| os.remove(tmp) |
| time.sleep(3) |
| return False |
|
|
|
|
| def cmd_rank(args): |
| root = getattr(args, "outdir", None) or BILI_RAW |
| meta_path = os.path.join(root, "sources_bili.json") |
| os.makedirs(root, exist_ok=True) |
| try: |
| _op.open("https://www.bilibili.com/", timeout=20).read(2048) |
| except Exception: |
| pass |
| meta = json.load(open(meta_path, "r", encoding="utf-8")) if os.path.exists(meta_path) else [] |
| done = {m["bvid"] for m in meta} |
| seen, cands = set(), [] |
|
|
| def add(items): |
| for it in items: |
| b = it.get("bvid"); dur = int(it.get("duration") or 0); title = it.get("title", "") |
| if not b or b in seen or b in done or not (args.min_sec <= dur <= args.max_sec): |
| continue |
| if any(x in title for x in ("直播", "回放", "合集", "教程", "reaction", "解说")): |
| continue |
| seen.add(b) |
| cands.append({"bvid": b, "title": title, "author": (it.get("owner") or {}).get("name", ""), |
| "duration": dur}) |
|
|
| for rid in (3, 129, 1, 181, 71, 160, 4): |
| try: |
| d = bjson("https://api.bilibili.com/x/web-interface/ranking/v2", {"rid": rid}) |
| add((d.get("data") or {}).get("list") or []) |
| except Exception as e: |
| print("RANK_FAIL", rid, e, flush=True) |
| time.sleep(1.2) |
| for pn in range(1, args.pages + 1): |
| try: |
| d = bjson("https://api.bilibili.com/x/web-interface/popular", {"ps": 20, "pn": pn}) |
| add((d.get("data") or {}).get("list") or []) |
| except Exception as e: |
| print("POPULAR_FAIL", pn, e, flush=True) |
| time.sleep(1.2) |
| for rid in (3, 129, 1, 181): |
| for pn in range(1, 6): |
| try: |
| d = bjson("https://api.bilibili.com/x/web-interface/dynamic/region", |
| {"rid": rid, "ps": 20, "pn": pn}) |
| add((d.get("data") or {}).get("archives") or []) |
| except Exception as e: |
| print("REGION_FAIL", rid, pn, e, flush=True); break |
| time.sleep(1.0) |
| existing = list(done); random.shuffle(existing) |
| seeds = ([c["bvid"] for c in cands] + existing[:args.seed_existing])[:args.seeds + args.seed_existing] |
| for b in seeds: |
| try: |
| d = bjson("https://api.bilibili.com/x/web-interface/archive/related", {"bvid": b}) |
| add(d.get("data") or []) |
| except Exception as e: |
| print("RELATED_FAIL", b, e, flush=True) |
| time.sleep(0.5) |
| print("CANDIDATES", len(cands), flush=True) |
| saved = 0 |
| for it in cands: |
| if saved >= args.target: |
| break |
| v = bview(it["bvid"]) |
| if not v: |
| continue |
| pu = None |
| for qn in (80, 64, 32): |
| pu = bplay(it["bvid"], v["cid"], qn) |
| if pu: |
| break |
| if not pu or not pu.get("url"): |
| continue |
| name = re.sub(r"[^0-9A-Za-z._-]+", "_", it["title"])[:60] or it["bvid"] |
| ext = ".flv" if ".flv" in pu["url"] else ".mp4" |
| fname = f"{it['bvid']}_{name}{ext}" |
| dst = os.path.join(root, fname) |
| if os.path.exists(dst): |
| continue |
| if not bdownload(pu["url"], dst): |
| print("DL_FAIL", it["bvid"], flush=True); continue |
| meta.append({**it, "file": fname, "cid": v["cid"], "qn": pu["quality"], |
| "license": "research use only, not redistributed"}) |
| saved += 1 |
| print("OK", fname, flush=True) |
| json.dump(meta, open(meta_path, "w", encoding="utf-8"), ensure_ascii=False, indent=2) |
| print("TOTAL_FILES", len(meta), flush=True) |
|
|
|
|
| |
| PATTERNS = {"1-1": [1], "2-2": [2], "3-3": [3], "mix": [2, 3, 2, 1, 3, 2]} |
| VEXT = (".webm", ".ogg", ".ogv", ".mp4", ".mov", ".mkv", ".flv") |
|
|
|
|
| def clip_id(stem, seg, pat): |
| s = re.sub(r"[^0-9A-Za-z._-]+", "_", stem).strip("_")[:40] |
| if any(ord(c) > 127 for c in stem) or len(s) < 4: |
| s = (s + "_" if s else "mv_") + hashlib.sha1(stem.encode("utf-8")).hexdigest()[:8] |
| return f"{s}_s{seg:02d}_{pat}" |
|
|
|
|
| def pulldown(frames, rounds, n): |
| out, i = [], 0 |
| while len(out) < n: |
| out.extend([frames[i % len(frames)]] * rounds[i % len(rounds)]); i += 1 |
| return out[:n] |
|
|
|
|
| def cmd_tiles(args): |
| rng = random.Random(args.seed); os.makedirs(args.out, exist_ok=True) |
| idx_path = os.path.join(args.out, "index.jsonl") |
| done = set() |
| if os.path.exists(idx_path): |
| for line in open(idx_path, "r", encoding="utf-8"): |
| try: |
| done.add(json.loads(line)["clip"]) |
| except Exception: |
| pass |
| videos = [n for n in sorted(os.listdir(args.raw)) if os.path.splitext(n)[1].lower() in VEXT] |
| print("VIDEOS", len(videos), "ALREADY", len(done), flush=True) |
| total = 0 |
| with open(idx_path, "a", encoding="utf-8") as idx: |
| for vi, name in enumerate(videos): |
| cap = cv2.VideoCapture(os.path.join(args.raw, name)) |
| if not cap.isOpened(): |
| print("OPEN_FAIL", name, flush=True); continue |
| nf = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0) |
| fps = float(cap.get(cv2.CAP_PROP_FPS) or 0) |
| w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) or 0) |
| h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) or 0) |
| if nf < args.frames * 3 or w < max(args.tile, args.min_width) or h < args.tile: |
| print("SKIP_DIM", name, w, h, nf, flush=True); cap.release(); continue |
| max_start = max(0, nf - args.frames * 4) |
| for si in range(args.segments): |
| start = rng.randint(0, max_start) |
| x = rng.randint(0, w - args.tile); y = rng.randint(0, h - args.tile) |
| cap.set(cv2.CAP_PROP_POS_FRAMES, start) |
| frames = [] |
| for _ in range(max(4, args.frames // 2)): |
| ok, fr = cap.read() |
| if not ok: |
| break |
| frames.append(cv2.cvtColor(fr[y:y + args.tile, x:x + args.tile], cv2.COLOR_BGR2GRAY)) |
| if len(frames) < 4: |
| continue |
| for pat, rounds in PATTERNS.items(): |
| clip = clip_id(os.path.splitext(name)[0], si, pat) |
| if clip in done: |
| continue |
| cdir = os.path.join(args.out, clip); os.makedirs(cdir, exist_ok=True) |
| for k, img in enumerate(pulldown(frames, rounds, args.frames)): |
| cv2.imwrite(os.path.join(cdir, f"f{k:04d}.png"), img) |
| idx.write(json.dumps({"clip": clip, "src": name, "pattern": pat, |
| "start_frame": start, "x": x, "y": y, |
| "tile": args.tile, "frames": args.frames, "fps": fps}, |
| ensure_ascii=False) + "\n") |
| idx.flush(); total += 1 |
| cap.release() |
| if (vi + 1) % 10 == 0: |
| print("PROGRESS", vi + 1, "/", len(videos), "clips", total, flush=True) |
| print("TOTAL_CLIPS", total, flush=True) |
|
|
|
|
| def cmd_teacher(args): |
| os.makedirs(args.labels, exist_ok=True) |
| print("SCAN_START", args.tiles, flush=True) |
| dirs = [d for d in sorted(os.listdir(args.tiles)) if os.path.isdir(os.path.join(args.tiles, d))] |
| print("SCAN_DIRS", len(dirs), flush=True) |
| todo = [] |
| for i, d in enumerate(dirs): |
| if not os.path.exists(os.path.join(args.labels, d + ".csv")): |
| todo.append(d) |
| if (i + 1) % 2000 == 0: |
| print("SCAN", i + 1, "/", len(dirs), flush=True) |
| print("DIRS", len(dirs), "TODO", len(todo), flush=True) |
| for i, d in enumerate(todo): |
| subprocess.run([TEACHER, os.path.join(args.tiles, d), os.path.join(args.labels, d + ".csv"), |
| "--mode", "uhd", "--crop", "128", "--per-pair", "4", "--limit", "0"], |
| stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) |
| if (i + 1) % 10 == 0: |
| print("PROGRESS", i + 1, "/", len(todo), flush=True) |
| print("DONE", flush=True) |
|
|
|
|
| |
| def soft_from_flow(flow): |
| return max(0.0, min(1.0, 1.0 - flow / 2.0)) |
|
|
|
|
| def cmd_label(args): |
| rng = random.Random(args.seed) |
| clip_src = {} |
| p = os.path.join(TILES, "index.jsonl") |
| if os.path.exists(p): |
| for line in open(p, "r", encoding="utf-8"): |
| try: |
| r = json.loads(line); clip_src[r["clip"]] = r["src"] |
| except Exception: |
| pass |
| rows = [] |
| for name in sorted(os.listdir(LABELS)): |
| if not name.endswith(".csv"): |
| continue |
| clip = name[:-4]; src = clip_src.get(clip, "unknown") |
| with open(os.path.join(LABELS, name), "r", encoding="utf-8", newline="") as f: |
| for r in csv.DictReader(f): |
| try: |
| mad, flow, var = float(r["mad"]), float(r["flow"]), float(r["var"]) |
| stage = r.get("stage", "") or "" |
| dup = int(float(r.get("dup", 0) or 0)) |
| x, y = int(float(r["x"])), int(float(r["y"])) |
| except (KeyError, ValueError, TypeError): |
| continue |
| if dup == 1: |
| soft = 1.0 if (flow == 0.0 and var == 0.0) else soft_from_flow(flow) |
| elif stage == "Flow": |
| if 0.1 <= flow <= 0.4: |
| continue |
| soft = soft_from_flow(flow) |
| else: |
| soft = 0.0 |
| rows.append({"a": os.path.join(clip, r["frame_a"]), "b": os.path.join(clip, r["frame_b"]), |
| "x": x, "y": y, "crop": 128, |
| "soft": round(soft, 4), "hard": 1 if (stage == "Flow" and mad < 1.0) else 0, |
| "stage": stage, "src": src, "clip": clip}) |
| pos = [r for r in rows if r["soft"] >= 0.9] |
| neg = [r for r in rows if r["soft"] <= 0.1] |
| mid = [r for r in rows if 0.1 < r["soft"] < 0.9] |
| rng.shuffle(pos); rng.shuffle(neg); rng.shuffle(mid) |
| neg.sort(key=lambda r: -r["hard"]) |
| m = int(len(pos) * args.max_neg_ratio) |
| neg_keep = neg[:m] if len(neg) > m else neg |
| if neg_keep and len(pos) > len(neg_keep): |
| pos = pos[:len(neg_keep)] |
| mid_keep = mid[:int((len(pos) + len(neg_keep)) * args.mid_ratio)] |
| data = pos + neg_keep + mid_keep; rng.shuffle(data) |
| val_f = lambda r: sum(ord(c) for c in r["src"]) % 7 == 0 |
| train = [r for r in data if not val_f(r)]; val = [r for r in data if val_f(r)] |
| for path, items in ((os.path.join(DATA, "train.jsonl"), train), (os.path.join(DATA, "val.jsonl"), val)): |
| with open(path, "w", encoding="utf-8") as f: |
| for it in items: |
| f.write(json.dumps(it, ensure_ascii=False) + "\n") |
| st = {"raw": len(rows), "kept_pos": len(pos), "kept_neg": len(neg_keep), |
| "kept_mid": len(mid_keep), "hard_mined": sum(1 for r in data if r["hard"]), |
| "train": len(train), "val": len(val)} |
| json.dump(st, open(os.path.join(DATA, "label_stats.json"), "w", encoding="utf-8"), |
| ensure_ascii=False, indent=2) |
| print(json.dumps(st, ensure_ascii=False), flush=True) |
|
|
|
|
| def eval_model(model, loader, device): |
| model.eval(); se = n = correct = pe = pn = ne = nn_ = 0 |
| with torch.no_grad(): |
| for a, b, y, _ in loader: |
| a, b, y = a.to(device), b.to(device), y.to(device) |
| p = model(a, b) |
| se += torch.abs(p - y).sum().item(); n += y.numel() |
| correct += ((p >= 0.5) == (y >= 0.5)).sum().item() |
| mm = y >= 0.9 |
| if mm.any(): |
| pe += torch.abs(p[mm] - y[mm]).sum().item(); pn += int(mm.sum()) |
| mm = y <= 0.1 |
| if mm.any(): |
| ne += torch.abs(p[mm] - y[mm]).sum().item(); nn_ += int(mm.sum()) |
| return {"mae": se / max(1, n), "acc": correct / max(1, n), |
| "pos_mae": pe / max(1, pn), "neg_mae": ne / max(1, nn_), "n": n} |
|
|
|
|
| class ClipGroupSampler(Sampler): |
| """按片段分组取样：同一 clip 的切片连续读，机械盘上把随机读变成顺序读。 |
| 片段之间每个 epoch 重新打乱、片段内部也打乱，样本分布不变。""" |
|
|
| def __init__(self, dataset, shuffle=True, seed=1234): |
| self.groups = {} |
| for i, it in enumerate(dataset.items): |
| self.groups.setdefault(it["clip"], []).append(i) |
| self.clips = sorted(self.groups.keys()) |
| self.shuffle = shuffle |
| self.seed = seed |
| self.epoch = 0 |
| self.n = len(dataset) |
|
|
| def set_epoch(self, e): |
| self.epoch = e |
|
|
| def __len__(self): |
| return self.n |
|
|
| def __iter__(self): |
| rng = random.Random(self.seed + self.epoch) |
| clips = self.clips[:] |
| if self.shuffle: |
| rng.shuffle(clips) |
| for c in clips: |
| idx = self.groups[c] |
| if self.shuffle: |
| idx = idx[:] |
| rng.shuffle(idx) |
| for i in idx: |
| yield i |
|
|
|
|
| def _worker_init(_): |
| |
| |
| try: |
| torch.set_num_threads(1) |
| except Exception: |
| pass |
|
|
|
|
| def cmd_train(args): |
| random.seed(args.seed); np.random.seed(args.seed) |
| torch.manual_seed(args.seed); torch.cuda.manual_seed_all(args.seed) |
| torch.backends.cudnn.benchmark = True |
| os.makedirs(args.out, exist_ok=True) |
| dev = "cuda" if torch.cuda.is_available() else "cpu" |
| print("DEVICE", dev, torch.cuda.get_device_name(0) if dev == "cuda" else "", flush=True) |
| tr = PairDataset(os.path.join(DATA, "train.jsonl"), TILES, augment=bool(args.augment)) |
| va = PairDataset(os.path.join(DATA, "val.jsonl"), TILES) |
| print("SAMPLES", len(tr), len(va), flush=True) |
| tl = DataLoader(tr, batch_size=args.batch, sampler=ClipGroupSampler(tr, shuffle=True), |
| num_workers=args.workers, worker_init_fn=_worker_init, prefetch_factor=2, |
| pin_memory=(dev == "cuda"), drop_last=True, persistent_workers=args.workers > 0) |
| vl = DataLoader(va, batch_size=args.batch, sampler=ClipGroupSampler(va, shuffle=False), |
| num_workers=max(2, args.workers // 2), worker_init_fn=_worker_init, |
| pin_memory=(dev == "cuda")) |
| model = SiameseSFDP().to(dev); print("PARAMS", count_params(model), flush=True) |
| opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4) |
| steps = max(1, len(tl)); total = args.epochs * steps |
|
|
| def lr_at(step): |
| warm = args.warmup * steps |
| if step < warm: |
| return (step + 1) / max(1, warm) |
| return 0.5 * (1 + math.cos(math.pi * (step - warm) / max(1, total - warm))) |
|
|
| sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_at) |
| best = 1e9 |
| if args.resume and os.path.exists(args.resume): |
| ck = torch.load(args.resume, map_location=dev, weights_only=True) |
| model.load_state_dict(ck["model"]); opt.load_state_dict(ck["opt"]) |
| best = ck.get("best", best) |
| print("RESUME", args.resume, "prev", ck.get("epoch"), flush=True) |
| log = os.path.join(args.out, "train_log.csv") |
| if not os.path.exists(log): |
| with open(log, "w", newline="", encoding="utf-8") as f: |
| csv.writer(f).writerow(["epoch", "lr", "loss", "val_mae", "val_acc", "pos_mae", "neg_mae", "sec"]) |
| for ep in range(args.epochs): |
| if hasattr(tl.sampler, "set_epoch"): |
| tl.sampler.set_epoch(ep) |
| model.train(); t0 = time.time(); run = cnt = 0 |
| se = 0 |
| step_i = 0 |
| for a, b, y, hard in tl: |
| a, b, y, hard = a.to(dev), b.to(dev), y.to(dev), hard.to(dev) |
| p = model(a, b); w = 1.0 + (args.hard_weight - 1.0) * hard |
| loss = (w * (p - y) ** 2).mean() |
| opt.zero_grad(set_to_none=True); loss.backward() |
| nn.utils.clip_grad_norm_(model.parameters(), 5.0) |
| opt.step(); sched.step() |
| run += loss.item() * a.size(0); cnt += a.size(0) |
| se += torch.abs(p - y).sum().item() |
| step_i += 1 |
| if step_i % 50 == 0: |
| print(json.dumps({"step": step_i, "steps": len(tl), "ep": ep, |
| "loss": round(run / max(1, cnt), 5), |
| "pmae": round(se / max(1, cnt), 5), |
| "elapsed": round(time.time() - t0, 1)}), flush=True) |
| m = eval_model(model, vl, dev); sec = time.time() - t0; lr = opt.param_groups[0]["lr"] |
| print(json.dumps({"epoch": ep, "loss": round(run / max(1, cnt), 5), "lr": round(lr, 6), |
| **{k: round(v, 5) for k, v in m.items()}, "sec": round(sec, 1)}), flush=True) |
| with open(log, "a", newline="", encoding="utf-8") as f: |
| csv.writer(f).writerow([ep, round(lr, 6), round(run / max(1, cnt), 5), round(m["mae"], 5), |
| round(m["acc"], 5), round(m["pos_mae"], 5), round(m["neg_mae"], 5), |
| round(sec, 1)]) |
| ck = {"model": model.state_dict(), "opt": opt.state_dict(), "epoch": ep, "best": best} |
| torch.save(ck, os.path.join(args.out, "last.pt")) |
| if m["mae"] < best: |
| best = m["mae"]; ck["best"] = best |
| torch.save(ck, os.path.join(args.out, "best.pt")) |
| print("BEST", round(best, 5), "epoch", ep, flush=True) |
| print("TRAIN_DONE best_mae", round(best, 5), flush=True) |
| ck_best = os.path.join(args.out, "best.pt") |
| if not os.path.exists(ck_best): |
| ck_best = os.path.join(args.out, "last.pt") |
| print("EXPORT_SRC", ck_best, flush=True) |
| export_onnx(ck_best, os.path.join(args.out, "easy4k_sfdp_siamese_fp32.onnx")) |
|
|
|
|
| def export_onnx(ckpt, out, opset=17): |
| import onnx |
| ck = torch.load(ckpt, map_location="cpu", weights_only=True) |
| m = SiameseSFDP(); m.load_state_dict(ck["model"]); m.eval() |
| a = torch.rand(2, 1, 128, 128); b = torch.rand(2, 1, 128, 128) |
| torch.onnx.export(m, (a, b), out, input_names=["frame_a", "frame_b"], output_names=["score"], |
| dynamic_axes={"frame_a": {0: "batch"}, "frame_b": {0: "batch"}, "score": {0: "batch"}}, |
| opset_version=opset, do_constant_folding=True) |
| o = onnx.load(out); onnx.checker.check_model(o) |
| meta = {"model_name": "Easy4K SFDP - Siamese NET AI", "model_version": "0.6.0", |
| "task": "frame-duplicate similarity regression", |
| "input": "two grayscale 128x128 patches, float32, [0,1]", |
| "output": "score in [0,1]; 1 = duplicate/same picture", |
| "training": "knowledge distillation from Easy4K SFDP four-stage judge", |
| "runtime": "onnxruntime DirectML (NVIDIA / AMD / Intel) or CPU", |
| "license": "cc-by-nc-4.0", "params": str(count_params(m))} |
| have = {p.key for p in o.metadata_props} |
| for k, v in meta.items(): |
| if k not in have: |
| p = o.metadata_props.add(); p.key, p.value = k, v |
| onnx.save(o, out) |
| print("EXPORTED", out, "size_mb", round(os.path.getsize(out) / 1048576, 2), flush=True) |
|
|
| |
| try: |
| from onnxconverter_common import float16 as _f16 |
| o16 = _f16.convert_float_to_float16(onnx.load(out), keep_io_types=True) |
| p16 = out.replace("_fp32.onnx", "_fp16.onnx") if out.endswith("_fp32.onnx") else (out[:-5] + "_fp16.onnx") |
| onnx.save(o16, p16) |
| print("EXPORTED_FP16", p16, "size_mb", round(os.path.getsize(p16) / 1048576, 2), flush=True) |
| except Exception as e: |
| print("EXPORTED_FP16_SKIP", repr(e), flush=True) |
|
|
|
|
| def cmd_eval(args): |
| ds = PairDataset(os.path.join(DATA, "val.jsonl"), TILES) |
| dl = DataLoader(ds, batch_size=512, shuffle=False, num_workers=6) |
| ck = torch.load(args.ckpt, map_location="cpu", weights_only=True) |
| net = SiameseSFDP(); net.load_state_dict(ck["model"]); net.eval() |
| S, Y, ST = [], [], [] |
| with torch.no_grad(): |
| off = 0 |
| for a, b, y, _ in dl: |
| p = net(a, b).squeeze(1).numpy(); S.append(p); Y.append(y.squeeze(1).numpy()) |
| ST.extend([ds.items[off + i]["stage"] for i in range(len(p))]); off += len(p) |
| s = np.concatenate(S); y = np.concatenate(Y); t = (y >= 0.5).astype(int); p = (s >= 0.5).astype(int) |
| tp = int(((p == 1) & (t == 1)).sum()); fp = int(((p == 1) & (t == 0)).sum()) |
| fn = int(((p == 0) & (t == 1)).sum()); tn = int(((p == 0) & (t == 0)).sum()) |
| print(json.dumps({"n": len(s), "acc": round((tp + tn) / len(s), 5), |
| "precision": round(tp / max(1, tp + fp), 5), |
| "recall": round(tp / max(1, tp + fn), 5)}), flush=True) |
| by = {} |
| for st in sorted(set(ST)): |
| mm = np.array([x == st for x in ST]) |
| by[st] = {"n": int(mm.sum()), "mean_score": round(float(s[mm].mean()), 4), |
| "mean_target": round(float(y[mm].mean()), 4)} |
| print("BY_STAGE", json.dumps(by, ensure_ascii=False), flush=True) |
| for th in (0.4, 0.5, 0.6, 0.7, 0.8): |
| pp = (s >= th).astype(int) |
| tp_ = int(((pp == 1) & (t == 1)).sum()); fp_ = int(((pp == 1) & (t == 0)).sum()) |
| fn_ = int(((pp == 0) & (t == 1)).sum()) |
| pr = tp_ / max(1, tp_ + fp_); rc = tp_ / max(1, tp_ + fn_) |
| print(json.dumps({"t": th, "acc": round(float((pp == t).mean()), 5), "precision": round(pr, 5), |
| "recall": round(rc, 5), "f1": round(2 * pr * rc / max(1e-9, pr + rc), 5)}), flush=True) |
|
|
|
|
| def cmd_compare(args): |
| items = [] |
| for name in sorted(os.listdir(TEST_LABELS)): |
| if not name.endswith(".csv"): |
| continue |
| clip = name[:-4] |
| for r in csv.DictReader(open(os.path.join(TEST_LABELS, name), "r", encoding="utf-8", newline="")): |
| flow, var = float(r["flow"]), float(r["var"]); dup = int(float(r.get("dup", 0) or 0)) |
| stage = r.get("stage", "") |
| soft = (1.0 if (flow == 0.0 and var == 0.0) else soft_from_flow(flow)) if dup == 1 else \ |
| (soft_from_flow(flow) if stage == "Flow" else 0.0) |
| items.append((os.path.join(clip, r["frame_a"]), os.path.join(clip, r["frame_b"]), |
| int(float(r["x"])), int(float(r["y"])), soft)) |
| print("TEST_ITEMS", len(items), flush=True) |
| A, B = [], [] |
| for pa, pb, x, y, _ in items: |
| a = load_gray(os.path.join(TEST_TILES, pa)); b = load_gray(os.path.join(TEST_TILES, pb)) |
| h, w = a.shape; x = min(max(0, x), w - 128); y = min(max(0, y), h - 128) |
| A.append(torch.from_numpy(np.ascontiguousarray(a[y:y + 128, x:x + 128])).float().div_(255.0).unsqueeze(0)) |
| B.append(torch.from_numpy(np.ascontiguousarray(b[y:y + 128, x:x + 128])).float().div_(255.0).unsqueeze(0)) |
| A = torch.stack(A); B = torch.stack(B) |
| y = torch.tensor([it[4] for it in items], dtype=torch.float32).unsqueeze(1) |
| t = (y >= 0.5).int().numpy() |
| for path in args.ckpts.split(","): |
| path = path.strip() |
| if not os.path.exists(path): |
| print("MISSING", path, flush=True); continue |
| ck = torch.load(path, map_location="cpu", weights_only=True) |
| net = SiameseSFDP(); net.load_state_dict(ck["model"]); net.eval() |
| outs = [] |
| with torch.no_grad(): |
| for i in range(0, len(A), 512): |
| outs.append(net(A[i:i + 512], B[i:i + 512]).numpy()) |
| s = np.concatenate(outs) |
| rec = {"ckpt": os.path.basename(os.path.dirname(path)), "epoch": ck.get("epoch"), |
| "mae": round(float(np.abs(s - y.numpy()).mean()), 5)} |
| for th in (0.5, 0.6, 0.7): |
| pp = (s >= th).astype(int) |
| tp = int(((pp == 1) & (t == 1)).sum()); fp = int(((pp == 1) & (t == 0)).sum()) |
| fn = int(((pp == 0) & (t == 1)).sum()); tn = int(((pp == 0) & (t == 0)).sum()) |
| rec[f"acc@{th}"] = round((tp + tn) / len(s), 4) |
| print(rec, flush=True) |
|
|
|
|
| def cmd_pack(args): |
| import shutil |
| card = open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "card.md"), "r", encoding="utf-8").read() |
| os.makedirs(args.out, exist_ok=True) |
| shutil.copy2(args.onnx, os.path.join(args.out, "model.onnx")) |
| _p16 = args.onnx.replace("_fp32.onnx", "_fp16.onnx") |
| if os.path.exists(_p16): |
| shutil.copy2(_p16, os.path.join(args.out, "model_fp16.onnx")) |
| size = round(os.path.getsize(args.onnx) / 1048576, 2) |
| best = {} |
| if os.path.exists(args.log): |
| for row in csv.DictReader(open(args.log, "r", encoding="utf-8")): |
| if not best or float(row["val_mae"]) < float(best["val_mae"]): |
| best = row |
| ev = {} |
| if os.path.exists(args.eval): |
| for enc in ("utf-8-sig", "utf-16", "utf-16-le"): |
| try: |
| for line in open(args.eval, "r", encoding=enc, errors="ignore"): |
| line = line.strip() |
| if line.startswith("{") and '"precision"' in line: |
| ev = json.loads(line); break |
| except Exception: |
| pass |
| if ev: |
| break |
| st = json.load(open(os.path.join(DATA, "label_stats.json"), "r", encoding="utf-8")) |
| clips = len([d for d in os.listdir(TILES) if os.path.isdir(os.path.join(TILES, d))]) |
| pct = lambda x: ("%.2f%%" % (float(x) * 100)) if str(x) not in ("", "n/a") else "n/a" |
| vals = {"params": args.params, "opset": 17, "size_mb": size, "epochs": args.epochs, "clips": clips, |
| "samples": st.get("train", 0) + st.get("val", 0), "train": st.get("train"), "val": st.get("val"), |
| "ratio": "1 : %.2f" % (st.get("kept_neg", 0) / max(1, st.get("kept_pos", 1))), |
| "hard": st.get("hard_mined"), "mae": best.get("val_mae", "n/a"), |
| "acc": pct(best.get("val_acc", "n/a")), "precision": pct(ev.get("precision", "n/a")), |
| "recall": pct(ev.get("recall", "n/a")), "pos_mae": best.get("pos_mae", "n/a"), |
| "neg_mae": best.get("neg_mae", "n/a"), "amd_ms": args.amd_ms, "cpu_ms": args.cpu_ms} |
| for k, v in vals.items(): |
| card = card.replace("{" + k + "}", str(v)) |
| cfg = {"model_name": "Easy4K SFDP - Siamese NET AI", "model_version": "0.6.0", |
| "architecture": "siamese depthwise-separable CNN + difference/product fusion head", |
| "input": {"names": ["frame_a", "frame_b"], "shape": ["batch", 1, 128, 128], |
| "dtype": "float32", "range": [0, 1], "color": "grayscale"}, |
| "output": {"name": "score", "shape": ["batch", 1], "meaning": "1 = duplicate"}, |
| "opset": 17, "providers": ["DmlExecutionProvider", "CPUExecutionProvider"], |
| "recommended_threshold": 0.5, |
| "metrics": {"val_mae": float(best.get("val_mae", 0)), "val_acc": float(best.get("val_acc", 0)), |
| "precision": float(ev.get("precision", 0)), "recall": float(ev.get("recall", 0))}, |
| "license": "cc-by-nc-4.0"} |
| json.dump(cfg, open(os.path.join(args.out, "config.json"), "w", encoding="utf-8"), |
| ensure_ascii=False, indent=2) |
| open(os.path.join(args.out, "README.md"), "w", encoding="utf-8").write(card) |
| open(os.path.join(args.out, "LICENSE"), "w", encoding="utf-8").write( |
| "Easy4K SFDP - Siamese NET AI\nCopyright (c) 2026 Easy4K\nCC BY-NC 4.0\n" |
| "https://creativecommons.org/licenses/by-nc/4.0/legalcode\n") |
| open(os.path.join(args.out, "requirements.txt"), "w", encoding="utf-8").write( |
| "onnxruntime-directml>=1.18\nnumpy>=1.26\n") |
| src = os.path.join(RAW, "sources.json") |
| if os.path.exists(src): |
| shutil.copy2(src, os.path.join(args.out, "sources.json")) |
| print("PACKED", args.out, flush=True) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(); sub = ap.add_subparsers(dest="cmd", required=True) |
| p = sub.add_parser("rank") |
| p.add_argument("--target", type=int, default=40); p.add_argument("--pages", type=int, default=6) |
| p.add_argument("--min-sec", type=int, default=60); p.add_argument("--max-sec", type=int, default=420) |
| p.add_argument("--seeds", type=int, default=60); p.add_argument("--seed-existing", type=int, default=300) |
| p.add_argument("--outdir", default=BILI_RAW); p.set_defaults(func=cmd_rank) |
| p = sub.add_parser("tiles") |
| p.add_argument("--raw", required=True); p.add_argument("--out", default=TILES) |
| p.add_argument("--tile", type=int, default=320); p.add_argument("--frames", type=int, default=32) |
| p.add_argument("--segments", type=int, default=4); p.add_argument("--min-width", type=int, default=640) |
| p.add_argument("--seed", type=int, default=1234); p.set_defaults(func=cmd_tiles) |
| p = sub.add_parser("teacher") |
| p.add_argument("--tiles", default=TILES); p.add_argument("--labels", default=LABELS) |
| p.set_defaults(func=cmd_teacher) |
| p = sub.add_parser("label") |
| p.add_argument("--max-neg-ratio", type=float, default=2.0); p.add_argument("--mid-ratio", type=float, default=0.25) |
| p.add_argument("--seed", type=int, default=2026); p.set_defaults(func=cmd_label) |
| p = sub.add_parser("train") |
| p.add_argument("--out", required=True); p.add_argument("--epochs", type=int, default=8) |
| p.add_argument("--batch", type=int, default=256); p.add_argument("--lr", type=float, default=2e-3) |
| p.add_argument("--workers", type=int, default=8); p.add_argument("--hard-weight", type=float, default=2.0) |
| p.add_argument("--warmup", type=int, default=2); p.add_argument("--augment", type=int, default=0) |
| p.add_argument("--seed", type=int, default=2026); p.add_argument("--resume", default="") |
| p.set_defaults(func=cmd_train) |
| p = sub.add_parser("eval"); p.add_argument("--ckpt", required=True); p.set_defaults(func=cmd_eval) |
| p = sub.add_parser("compare"); p.add_argument("--ckpts", required=True); p.set_defaults(func=cmd_compare) |
| p = sub.add_parser("pack") |
| p.add_argument("--onnx", required=True); p.add_argument("--out", required=True) |
| p.add_argument("--log", required=True); p.add_argument("--eval", required=True) |
| p.add_argument("--params", default="2602017"); p.add_argument("--epochs", default="8") |
| p.add_argument("--amd-ms", default="0.56"); p.add_argument("--cpu-ms", default="0.317") |
| p.set_defaults(func=cmd_pack) |
| args = ap.parse_args(); args.func(args) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
