COCO: Super-resolution

This notebook continues the data preprocessing work, but specifically focused on preparing the data for super-resolution.

try:
    with open("data/coco_color_fps.json") as f:
        color_fps = json.load(f)["fps_color"]
except FileNotFoundError:
    fps = [str(fi) for fi in Path("data/train2017").glob("**/*.jpg")]
    color_fps = []
    with multiprocessing.Pool() as p:
        iter_ = zip(p.imap(black_and_white, fps), fps)
        for b_w, fp in tqdm(iter_, total=len(fps)):
            if not b_w:
                color_fps.append(fp)
    with open("data/coco_color_fps.json", "wt") as f:
        color_fps = json.dump({"fps_color": fps_color}, f)

---

source

coco_2017_trn

 coco_2017_trn (fps=None, n=None, remove_bw=True)

ds = coco_2017_trn(fps_color)

---

source

crop_to_box

 crop_to_box (img:<module'PIL.Image'from'/opt/hostedtoolcache/Python/3.10.
              14/x64/lib/python3.10/site-packages/PIL/Image.py'>)

---

source

preprocess_ddpm

 preprocess_ddpm (examples, pipe, extra_blur=False)

rows = ds["train"][:6]
fig, axes = plt.subplots(4, 5, figsize=(10, 8))
for ax in axes.flatten():
    ax.set_xticks([])
    ax.set_yticks([])
trn = preprocess_ddpm(rows, pipe=trn_preprocess_super_rez)
tst = preprocess_ddpm(rows, pipe=tst_preprocess_super_rez)
for im_trn_hi, im_trn_lo, im_test, im_org, ax_col in zip(
    trn["image_high_rez"],
    trn["image_low_rez"],
    tst["image_high_rez"],
    rows["image"],
    axes.T,
):
    for ax, im in zip(ax_col, (im_trn_lo, im_trn_hi, im_test)):
        ax.imshow(denorm(im).permute(1, 2, 0))
    ax_col[3].imshow(im_org)
axes[0, 0].set(title="Train (Low Rez)")
axes[1, 0].set(title="Train (Hi Rez)")
axes[2, 0].set(title="Test")
axes[3, 0].set(title="Original")
fig.tight_layout()

---

source

get_coco_dataset

 get_coco_dataset (fac, trn, tst, fp='data/train2017', bs=512, n=None,
                   columns=['image_low_rez', 'image_high_rez'])

dls = get_coco_dataset_super_rez(n=100)

CPU times: user 320 ms, sys: 35.5 ms, total: 355 ms
Wall time: 173 ms

We also want to do colorization

---

source

preprocess_colorization

 preprocess_colorization (examples, pipe)

dls = get_coco_dataset_colorization(n=100)

CPU times: user 319 ms, sys: 25 ms, total: 344 ms
Wall time: 156 ms

xb, yb = dls.peek()

denorm(xb).max()

tensor(0.9961)

show_images(denorm(xb[:6, ...]), imsize=(1.6))

show_images(denorm(yb[:6, ...]), imsize=(1.6))