IOAI Training: Learn AI by solving problems

The problem in brief

Each $112\times112$ RGB image shows one filled shape on a white background: a circle, a triangle (not necessarily equilateral) or a rectangle. Each shape has a random position, size, rotation and colour. You get 1 000 training images, of which only the first 200 are labelled, and must classify 300 test images. Pretrained models are forbidden, and so is labelling images by eye. This is Problem 4 of the 2026 USA AI Olympiad Round 2, Day 2. The Hugging Face copy of the data also contains 300 labelled validation images, which is what we report on, because the test labels are private.

One caveat about the data files: train_labels.csv actually contains labels for all 1 000 training images, although the statement says images 200 to 999 are unlabelled. We follow the rules: our models only ever see the first 200 labels. We used the other 800 labels for a single diagnostic, measuring how accurate our pseudo-labels were.

How you are scored

The metric is macro F1 over the 3 classes,

F_1^{\text{macro}} = \frac{1}{3}\sum_{c} \frac{2 P_c R_c}{P_c + R_c},

and the mark is relative to the best contestant: $100\,\frac{X - X_\text{baseline}}{X_\text{best} - X_\text{baseline}}$ , with a baseline of 0 (no prediction). The best contestant's score is not published, so there is no fixed target. The ceiling is $F_1 = 1$ , and since this is a relative mark, anything below the best contestant loses points.

First look at the data

Every image has exactly two colours: white, and the shape's colour (no anti-aliasing, no noise).
Colour, size and position are pure nuisance variables; only the outline carries the class.
The classes are roughly balanced (362 / 327 / 311 among the 1 000 training images, 111 / 97 / 92 in validation).

So the first thing to do is throw the colour away: a pixel is "ink" if it is not white. That leaves a binary mask, and the problem becomes classical geometry.

A first baseline

A small CNN trained from scratch on the raw RGB images of the 200 labelled examples, without augmentation, is unstable. Its validation macro F1 is 0.955, 0.778 and 0.339 for three random seeds. With 200 examples, the network can latch onto colour or position, which mean nothing here, and its training can fall apart.

The key idea

Remove what does not matter, then build in what should not matter.

Ink mask: one binary channel. Colour is gone by construction.
Invariance by augmentation: during training, every mask is randomly rotated (any angle), rescaled (0.8 to 1.2) and shifted. The network then sees endless variations of the same 200 shapes and has to learn the shape itself, not its pose.
Semi-supervision by self-training: the trained network labels the 800 unlabelled images, the confident labels (probability above 0.95) are added to the training set, and the network is retrained.

There is also a model-free check, which shows why the task is easy once colour is gone. Seen from its centroid, the boundary of a shape traces a radial signature: a flat line for a circle, 3 peaks for a triangle, 4 for a rectangle. Two numbers already separate the classes:

\rho = \frac{\text{area}}{\pi\, d_{\max}^2} \quad (\text{1 for a circle, about } 0.64 \text{ for a square, under } 0.42 \text{ for a triangle}), \qquad \frac{\operatorname{std}(d)}{\operatorname{mean}(d)},

where $d$ is the distance from the centroid to the boundary pixels. A square inscribed in a circle of radius $d_{\max}$ has area $2d_{\max}^2$ , so $\rho = 2/\pi \approx 0.64$ . An equilateral triangle has $\rho = \frac{3\sqrt3}{4\pi} \approx 0.41$ , and less if it is elongated.

Ink masks, radial signatures and two geometric features

The full solution

