SwinVQColor/inference.py at main · navaneet625/SwinVQColor · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
import os
import torch
from torchvision.utils import save_image
from config import Config
from models.colorizer import Colorizer
from data.dataloader import get_loader
import utils
from tqdm import tqdm
from collections import OrderedDict
import warnings

warnings.filterwarnings("ignore", category=UserWarning)

@torch.no_grad()
def inference():
    cfg = Config
    cfg.device = (
        "cuda" if torch.cuda.is_available() else
        "mps" if torch.backends.mps.is_available() else
        "cpu"
    )

    # === Initialize model ===
    model = Colorizer(cfg).to(cfg.device)

    # === Load checkpoint safely ===
    ckpt_path = os.path.join(cfg.checkpoints_dir, "latest.pth")
    print(f"🔍 Loading checkpoint from: {ckpt_path}")
    assert os.path.exists(ckpt_path), f"❌ No checkpoint found at {ckpt_path}"

    ckpt = torch.load(ckpt_path, map_location=cfg.device)

    # If checkpoint was saved during training with optimizer/scaler
    if "model" in ckpt:
        ckpt = ckpt["model"]

    # Handle multi-GPU checkpoints (remove "module." prefixes)
    new_state_dict = OrderedDict()
    for k, v in ckpt.items():
        new_key = k.replace("module.", "")
        new_state_dict[new_key] = v

    missing, unexpected = model.load_state_dict(new_state_dict, strict=False)
    print(f"✅ Model loaded successfully.")
    if missing:
        print(f"⚠️ Missing keys: {len(missing)} → {missing[:5]} ...")
    if unexpected:
        print(f"⚠️ Unexpected keys: {len(unexpected)} → {unexpected[:5]} ...")

    model.eval()

    # === Data Loader ===
    val_loader = get_loader(
        root=cfg.data_root,
        split="val",
        bs=cfg.batch_size,
        image_size=cfg.img_size,
        max_samples=8000
    )

    # === Output Directory ===
    save_dir = os.path.join(cfg.results_dir, "inference")
    os.makedirs(save_dir, exist_ok=True)

    img_counter = 0

    # === Run Inference ===
    for idx, (L, ab) in enumerate(tqdm(val_loader, desc="🎨 Generating colorized images")):
        L = L.to(cfg.device)
        ab = ab.to(cfg.device)

        ab_pred, _ = model(L)
        rgb_pred = utils.lab2rgb_torch(L, ab_pred)
        rgb_true = utils.lab2rgb_torch(L, ab)

        # Save each image in the batch individually
        for b in range(rgb_pred.size(0)):
            gt_img = rgb_true[b].unsqueeze(0)
            pred_img = rgb_pred[b].unsqueeze(0)

            save_image(pred_img, os.path.join(save_dir, f"colorized_{img_counter:04d}.png"))
            save_image(gt_img, os.path.join(save_dir, f"gt_{img_counter:04d}.png"))
            save_image(torch.cat([gt_img, pred_img], dim=3),
                       os.path.join(save_dir, f"compare_{img_counter:04d}.png"))

            img_counter += 1

        if idx < 1:
            print(f"🖼️ Saved batch {idx} → images {img_counter - rgb_pred.size(0)} to {img_counter - 1}")

    print(f"\n✨ All images saved successfully to: {save_dir}\n")


if __name__ == "__main__":
    inference()