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)

---

coco_2017_trn

def coco_2017_trn(
    fps:NoneType=None, n:NoneType=None, remove_bw:bool=True
):

ds = coco_2017_trn(fps_color)

---

crop_to_box

def crop_to_box(
    img:Image
):

---

preprocess_ddpm

def preprocess_ddpm(
    examples, pipe, extra_blur:bool=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()

---

get_coco_dataset

def get_coco_dataset(
    fac, trn, tst, fp:str='data/train2017', bs:int=512, n:NoneType=None,
    columns:list=['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

---

preprocess_colorization

def 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))