| """Feature caching and dataset loading utilities.""" |
|
|
| import json |
| import os |
| import random |
| from typing import List, Tuple |
|
|
| import numpy as np |
| import torch |
| from PIL import Image |
| from torchvision.transforms import v2 |
|
|
| RESOLUTION = 640 |
| HOOK_BLOCKS = [2, 5, 8, 11] |
| N_PREFIX = 5 |
|
|
| COCO_CONTIG_TO_CAT = [ |
| 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, |
| 27, 28, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 48, 49, 50, |
| 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 67, 70, 72, 73, 74, 75, |
| 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90, |
| ] |
| COCO_CAT_TO_CONTIG = {cat: i for i, cat in enumerate(COCO_CONTIG_TO_CAT)} |
|
|
|
|
| def letterbox(image: Image.Image, res: int): |
| W0, H0 = image.size |
| scale = res / max(H0, W0) |
| new_w, new_h = int(round(W0 * scale)), int(round(H0 * scale)) |
| resized = image.resize((new_w, new_h), Image.BILINEAR) |
| canvas = Image.new("RGB", (res, res), (0, 0, 0)) |
| canvas.paste(resized, (0, 0)) |
| return canvas, scale |
|
|
|
|
| def load_coco_subset(split, n, coco_root, min_anns=0): |
| img_dir = os.path.join(coco_root, f"{split}2017") |
| ann_file = os.path.join(coco_root, "annotations", f"instances_{split}2017.json") |
| with open(ann_file) as f: |
| coco = json.load(f) |
| id_to_anns = {} |
| for a in coco["annotations"]: |
| if a["iscrowd"]: |
| continue |
| cat = a["category_id"] |
| if cat not in COCO_CAT_TO_CONTIG: |
| continue |
| id_to_anns.setdefault(a["image_id"], []).append(a) |
| id_to_info = {img["id"]: img for img in coco["images"]} |
| candidates = [(iid, anns) for iid, anns in id_to_anns.items() if len(anns) >= min_anns] |
| if split == "val": |
| candidates.sort(key=lambda x: -len(x[1])) |
| else: |
| random.seed(42) |
| random.shuffle(candidates) |
| items = [] |
| for iid, anns in candidates[:n]: |
| info = id_to_info[iid] |
| path = os.path.join(img_dir, info["file_name"]) |
| boxes, labels = [], [] |
| for a in anns: |
| x, y, w, h = a["bbox"] |
| if w < 1 or h < 1: |
| continue |
| boxes.append([x, y, x + w, y + h]) |
| labels.append(COCO_CAT_TO_CONTIG[a["category_id"]]) |
| if boxes: |
| items.append({"path": path, "boxes": boxes, "labels": labels, |
| "width": info["width"], "height": info["height"]}) |
| return items |
|
|
|
|
| def cache_features(backbone, items, cache_path, resolution=RESOLUTION): |
| """Cache backbone features for a list of image items.""" |
| if os.path.isfile(cache_path): |
| return |
|
|
| normalize = v2.Compose([ |
| v2.ToImage(), v2.ToDtype(torch.float32, scale=True), |
| v2.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), |
| ]) |
|
|
| hooks, intermediates = [], {} |
| for idx in HOOK_BLOCKS: |
| def _hook(block_idx): |
| def fn(mod, inp, out): |
| intermediates[block_idx] = (out[0] if isinstance(out, list) else out).detach() |
| return fn |
| hooks.append(backbone.blocks[idx].register_forward_hook(_hook(idx))) |
|
|
| cached = [] |
| for item in items: |
| img = Image.open(item["path"]).convert("RGB") |
| canvas, scale = letterbox(img, resolution) |
| x = normalize(canvas).unsqueeze(0).cuda() |
|
|
| intermediates.clear() |
| with torch.no_grad(): |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| out = backbone.forward_features(x) |
| patches = out["x_norm_patchtokens"].float() |
| B, N, D = patches.shape |
| h = w = int(N ** 0.5) |
| spatial = patches[0].permute(1, 0).reshape(D, h, w).half().cpu() |
| inter = [intermediates[idx][0].half().cpu() for idx in HOOK_BLOCKS] |
|
|
| boxes = torch.tensor(item["boxes"], dtype=torch.float32) * scale |
| labels = torch.tensor(item["labels"], dtype=torch.long) |
| cached.append({ |
| "spatial": spatial, "intermediates": inter, |
| "boxes": boxes, "labels": labels, |
| }) |
|
|
| for h in hooks: |
| h.remove() |
|
|
| os.makedirs(os.path.dirname(cache_path), exist_ok=True) |
| torch.save(cached, cache_path) |
|
|
