IOAI Training: Learn AI by solving problems

The problem in brief

You receive a Vision Transformer (ViT) already trained to recognise 15 animals, from beetles to zebras, and it used to be almost perfect. The new test images (1 000 of them) break it. The training set is tiny on purpose: 15 clean photos covering only 5 of the 15 classes. You may not use any other pretrained weights or external data. The statement drops two hints: the ViT reads each 224 × 224 image as a 14 × 14 grid of patches, and treating it as a black box will not be enough.

The data are the official Kaggle files. The test labels are private and we did not submit anything. Because 15 clean training images cannot be split into a meaningful validation set, we score ourselves on self-generated test images built to mimic the test (see below), and we say clearly how far that proxy can be trusted.

How you are scored

Plain accuracy over the 1 000 test images:

\text{accuracy} = \frac{1}{1000}\sum_{i} \mathbf{1}[\hat y_i = y_i].

What to aim for. The Evaluation page gives the organisers' baseline at 0.370 (the ViT applied to the raw test images) and their reference solution at 0.914. On the public leaderboard the top entry, named "Solution 2.csv" and apparently an organiser's file, has 0.930, and the best contestant has 0.914.

First look at the data

The ViT classifies the 15 training images perfectly. So the model is fine, and the test images must be different. Looking at them explains everything:

Test images, and what is left after keeping only patches with a smooth border

Each test image is a patchwork on the 16-pixel grid. A part of the grid, typically a few full rows in the upper half, is a genuine photo of the animal, left in place. Every other square has been replaced by a 16 × 16 patch cut from some other animal picture, often zebra stripes, panda black-and-white or tiger fur. The ViT looks at all 196 patches, and the pasted ones outvote the real ones: on the raw test images it predicts "Panda" for 364 images out of 1 000 and "Tiger" for 123, while the elephant class gets none (red bars below).

Predicted class counts

Averaging our genuine-patch detector (described below) over all test images shows where the real photo usually sits: rows 3 to 6 are genuine in about 70 % of the images, rows 0 and 1 in about 25 %, and rows 8 to 13 almost never.

A first baseline

Applying the ViT to the raw image is the organisers' baseline (0.370 on the leaderboard). We cannot reproduce that number without test labels, but our self-generated test images give the same picture: 0.553 for the raw image against 1.000 on clean photos.

The key idea

Two observations unlock the problem.

1. Genuine patches continue into their neighbours; pasted patches do not. In a real photo, the pixels on both sides of a patch border are almost equal, because the border is an arbitrary line through a smooth image. A pasted patch has nothing to do with its neighbour, so the colours jump. For the border between patch $A$ (left) and patch $B$ (right), we extrapolate each side linearly from its last two pixel columns and compare with the other side:

c_{AB} = \tfrac12\left(\,\big|2A_{-1} - A_{-2} - B_{0}\big| + \big|2B_{0} - B_{1} - A_{-1}\big|\,\right),

averaged over the 16 pixels and 3 colour channels of the border ( $A_{-1}$ is the last column of $A$ , $B_0$ the first column of $B$ ). Same thing for vertical borders. Inside a clean training photo the median cost is about 2. Between unrelated patches it is typically 50 to 150. A patch with at least one border below 15 is declared genuine. This one rule recovers the photo region in the examples above with very few false alarms (uniform patches, like a black bar, occasionally match each other).

2. A ViT does not need all 196 patches. This is the part of the hint about not treating the model as a black box. A ViT turns each patch into a token, adds a position embedding that tells the model where the patch was, and then runs self-attention over the set of tokens. Self-attention works on a set of any size. So we can simply delete the tokens of the pasted patches and run the transformer on the CLS token plus the genuine tokens, each keeping its own position embedding. The model sees a partial image with the real geometry, and the distractors are not merely hidden: they are absent.

Why is deleting better than painting the pasted squares grey? A grey square is still a token, and a strange one: the model never saw large grey holes in training. On our self-generated images, grey filling dropped accuracy to about 0.2, below the raw image. Deleting the tokens has no such side effect.

The full solution

Border costs for the 14 × 13 vertical and 13 × 14 horizontal patch borders (pure NumPy, a few milliseconds per image).
Keep the patches with at least one border cost below 15.
ViT on the kept tokens only:

tok = emb.patch_embeddings(x) + pos[:, 1:]       # 196 patch tokens with their positions
tok = tok[:, keep]                               # drop the pasted patches
h = torch.cat([emb.cls_token + pos[:, :1], tok], 1)
for layer in model.vit.encoder.layer:
    h = layer(h)
logits = model.classifier(model.vit.layernorm(h)[:, 0])

Predict the arg-max class. No training at all. The 15 training images are not used to fit anything, which is just as well, since 10 classes have no example.

How we validated. We built 300 self-generated test images. Each takes one clean training photo, keeps it inside a realistic genuine-region mask (masks copied from the structure we detect on real test images, holes filled), and fills every other square with a random patch from the other 14 training photos. This copies the test's construction, but with only 15 source photos and 5 classes, and with masks that do not always fall on the animal. It is a harsher test than the real one: even with the true mask, the ViT gets only 0.69 there.

