← Building the verification databases
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.
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:
numpy | The 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. |
duckdb | A 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. |
pandas | A 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. |
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.
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.
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.
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.