Masks: pixel_sum < 3 * 250.
CNN: four blocks of [conv 3x3, batch norm, ReLU, max-pool] with 16, 32, 64 and 64 channels, then global average pooling and a linear layer with 3 outputs. About 60 k parameters, trained from scratch with AdamW (learning rate $3\times10^{-3}$ ) for 1 500 steps of 64 augmented masks, using the 200 labelled images.
Self-training: 794 of the 800 unlabelled images get a confident pseudo-label; all 794 turned out correct when checked against the file's hidden labels. Retrain for 2 000 steps on 994 images and predict the test set.
Independent check: 12 geometric features (the two above, the minimum of the radial signature, the eigenvalue ratio of the shape's covariance, and 7 Fourier coefficients of the signature) with logistic regression, also fitted on the 200 labels only.

Results

Macro F1 on the 300 labelled validation images (test labels are private).

approach	validation macro F1
no prediction (official baseline)	0
CNN on raw RGB, 200 labels, no augmentation (3 seeds)	0.339 / 0.778 / 0.955
CNN on ink masks, 200 labels, no augmentation	0.986
geometric features + logistic regression, 200 labels	1.000
CNN on masks + augmentation, 200 labels	1.000
CNN on masks + augmentation + self-training (final)	1.000

The final CNN and the geometric model agree on all 300 test images. Two completely different methods making identical predictions is good evidence that the test predictions are correct too.

What to remember

Strip nuisance variables at the input (colour here). The cheapest invariance is the one you build into the data representation.
With few labels, augmentation that encodes the true symmetries (rotation, scale, translation) is worth more than any architecture choice.
Self-training with a confidence threshold turns unlabelled data into extra labelled data, at almost no risk when the base model is already accurate.
A tiny hand-crafted geometric model is a powerful sanity check: when two unrelated methods agree on every test sample, you can trust the predictions.
Read the data files critically: one of them contained more labels than the rules allow you to use.

"""USAAIO 2026 Round 2, Problem 4: circle / triangle / rectangle from 112x112 RGB images,
200 labelled training images, no pretrained models.

Key step: colour is noise, so turn every image into a binary ink mask (pixel != white).
Model A (final): small CNN trained from scratch on the 200 labelled masks with random
rotation / scale / shift augmentation, then self-training on the 800 unlabelled images.
Model B (independent check): 12 hand-made geometric features + logistic regression.
Evaluated on the 300 labelled validation images of the dataset (never used for fitting).
  run from the folder holding Shapes/ (HF usaaio-official/2026_USAAIO_Round2):  python3 solution.py
"""
import numpy as np, pandas as pd, torch, torch.nn as nn, torch.nn.functional as F
from PIL import Image
from scipy import ndimage
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import f1_score
torch.set_num_threads(2); torch.manual_seed(0)
DEV = 'mps' if torch.backends.mps.is_available() else 'cpu'
N_LABELLED = 200                          # the statement: only train_0 .. train_199 have labels


def load(split, n):
    return np.stack([np.array(Image.open(f'Shapes/{split}_images/{split}_{i}.png').convert('RGB'))
                     for i in range(n)])


def ink(X):
    return X.astype(int).sum(-1) < 3 * 250                         # True where the shape is


def geometric_features(M):
    out = []
    for m in M:
        m = ndimage.binary_fill_holes(m); A = m.sum(); ys, xs = np.nonzero(m); cy, cx = ys.mean(), xs.mean()
        ey, ex = np.nonzero(m & ~ndimage.binary_erosion(m))         # boundary pixels
        d = np.hypot(ey - cy, ex - cx); ang = np.arctan2(ey - cy, ex - cx)
        sig = np.zeros(72); np.maximum.at(sig, ((ang + np.pi) / (2 * np.pi) * 72).astype(int) % 72, d)
        sig /= sig.max() + 1e-9                                     # radial signature (corners = peaks)
        spec = np.abs(np.fft.rfft(sig - sig.mean()))[1:8]; spec /= spec.sum() + 1e-9
        ev = np.linalg.eigvalsh(np.cov(np.vstack([ys, xs])))
        out.append(np.r_[A / (np.pi * d.max() ** 2), A / (np.pi * d.mean() ** 2), d.std() / d.mean(),
                         sig.min(), ev[0] / ev[1], spec])
    return np.array(out)


def make_cnn():
    blk = lambda i, o: nn.Sequential(nn.Conv2d(i, o, 3, padding=1), nn.BatchNorm2d(o), nn.ReLU(), nn.MaxPool2d(2))
    return nn.Sequential(blk(1, 16), blk(16, 32), blk(32, 64), blk(64, 64),
                         nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(64, 3)).to(DEV)


def augment(x):
    n = x.shape[0]; a = torch.rand(n) * 2 * np.pi; s = 0.8 + 0.4 * torch.rand(n); t = (torch.rand(n, 2) - 0.5) * 0.4
    th = torch.stack([torch.stack([torch.cos(a) / s, -torch.sin(a) / s, t[:, 0]], 1),
                      torch.stack([torch.sin(a) / s, torch.cos(a) / s, t[:, 1]], 1)], 1)
    return F.grid_sample(x, F.affine_grid(th.to(x.device), x.shape, align_corners=False), align_corners=False)


def train_cnn(X, y, steps=1500):
    net = make_cnn(); opt = torch.optim.AdamW(net.parameters(), 3e-3, weight_decay=1e-4)
    X, y = X.to(DEV), y.to(DEV)
    for _ in range(steps):
        b = torch.randint(0, len(y), (64,))
        loss = F.cross_entropy(net(augment(X[b])), y[b]); opt.zero_grad(); loss.backward(); opt.step()
    net.eval(); return net


@torch.no_grad()
def proba(net, X):
    return torch.softmax(net(X.to(DEV)), -1).cpu()


if __name__ == '__main__':
    tr = pd.read_csv('Shapes/train_labels.csv'); va = pd.read_csv('Shapes/val_labels.csv')
    sub = pd.read_csv('Shapes/sample_submission.csv')
    Mtr, Mva, Mte = ink(load('train', 1000)), ink(load('val', 300)), ink(load('test', len(sub)))
    y = torch.tensor(tr.label.values[:N_LABELLED]); yv = va.label.values      # labels >= 200 never used
    T = lambda M: torch.tensor(M, dtype=torch.float32)[:, None]

    # Model B: geometry
    lr = LogisticRegression(C=10, max_iter=5000).fit(geometric_features(Mtr[:N_LABELLED]), y.numpy())
    pb_val = lr.predict(geometric_features(Mva)); pb_test = lr.predict(geometric_features(Mte))
    print('B geometric features + logistic regression: val macro F1 %.4f' % f1_score(yv, pb_val, average='macro'))

    # Model A: CNN on masks, then one round of self-training on confident unlabelled images
    net = train_cnn(T(Mtr[:N_LABELLED]), y)
    print('A0 CNN, 200 labels + augmentation:        val macro F1 %.4f'
          % f1_score(yv, proba(net, T(Mva)).argmax(1).numpy(), average='macro'))
    P = proba(net, T(Mtr[N_LABELLED:])); conf, pl = P.max(1); keep = torch.where(conf > 0.95)[0]
    Xst = torch.cat([T(Mtr[:N_LABELLED]), T(Mtr[N_LABELLED:])[keep]]); yst = torch.cat([y, pl[keep]])
    print(f'   pseudo-labelled {len(keep)} of {len(P)} unlabelled images')
    net2 = train_cnn(Xst, yst, 2000)
    pa_val = proba(net2, T(Mva)).argmax(1).numpy(); pa_test = proba(net2, T(Mte)).argmax(1).numpy()
    print('A  CNN after self-training:               val macro F1 %.4f' % f1_score(yv, pa_val, average='macro'))
    print('test agreement between A and B: %.4f' % (pa_test == pb_test).mean())
    sub['label'] = pa_test; sub.to_csv('submission.csv', index=False)

Classifying Geometric Shapes from RGB Images