← Building the verification databases

build_notta_tta_comp_db.py, explained for beginners

The previous page covers how to run this script. This page is for readers who want to understand how it works and why it's written the way it is — no prior familiarity with this codebase assumed. Released in full, unedited, as DOI 1.19.1.

What each library is for

The script depends on six third-party libraries (see requirements.txt in the release). Here's what each one actually does in this code, not just its name:

numpyThe basic library for working with grids of numbers efficiently. An image is, underneath, just a grid of pixel brightness values — numpy lets the code do math on an entire 8×8 block (or a whole image) at once, instead of writing a slow loop over every individual pixel.
opencv-python-headless (imported as cv2)A widely used computer-vision library. "Headless" means it has no screen-display features built in, which is exactly right here since this script never needs to show an image on screen — it only needs to compute numbers from one. Used (inside heatmap_tool.py) for specific measurements like edge sharpness and noise level.
scikit-image (imported as parts of skimage)Another image-processing library, used here specifically for one texture-analysis technique (a "gray-level co-occurrence matrix," which looks at how often certain pairs of pixel brightnesses appear next to each other — a standard way to quantify how "textured" a patch looks).
pillow (imported as PIL)The standard Python library for opening image files (PNG, JPEG, etc.) and turning them into pixel data the other libraries can work with.
duckdbA database engine that runs locally, with no separate server to install or manage. This is where every computed number ultimately gets stored, in a form that can later be queried directly with SQL.
pandasA library for working with data arranged in rows and columns (like a spreadsheet, in memory). Used here purely as a convenient stepping stone: after computing tens of thousands of numbers for one image, they're assembled into a table first, then handed to duckdb in one batch write — much faster than inserting one row at a time.

The code, section by section

Finding which files belong together

Every image exists as three separate files — one ground-truth version, one reconstructed without test-time augmentation (notta), and one reconstructed with it (tta) — usually in three different folders, often with slightly different filename conventions per model. The script's first job is figuring out, for each image ID, whether all three files actually exist.

def build_lookup(model_dir):
    lookup = {}
    for path in glob.glob(os.path.join(model_dir, '*.png')):
        base = os.path.basename(path)
        m = IMAGE_ID_RE.match(base)
        if m:
            image_id = m.group(1)
        elif '_LR' in base:
            image_id = base.split('_LR')[0]
        else:
            image_id = os.path.splitext(base)[0]
        lookup[image_id] = path
    return lookup

This function lists a folder's contents once and builds a lookup table (image ID → file path) in memory. The alternative — searching the folder again for every single image — sounds harmless, but with thousands of files in the folder it turns into thousands of full folder scans, which gets dramatically slower the more images there are, especially over a network or mounted drive. Listing once and looking things up in memory afterward avoids that entirely.

The main per-image computation

This is the function that actually does the work for one image. It runs independently for every image, in parallel (more on that below).

def process_one_image(args):
    image_id, gt_path, notta_path, tta_path = args
    try:
        gt_u8, gt_f = load_gray(gt_path)
        notta_u8, notta_f = load_gray(notta_path)
        tta_u8, tta_f = load_gray(tta_path)
        if not (gt_u8.shape == notta_u8.shape == tta_u8.shape):
            return image_id, None, f'size mismatch ...'
        h, w = gt_u8.shape
        if h % BLOCK_SIZE or w % BLOCK_SIZE:
            return image_id, None, f'size {h}x{w} not a multiple of {BLOCK_SIZE}'

        n, grid_notta = compute_ssim_grid(gt_f, notta_f)
        _, grid_tta = compute_ssim_grid(gt_f, tta_f)
        comp_gt = compute_comp_grid(gt_u8, n)
        comp_notta = compute_comp_grid(notta_u8, n)
        comp_tta = compute_comp_grid(tta_u8, n)
        ...
        df = pd.DataFrame(data)
        img_info = (image_id, gt_path, notta_path, tta_path, w, h)
        return image_id, (img_info, df), None
    except Exception as e:
        return image_id, None, f'{type(e).__name__}: {e}'

Step by step: load all three images as grayscale pixel arrays (once as whole numbers for the comp_* features, once as decimal numbers for the SSIM math — load_gray returns both). Confirm the three images are actually the same size, and that the size divides evenly into 8×8 blocks with nothing left over. Then compute SSIM twice (ground truth vs. notta, and ground truth vs. tta) and the eight comp_* features three times (once each for the ground-truth, notta, and tta pixels) — all of it done for every 8×8 block in the image at once, because these functions are written using numpy's array operations rather than a per-block loop.

Notice the whole function is wrapped in a try/except: if anything goes wrong for one image (a corrupted file, an unexpected size), that image is skipped with an error message rather than crashing the entire run partway through thousands of images.

Setting up the database structure

def ensure_schema(con):
    con.execute("""
        CREATE TABLE IF NOT EXISTS images (
            image_id VARCHAR PRIMARY KEY,
            path_gt VARCHAR, path_notta VARCHAR, path_tta VARCHAR,
            width INTEGER, height INTEGER
        );
    """)
    ...
    con.execute(f"""
        CREATE TABLE IF NOT EXISTS blocks (
            image_id VARCHAR, row INTEGER, col INTEGER,
            sigma_gt DOUBLE, mu_gt DOUBLE,
            ...
            PRIMARY KEY (image_id, row, col)
        );
    """)

This creates the two database tables described on the previous page, if they don't already exist. IF NOT EXISTS matters here: it's what makes it safe to run this same function again every time the script restarts, without wiping out data from a previous, interrupted run.

Tying it together: resumability and parallelism

The main() function is where the resumability described on the previous page is actually implemented:

done = set(r[0] for r in con.execute('SELECT image_id FROM images').fetchall())
...
for image_id in all_ids:
    if image_id in done:
        continue
    ...
    tasks.append((image_id, ...))

if args.limit is not None:
    tasks = tasks[:args.limit]

Before doing any work, the script reads which image IDs are already recorded in the database and puts them in a set called done. The loop right after it then simply skips any image already in that set. This is the entire resumability mechanism — nothing more elaborate than "check what's already there, don't redo it." Because --limit is applied after this filtering, running the same command repeatedly with a small --limit processes a fresh batch of new images each time, rather than reprocessing the same ones.

for chunk_start in range(0, len(tasks), args.chunk_size):
    chunk = tasks[chunk_start:chunk_start + args.chunk_size]
    with ProcessPoolExecutor(max_workers=args.workers) as executor:
        results = list(executor.map(process_one_image, chunk))

    for image_id, payload, error in results:
        ...
        con.execute('INSERT INTO images VALUES (?, ?, ?, ?, ?, ?)', img_info)
        con.append('blocks', df)
        n_done += 1

Images are processed in small batches ("chunks," 20 images by default). Within a chunk, ProcessPoolExecutor hands the images out across several worker processes running at the same time — this is what lets the script use multiple CPU cores instead of processing one image after another. Only once an entire chunk's worth of images has finished does the script write their results into the database and print a progress line. This is also why interrupting a run mid-chunk loses that chunk's progress (nothing in it has been written yet) while already-completed chunks stay safely recorded — exactly what the demonstration video shows happening.