Results

Accuracy on our 300 self-generated test images. The official test labels are private.

approach	self-generated accuracy
ViT on the raw image (organisers' baseline idea)	0.553
pasted squares painted grey, true mask	about 0.20
largest connected group of smooth borders only	0.48 (first, 150-image version)
raster-order runs of genuine patches, holes filled	0.72 (150-image version)
genuine-patch detector, tokens deleted (final)	0.710
tokens deleted with the true mask (oracle)	0.690

The detector even edges out the "true" mask. That is not a bug. The true mask contains flat background squares of the source photo (sky, white walls), which carry no information about the animal, and the detector tends to drop them.

On the real test images, the absolute accuracy is out of reach without labels, but one check is available. With the raw image the predicted classes are wildly unbalanced (364 Pandas, 0 Elephants). With the genuine patches only, every class receives between 54 and 84 images (blue bars), close to the 66.7 per class that a balanced test set would give. This is consistent with the organisers' reference at 0.914, but it does not prove it. If the test set is indeed balanced, the counts imply that at least about 5 % of the images are misclassified.

Going further: the author's approach

The author of the task, Megazhan, kindly told us how the reference solution scores higher. Instead of a hand-written rule on borders, it trains a segmentation model that decides, for every patch, whether it is genuine or pasted, and then runs the ViT on the patches it keeps, as above.

The training data for that detector costs nothing to build:

take the clean training photos and turn them into patchworks yourself, by pasting random squares over them: the label of every square (genuine or pasted) is then known by construction;
use patches cut from the test images as the pasted material, so that the distractors in this home-made training set look exactly like the real ones.

A learned detector combines far more cues than border continuity (texture, colour statistics, the position in the grid, the context of the neighbours), so it also handles the cases where our rule fails: uniform squares that happen to match each other, genuine squares on the edge of the animal with no smooth border, isolated genuine squares. To make it work, generate patchworks whose genuine regions look like the test ones (mostly whole rows in the upper half), and keep a few test images aside to check its masks by eye.

Thanks to Megazhan for this tip.

What to remember

When a strong model suddenly fails, look at the inputs before touching the model. Two minutes of looking at the test images showed the patchwork; no amount of fine-tuning on 15 images would have found it.
Natural images are smooth across arbitrary lines. Border continuity is a cheap, training-free test for "do these two pieces belong together", and it also underlies jigsaw solvers.
Know your architecture. A ViT is a transformer over a set of patch tokens with position embeddings, so it can run on any subset of patches. Removing tokens is cleaner than masking pixels, which creates inputs the model never saw.
When you cannot validate on the real test, build a test-like set and be explicit about its biases. Here it was harsher than the real test, so it ranks methods but underestimates the final score.
Predicted class counts are a free sanity check on an unlabelled test set: a model that sends a third of the images to one class is broken, whatever its confidence.

"""AICC Round 11, "Guessing Animals Again": worked solution.

Each test image is a 14 x 14 grid of 16-px patches. Only part of the grid comes from the real photo;
the other patches were pasted from other animal pictures (lots of zebra and panda stripes), which
fools the provided ViT. Idea:
1. A genuine patch continues smoothly into its genuine neighbours; a pasted patch does not. We measure
   the colour jump across every patch border (after a linear extrapolation of the last two pixel
   rows/columns) and keep the patches that have at least one smooth border.
2. A ViT is a set model: we feed ONLY the kept patch tokens (each with its own position embedding)
   plus the CLS token. The pasted patches are simply never seen.

Evaluation: the official test labels are private and the 15 training images are clean, so we
build self-generated test images: one training photo inside a realistic genuine-region mask (masks
taken from the structure detected on the test images), all other patches pasted at random from the
other training photos.

Usage: python solution.py [data_dir]   (data_dir holds train/, test/, model/, train.csv, test.csv)
"""
import os
import sys

os.environ.setdefault("OMP_NUM_THREADS", "2")
import numpy as np
import pandas as pd
import torch
from PIL import Image
from transformers import ViTForImageClassification

torch.set_num_threads(2)
DATA = sys.argv[1] if len(sys.argv) > 1 else "work"
P, N = 16, 14                    # patch size, patches per side
MODEL = ViTForImageClassification.from_pretrained(f"{DATA}/model").eval()


def load(path):
    return np.asarray(Image.open(path).convert("RGB"))


# ---------------------------------------------------------------- 1. which patches are genuine?
def border_costs(a):
    """Colour jump across each vertical border (H, 14x13) and horizontal border (V, 13x14).
    Each side is extrapolated linearly from its last two pixel lines and compared with the other side."""
    a = a.astype(float)
    H, V = np.zeros((N, N - 1)), np.zeros((N - 1, N))
    for j in range(N - 1):
        x = (j + 1) * P
        L1, L2, R1, R2 = a[:, x - 1], a[:, x - 2], a[:, x], a[:, x + 1]
        d = (np.abs(2 * L1 - L2 - R1) + np.abs(2 * R1 - R2 - L1)) / 2   # 224 x 3
        H[:, j] = d.reshape(N, P, 3).mean(axis=(1, 2))
    for i in range(N - 1):
        y = (i + 1) * P
        T1, T2, B1, B2 = a[y - 1], a[y - 2], a[y], a[y + 1]
        d = (np.abs(2 * T1 - T2 - B1) + np.abs(2 * B1 - B2 - T1)) / 2
        V[i, :] = d.reshape(N, P, 3).mean(axis=(1, 2))
    return H, V


def genuine(a, thr=15.0, need=1):
    """Patches with at least `need` smooth borders (cost < thr)."""
    H, V = border_costs(a)
    g = np.zeros((N, N))
    g[:, :-1] += H < thr
    g[:, 1:] += H < thr
    g[:-1] += V < thr
    g[1:] += V < thr
    return (g >= need).reshape(-1)


# ---------------------------------------------------------------- 2. ViT on a subset of tokens
@torch.no_grad()
def vit_logits(a, keep=None):
    x = (torch.tensor(a).permute(2, 0, 1).float() / 255 - 0.5) / 0.5
    emb = MODEL.vit.embeddings
    pos = emb.position_embeddings                               # 1 x 197 x 768
    tok = emb.patch_embeddings(x[None]) + pos[:, 1:]            # 1 x 196 x 768
    if keep is not None and keep.any():
        tok = tok[:, torch.from_numpy(keep)]                    # drop the pasted patches
    h = torch.cat([emb.cls_token + pos[:, :1], tok], 1)
    for layer in MODEL.vit.encoder.layer:
        out = layer(h)
        h = out[0] if isinstance(out, tuple) else out
    return MODEL.classifier(MODEL.vit.layernorm(h)[:, 0])[0].numpy()


def predict(a):
    return int(vit_logits(a, genuine(a)).argmax())


# ---------------------------------------------------------------- 3. self-generated validation
def to_patches(a):
    return a.reshape(N, P, N, P, 3).transpose(0, 2, 1, 3, 4).reshape(N * N, P, P, 3)


def from_patches(p):
    return p.reshape(N, N, P, P, 3).transpose(0, 2, 1, 3, 4).reshape(N * P, N * P, 3)


def span_fill(g, neg=0.5, min_len=8):
    """Genuine regions are runs in raster order: fill holes with repeated max-subarray search."""
    w, keep = np.where(g, 1.0, -neg), np.zeros(len(g), bool)
    for _ in range(3):
        best, cur, st = (0, 0, 0), 0.0, 0
        for i in range(len(g)):
            if keep[i]:
                cur, st = 0.0, i + 1
                continue
            if cur <= 0:
                cur, st = 0.0, i
            cur += w[i]
            if cur > best[0]:
                best = (cur, st, i + 1)
        if best[0] < min_len:
            break
        keep[best[1]:best[2]] = True
    return keep


def self_generated_eval(train_imgs, y, test_paths, n=300, seed=0):
    rng = np.random.default_rng(seed)
    masks = [span_fill(genuine(load(p), need=2)) for p in rng.choice(test_paths, 400, replace=False)]
    masks = [m for m in masks if m.sum() >= 20]
    pts = [to_patches(a) for a in train_imgs]
    hits = {"full image": 0, "true mask (oracle)": 0, "detected patches": 0}
    for k in range(n):
        main, keep = k % len(train_imgs), masks[rng.integers(len(masks))]
        others = [i for i in range(len(train_imgs)) if i != main]
        img = from_patches(np.array([pts[main][q] if keep[q] else
                                     pts[rng.choice(others)][rng.integers(N * N)] for q in range(N * N)]))
        hits["full image"] += vit_logits(img).argmax() == y[main]
        hits["true mask (oracle)"] += vit_logits(img, keep).argmax() == y[main]
        hits["detected patches"] += predict(img) == y[main]
    return {k: v / n for k, v in hits.items()}


if __name__ == "__main__":
    tr = pd.read_csv(f"{DATA}/train.csv")
    te = pd.read_csv(f"{DATA}/test.csv")
    train_imgs = [load(f"{DATA}/train/{i}.png") for i in tr.image_id]
    test_paths = [f"{DATA}/test/{i}.png" for i in te[te.columns[0]]]
    print("clean training images, full ViT:",
          np.mean([vit_logits(a).argmax() == t for a, t in zip(train_imgs, tr.label)]))
    for k, v in self_generated_eval(train_imgs, tr.label.values, test_paths).items():
        print(f"self-generated test images, {k}: {v:.3f}")
    preds = [predict(load(p)) for p in test_paths]
    print("predicted class counts on test:", np.bincount(preds, minlength=15))
    sub = pd.DataFrame({te.columns[0]: te[te.columns[0]], "label": preds})
    sub.to_csv("submission.csv", index=False)
    print("wrote submission.csv (not submitted)")