pablo-moreira/puzzle-piece-classifier
Puzzle Piece Classifier
A computer-vision model for identifying the position and orientation of an individual jigsaw puzzle piece within a set of candidate puzzle-piece crops.
The model is part of the PuzzleMap project and is designed to operate after puzzle-piece detection.
Unlike a conventional image classifier, the model does not classify a piece into a fixed semantic category. Instead, it performs two related tasks:
- Crop selection — identifies which candidate crop corresponds to the input puzzle piece.
- Angle prediction — determines the orientation of the input piece relative to the selected crop.
The model uses DINOv2 Base as its visual backbone and adds custom neural heads for crop ranking and angle estimation.
Model Details
Intended Use
The model is intended to answer the following question:
Given an isolated puzzle piece and several candidate positions from a puzzle image, which candidate corresponds to the piece, and in which orientation should the piece be placed?
A typical PuzzleMap pipeline is:
Puzzle Image
│
▼
Puzzle Piece Detection
│
▼
Bounding Boxes
│
▼
Individual Piece Crop
│
▼
Puzzle Piece Classifier
│
├─────────────────────┐
│ │
▼ ▼
Crop Selection Angle Prediction
│ │
└──────────┬──────────┘
▼
Piece Localization
+ Orientation
│
▼
Puzzle AssemblyThe detector provides the initial piece localization, while this model determines which position in the puzzle corresponds to the extracted piece.
Architecture
The model is based on facebook/dinov2-base.
DINOv2 produces:
CLS token
+
Patch tokensThe model uses the DINOv2 patch tokens as the primary representation for both crop matching and angle prediction.
Unlike version 19, which calculated cosine similarity between corresponding patch representations, version 20 uses Multi-Head Cross-Attention to allow the patches from the rotated puzzle piece to attend to the patches from each candidate crop.
The architecture can be summarized as:
Input Piece
│
┌───────────┴───────────┐
│ │
▼ │
Create 4 Rotations │
│ │
┌─────────┼─────────┬─────────┐ │
│ │ │ │ │
▼ ▼ ▼ ▼ │
0° 90° 180° 270° │
│ │ │ │ │
└─────────┴─────────┴─────────┘ │
│ │
▼ │
DINOv2 Backbone │
│ │
▼ │
Patch Tokens │
│ │
▼ │
Patch Projection │
│ │
▼ │
[B, 4, 256, D] │
│
│
Candidate Crops ────────────────────────────┘
│
▼
DINOv2 Backbone
│
▼
Patch Tokens
│
▼
Patch Projection
│
▼
[B, Nc, 256, D]
│
│
└──────────────────┐
│
▼
Pair every rotation
with every crop
│
▼
Multi-Head Cross
Attention
│
┌──────────┴──────────┐
│ │
Query = Piece Key = Crop
Value = Crop
│ │
└──────────┬──────────┘
│
▼
Residual +
LayerNorm
│
▼
Attention Pooling
│
▼
[B, Nc, 4, D]
│
▼
Crop Head
│
▼
[B, Nc, 4]
│
┌────────┴────────┐
│ │
▼ ▼
Max over 4 Winner Crop
│ │
▼ ▼
Crop Scores Angle Logits
[B, Nc] [B, 4]The important architectural change in version 20 is that each candidate crop is evaluated jointly with each of the four rotations of the input piece.
For every candidate, the model therefore produces four scores:
Candidate 0:
rotation 0°
rotation 90°
rotation 180°
rotation 270°
Candidate 1:
rotation 0°
rotation 90°
rotation 180°
rotation 270°
...This produces an intermediate tensor:
[B, Nc, 4]where:
B = batch size
Nc = number of candidate crops
4 = possible rotationsModel Components
DINOv2 Backbone
The backbone is loaded from:
facebook/dinov2-baseThe implementation uses AutoModel so that the backbone can be loaded from its Hugging Face configuration.
The hidden dimension is obtained dynamically from:
self.backbone.config.hidden_sizeThe model uses the DINOv2 patch tokens rather than the CLS token.
This allows the custom heads to be constructed according to the backbone configuration.
For the standard DINOv2 Base configuration, the output contains:
1 CLS token
+
256 patch tokensThe CLS token is removed:
patch_tokens = output[:, 1:]resulting in:
[B, 256, hidden_size]These patch representations are then passed through the shared patch_proj layer.
Patch Projection
The model applies the same learned projection to the patch tokens from both the puzzle piece and candidate crops.
The projection is:
Linear(hidden_size → hidden_size)
GELU
Linear(hidden_size → hidden_size)implemented as:
self.patch_proj = nn.Sequential(
nn.Linear(hidden_size, hidden_size),
nn.GELU(),
nn.Linear(hidden_size, hidden_size)
)The projection is shared between:
Piece patches
Crop patchesAfter projection, the model normalizes the patch representations:
F.normalize(patch_tokens, dim=-1)The resulting representations are used by the cross-attention layer.
Crop Ranking
The model evaluates every candidate crop against all four rotations of the input piece.
For each pair:
piece rotation × candidate cropthe model applies Multi-Head Cross-Attention.
The attention uses:
Query = rotated piece patch tokens
Key = candidate crop patch tokens
Value = candidate crop patch tokensThe attention layer is:
self.crop_attention = nn.MultiheadAttention(
embed_dim=hidden_size,
num_heads=8,
dropout=0.1,
batch_first=True
)This produces an attended representation of the piece conditioned on the candidate crop.
A residual connection followed by LayerNorm is then applied:
attention_output = self.attention_norm(
piece_tokens + attention_output
)The resulting representation is pooled using learned attention and passed through the crop head.
For every candidate crop, the model produces four scores:
[B, Nc, 4]representing:
candidate 0 → [score 0°, score 90°, score 180°, score 270°]
candidate 1 → [score 0°, score 90°, score 180°, score 270°]
...Cross Attention
The main architectural innovation introduced in version 20 is the use of Multi-Head Cross-Attention between the puzzle piece and each candidate crop.
For each candidate and each possible rotation:
Rotated Piece Patches
│
│ Query
▼
┌───────────────┐
│Cross-Attention│
└───────────────┘
▲
│
│ Key / Value
│
Crop PatchesThe query consists of the patch tokens from the rotated puzzle piece.
The key and value consist of the patch tokens from the candidate crop.
This allows the model to learn interactions between regions of the piece and regions of the candidate rather than relying only on independent embeddings or patch-to-patch similarity.
The attention configuration is:
Embedding dimension: hidden_size
Number of heads: 8
Dropout: 0.1
Batch first: TrueThe output has the same sequence length and hidden dimension as the piece representation:
[B × Nc × 4, 256, hidden_size]A residual connection is applied:
attention_output = self.attention_norm(
piece_tokens + attention_output
)This provides the angle/crop head with a representation that incorporates information from both the piece and the candidate crop.
Crop Head
Version 20 uses a single neural head to score every:
candidate crop × rotationcombination.
The head is:
Linear(hidden_size → 512)
GELU
Dropout(0.2)
Linear(512 → 256)
GELU
Dropout(0.2)
Linear(256 → 1)The input is:
[B, Nc, 4, hidden_size]and the output is:
[B, Nc, 4]The final dimension represents the four possible rotations:
0 → 0°
1 → 90°
2 → 180°
3 → 270°The implementation is:
all_angle_logits = self.crop_head(
pooled
).squeeze(-1)The name all_angle_logits reflects the fact that the model initially calculates an angle-aware score for every candidate crop.
Crop Selection
For each candidate crop, the model has four scores:
[B, Nc, 4]The crop score is defined as the best score among its four possible rotations:
crop_scores = all_angle_logits.max(dim=-1).valuesThis produces:
[B, Nc]The best candidate is then selected using:
winner = crop_scores.argmax(dim=1)Therefore:
all_angle_logits
│
▼
max over rotations
│
▼
crop_scores [B, Nc]
│
▼
argmax
│
▼
winner [B]The use of the maximum means that a candidate crop is considered strong if at least one of its four possible orientations produces a high compatibility score.
Angle Prediction
Angle prediction is directly coupled to crop scoring.
The model first computes:
[B, Nc, 4]scores for all candidate/rotation combinations.
After selecting the best candidate:
winner = crop_scores.argmax(dim=1)the four scores corresponding to that candidate are extracted:
angle_logits = all_angle_logits[
torch.arange(B, device=all_angle_logits.device),
winner
]The result is:
[B, 4]representing the compatibility of the selected crop with each possible rotation.
The predicted angle is:
angle_logits.argmax(dim=1)with the mapping:
0 → 0°
1 → 90°
2 → 180°
3 → 270°This means that the model does not have an independent angle classifier. Instead, angle prediction is derived from the same crop/rotation compatibility scores used for candidate selection.
Model Output
The model returns a custom PuzzlePieceClassifierOutput:
@dataclass
class PuzzlePieceClassifierOutput(ModelOutput):
crop_scores: torch.Tensor = None
angle_logits: torch.Tensor = NoneThe two outputs are:
crop_scores
Shape:
[B, Nc]Each value represents the best compatibility score for a candidate crop across its four possible rotations.
It is calculated as:
crop_scores = all_angle_logits.max(dim=-1).valuesThe predicted crop is:
crop_scores.argmax(dim=1)angle_logits
Shape:
[B, 4]These are the four rotation scores for the selected candidate crop.
The predicted angle class is:
angle_logits.argmax(dim=1)The complete internal flow is therefore:
[B, Nc, 4]
│
├── max(rotation) ──→ crop_scores [B, Nc]
│
└── select winner ──→ angle_logits [B, 4]Input
The model receives two tensors.
Piece
piece_pixel_valuesShape:
[B, C, H, W]This represents the isolated puzzle piece.
Candidate Crops
crop_pixel_valuesShape:
[B, Nc, C, H, W]where Nc is the number of candidate crops.
Inside forward(), the candidates are flattened before being processed by DINOv2:
[B, Nc, C, H, W]
│
▼
[B × Nc, C, H, W]The candidate representations are then reshaped back to:
[B, Nc, hidden_size]after feature extraction.
Dataset
The training dataset is composed of both:
Synthetic puzzle pieces
Real puzzle piecesReal pieces are selected from annotations containing:
valid = true
bbox != None
puzzle != NoneThe original image is cropped using the bounding box and rotated according to its annotation before being prepared for the model.
Synthetic pieces are generated from known puzzle layouts and masks.
Candidate Generation
Each training sample contains:
1 positive candidate
+
hard negatives
+
random negativesThe positive candidate is always included.
Hard negatives are taken from:
piece.top3_hard_negativesThe remaining candidates are sampled randomly from the other puzzle positions.
The candidates are then shuffled, and the label is the position of the positive candidate in the shuffled list.
The implementation supports a configurable number of candidates, with six candidates used by default in PieceSample.
Conceptually:
Candidate Set
│
├── Positive
├── Hard Negative
├── Hard Negative
├── Hard Negative
├── Random Negative
└── Random NegativeThis makes the crop-ranking task substantially harder than simply comparing a piece against arbitrary unrelated crops.
Hard Negative Mining
Hard negatives are updated dynamically during training.
After computing the crop scores, the model selects the top four candidates:
batch_topk = crop_scores.topk(4, dim=1)These candidates are used to update the hard-negative list for future samples.
The correct candidate is removed from the hard-negative list, and up to three difficult candidates are retained.
The process can therefore be represented as:
Current model
│
▼
Score candidates
│
▼
Top-K candidates
│
▼
Remove positive
│
▼
Store difficult negatives
│
▼
Future training samplesThis creates an online hard-negative mining process.
Data Augmentation
Training samples can receive image augmentations including:
- brightness;
- contrast;
- color;
- Gaussian noise;
- Gaussian blur;
- combinations of the above.
The combined random augmentation applies each transformation independently with a probability of 50%.
The training samples use random_example as their image transformer.
Train / Validation Split
The available puzzle pieces are split using:
train_test_split(
puzzles_pieces,
test_size=0.2,
random_state=86,
stratify=puzzles_pieces_labels
)The stratification label is the puzzle name.
Therefore, the split maintains the distribution of puzzle sources between training and validation.
The resulting datasets are:
80% training
20% validationDataset Class
The PyTorch dataset is:
PuzzlePieceClassifierDatasetEach item contains:
{
"crop_pixel_values": ...,
"piece_pixel_values": ...,
"crop_labels": ...,
"angle_labels": ...,
"indices": ...,
"candidate_indices": ...
}The implementation converts the selected crop and angle into separate labels.
The angle labels correspond directly to the four angle classes.
Loss
The model uses two Cross Entropy losses:
Crop Loss
Angle LossThe total loss is:
Loss =
Crop Loss
+
angle_weight × Angle LossThe two tasks contribute equally to the training objective.
Metrics
The training process evaluates the two tasks independently as well as jointly.
Crop Accuracy
Percentage of samples where the correct candidate is ranked first.
crop_scores.argmax(dim=1)Angle Accuracy
Percentage of samples where the correct angle is predicted.
angle_logits.argmax(dim=1)Joint Accuracy
A sample is considered correct only when:
crop prediction == correct crop
AND
angle prediction == correct angleThe implementation calculates all three metrics.
Crop Ranking Metrics
Because crop selection is a ranking problem, additional metrics are used.
Top-2 Accuracy
The correct crop must appear among the two highest-scoring candidates.
Top-3 Accuracy
The correct crop must appear among the three highest-scoring candidates.
Mean Reciprocal Rank
MRR evaluates how highly the correct candidate is ranked.
For example:
Rank 1 → 1.000
Rank 2 → 0.500
Rank 3 → 0.333
Rank 4 → 0.250The implementation calculates MRR directly from the ordering of crop_scores.
These metrics are particularly useful because a correct candidate ranked second or third can still be useful to a downstream puzzle-search system.
Training Strategy
Training is performed in two phases.
Phase 1 — Train Heads
The DINOv2 backbone is frozen:
backbone_requires_grad=FalseThe custom heads are trained using a relatively aggressive learning rate:
Learning rate: 1e-3The phase is configured for up to:
150 epochswith early stopping patience:
4 epochsThe best model is selected according to:
validation lossand saved together with the processor.
The purpose of this phase is to train the custom heads while preserving the pretrained DINOv2 representation.
Phase 2 — Fine-Tune DINOv2
The best checkpoint from Phase 1 is loaded and the DINOv2 backbone is unfrozen:
backbone_requires_grad=TrueA substantially smaller learning rate is used:
Learning rate: 5e-6The phase runs for up to:
40 epochswith early stopping patience:
4 epochsThe training batch size is reduced to:
2to accommodate the computational cost of running DINOv2 multiple times per sample.
The complete training strategy is therefore:
Phase 1
────────────────────────────
DINOv2 Frozen
Custom Heads Trainable
LR 1e-3
Max Epochs 150
Patience 4
│
▼
Best Checkpoint
│
▼
Phase 2
────────────────────────────
DINOv2 Trainable
Custom Heads Trainable
LR 5e-6
Max Epochs 40
Patience 4Loading the Model
Because the model implements a custom Hugging Face PreTrainedModel and defines an auto_map, it should be loaded with remote/custom code enabled when loading the repository.
from transformers import AutoModel, AutoImageProcessor
MODEL_ID = "pablo-moreira/puzzle-piece-classifier"
processor = AutoImageProcessor.from_pretrained(MODEL_ID)
model = AutoModel.from_pretrained(
MODEL_ID,
trust_remote_code=True,
)
model.eval()The custom model configuration maps:
AutoConfig
→ PuzzlePieceClassifierConfig
AutoModel
→ PuzzlePieceClassifierand the processor use default Hugging Face AutoImageProcessor implementation.
Inference
The model requires:
1 puzzle piece
+
N candidate cropsThe images are processed using the standard Hugging Face AutoImageProcessor.
Conceptually:
piece_inputs = processor(
images=[piece_image],
return_tensors="pt",
)
crop_inputs = processor(
images=crop_images,
return_tensors="pt",
)The tensors need to be arranged as:
piece_pixel_values:
[1, C, H, W]
crop_pixel_values:
[1, N, C, H, W]Then:
with torch.no_grad():
outputs = model(
piece_pixel_values=piece_pixel_values,
crop_pixel_values=crop_pixel_values,
)Internally, the model evaluates all four rotations against every candidate crop:
Candidate Crops
│
┌──────────────┼──────────────┐
▼ ▼ ▼
Crop 0 Crop 1 Crop N
│ │ │
┌────┼────┐ ┌────┼────┐ ┌────┼────┐
▼ ▼ ▼ ▼ ▼ ▼ ▼ ▼ ▼
0° 90° 180° 0° 90° 180° ... 270°
│
▼
Cross Attention
│
▼
Rotation Scores
│
▼
Best Crop + AngleReading the Prediction
Best Crop
crop_index = outputs.crop_scores.argmax(dim=1).item()The value identifies the candidate crop with the highest compatibility score.
For example:
crop_index = 4means that the fifth candidate crop is considered the best match.
Predicted Angle
angle_index = outputs.angle_logits.argmax(dim=1).item()Convert the class to degrees:
angle = angle_index * 90The mapping is:
0 → 0°
1 → 90°
2 → 180°
3 → 270°Complete Inference Example
import torch
from transformers import (
AutoModel,
AutoImageProcessor,
)
from PIL import Image
MODEL_ID = "pablo-moreira/puzzle-piece-classifier"
# Load processor
processor = AutoImageProcessor.from_pretrained(MODEL_ID)
# Load model
model = AutoModel.from_pretrained(
MODEL_ID,
trust_remote_code=True,
)
model.eval()
# Load the isolated puzzle piece
piece_image = Image.open(
"piece.png"
).convert("RGB")
# Load candidate crops
crop_images = [
Image.open("crop_0.png").convert("RGB"),
Image.open("crop_1.png").convert("RGB"),
Image.open("crop_2.png").convert("RGB"),
Image.open("crop_3.png").convert("RGB"),
Image.open("crop_4.png").convert("RGB"),
Image.open("crop_5.png").convert("RGB"),
]
# Process piece
piece_pixel_values = processor(
images=[piece_image],
return_tensors="pt",
)["pixel_values"]
# Process candidates
crop_pixel_values = processor(
images=crop_images,
return_tensors="pt",
)["pixel_values"]
# Add batch dimension to candidate crops
crop_pixel_values = crop_pixel_values.unsqueeze(0)
# Inference
with torch.no_grad():
outputs = model(
piece_pixel_values=piece_pixel_values,
crop_pixel_values=crop_pixel_values,
)
# Best candidate
crop_index = outputs.crop_scores.argmax(
dim=1
).item()
# Best angle
angle_index = outputs.angle_logits.argmax(
dim=1
).item()
angle = angle_index * 90
print("Best crop:", crop_index)
print("Angle:", angle)The model will therefore return something conceptually equivalent to:
Best crop: 4
Angle: 90meaning:
Candidate crop 4
should be used for the piece,
with a rotation of 90 degrees.Training vs. Inference
There is an important difference between training and inference.
During training, the correct candidate can be passed as:
crop_labelsThe model then uses the correct crop for the angle branch.
During inference, crop_labels is omitted:
outputs = model(
crop_pixel_values=crop_pixel_values,
piece_pixel_values=piece_pixel_values,
)The model selects:
argmax(crop_scores)and uses that candidate for angle prediction.
Therefore, the production pipeline is:
Piece
│
▼
DINOv2
│
├─────────────── Candidate 0
├─────────────── Candidate 1
├─────────────── Candidate 2
├─────────────── ...
└─────────────── Candidate N
│
▼
Crop Scores
│
▼
Best Candidate
│
▼
Selected Candidate + Piece
│
▼
DINOv2 Patch Tokens
│
├── Piece rotated 0°
├── Piece rotated 90°
├── Piece rotated 180°
└── Piece rotated 270°
│
▼
Cosine Similarity
│
▼
Attention Pooling
│
▼
Angle PredictionOutput Shapes
For:
batch_size = 1
number_of_candidates = 6the model produces:
outputs.crop_scores.shapetorch.Size([1, 6])and:
outputs.angle_logits.shapetorch.Size([1, 4])The dimensions therefore represent:
crop_scores
[batch, candidates]
angle_logits
[batch, angles]Limitations
Rotation Classes
The model predicts only four discrete rotations:
0°
90°
180°
270°It does not directly predict arbitrary continuous angles.
Puzzle Domain
The model is specialized for jigsaw puzzles.
It is not intended as a general image matching model.
PuzzleMap Pipeline
The classifier is designed to operate after the PuzzleMap detection model.
The complete workflow can be represented as:
Puzzle Image
│
▼
Puzzle Piece Detection
│
▼
Bounding Boxes
│
┌───────────┴───────────┐
│ │
▼ ▼
Piece Crop Puzzle Candidates
│ │
└───────────┬───────────┘
▼
Puzzle Piece Classifier
│
┌──────┴──────┐
│ │
▼ ▼
Best Crop Rotation
│ │
└──────┬──────┘
▼
Piece Localization
+ Orientation
│
▼
Other Puzzle Analysis
│
▼
Puzzle AssemblyThe detector answers:
Where are the pieces?The classifier answers:
Which puzzle position does this piece belong to?
How should it be rotated?Together they form the localization and matching stages of PuzzleMap.
Version Comparison
The following table summarizes the main milestones in the development of the Puzzle Piece Classifier.
Version 21
Version 21 focuses on dataset restructuring and improved train/validation stratification, while maintaining the patch-level cross-attention architecture introduced in Version 20.
The dataset was significantly reduced and rebalanced, increasing the relative proportion of real-world puzzle pieces while retaining synthetic examples for additional variation. This change was intended to improve dataset quality and provide a more representative evaluation of model generalization.
The train/validation splitting strategy was also redesigned to reduce potential data leakage:
- Real pieces are now grouped by their complete puzzle coordinate (
puzzle.row.column), ensuring that the same physical piece cannot appear in both the training and validation sets. - Synthetic pieces are grouped by puzzle and row (
puzzle.row), keeping highly correlated pieces generated from the same source structure within the same split. - This replaces the previous strategy, which stratified pieces only by puzzle name.
Together, these changes produce a more reliable validation set and a more realistic assessment of the model's ability to generalize, particularly to real-world puzzle pieces.
Version 20
Version 20 introduces a new architecture based on patch-level cross-attention between the puzzle piece and candidate crops.
- Replaced the previous CLS-based crop representation with DINOv2 patch tokens.
- Added a shared
patch_projprojection for piece and crop patch tokens. - Added an 8-head
MultiheadAttentionlayer for piece-to-crop cross-attention. - The rotated piece acts as the Query.
- Candidate crop patches act as Key and Value.
- Added a residual connection followed by LayerNorm after cross-attention.
- Added learned attention pooling over the cross-attended patch representations.
- The crop head now produces four scores per candidate, one for each possible rotation.
- The final crop score is the maximum score across the four rotations.
- The angle prediction is obtained from the four scores belonging to the selected crop.
- Maintained the four discrete angle classes:
0°,90°,180°, and270°. - Switched to the default Hugging Face
AutoImageProcessor, eliminating the need for custom processor code. - The processor no longer requires
trust_remote_code=True.
Version 19
- Introduced DINOv2 patch-token similarity for angle prediction.
- Added learned attention pooling over the patch-level similarities.
- The crop branch uses the DINOv2 CLS representation with dedicated projections.
- Maintained four discrete angle classes: 0°, 90°, 180°, and 270°.
- Achieved 97.36% crop accuracy and 94.60% angle accuracy.
- Achieved 98.68% MRR and 100% Top-2/Top-3 crop accuracy.
- Represents the final architecture documented for the model.
Version 18
- Fixed the rotation handling for 90° and 270°.
- This was the major breakthrough in the angle prediction task.
- Angle accuracy increased from 48.55% to 96.62%.
- Joint accuracy increased from 47.10% to 95.89%.
- Crop accuracy reached 99.03%.
- Achieved the lowest validation loss among the reported versions: 0.2295.
Version 17
- Introduced the combined Crop Head + Angle Head architecture.
- Added prediction of the four possible piece rotations.
- Crop matching was already highly accurate.
- Angle prediction remained close to random/ambiguous behavior, reaching only 48.55% validation accuracy.
- The experiment was therefore abandoned.
The recorded experiment history is maintained in the training notebook.
Citation
If you use this model in your project, please reference:
Pablo Moreira.
Puzzle Piece Classifier.
PuzzleMap project.License
This repository is released under the terms specified by the repository license.
The underlying facebook/dinov2-base model and the PuzzleMap dataset are subject to their respective licenses and terms of use.
Users are responsible for verifying the licensing requirements of the underlying pretrained model, dataset, and other dependencies used in their applications.