Sides Head — independently classifies the four sides of the puzzle piece.
Pattern Head — classifies the complete four-side configuration as a single pattern.

This model is part of the PuzzleMap project.

Model Details

Property	Value
Architecture	DINOv2 Base + multi-task classification heads
Backbone	`facebook/dinov2-base`
Framework	PyTorch
Integration	Hugging Face Transformers
Number of sides	4
Side classes	3
Pattern classes	60 output classes
Input size	`224 × 224`
Color format	RGBA/RGB converted by the image processor
Training strategy	Two-stage training
Augmentation	Rotation, mirroring, brightness, contrast, color, noise and blur

Intended Use

The model is intended to classify individual jigsaw puzzle pieces after they have been detected and cropped from a puzzle image.

For each piece, the model predicts:

text

TOP
RIGHT
BOTTOM
LEFT

Each side can be one of:

SMOOTH — a flat puzzle border.
OUTER — an outward protruding tab.
INNER — an inward indentation.

For example:

text

TOP    = SMOOTH
RIGHT  = OUTER
BOTTOM = INNER
LEFT   = SMOOTH

The corresponding pattern is:

SOIS

where each character represents the first letter of its corresponding side class:

text

S = SMOOTH
O = OUTER
I = INNER

Therefore:

[SMOOTH, OUTER, INNER, SMOOTH] -> SOIS

Architecture

The model uses facebook/dinov2-base as its visual backbone.

The CLS token from DINOv2 is used as the image representation:

text

Input Image
     │
     ▼
DINOv2 Base
     │
     ▼
CLS Embedding
     │
     ├─────────────────────────────┐
     │                             │
     ▼                             ▼
Sides Head                    Pattern Head
     │                             │
     ▼                             ▼
4 × 3 logits                 60 logits
     │                             │
     ▼                             ▼
TOP/RIGHT/BOTTOM/LEFT          Pattern

Sides Head

The sides head receives the DINOv2 hidden representation and produces:

4 × 3 = 12 logits

The output is reshaped to:

(batch_size, 4, 3)

The four positions correspond to:

text

= TOP
= RIGHT
= BOTTOM
= LEFT

Each position has three classes:

text

= SMOOTH
= OUTER
= INNER

The head architecture is:

text

LayerNorm
Linear(hidden_size → 512)
GELU
Dropout(0.2)
Linear(512 → 12)

Pattern Head

The pattern head receives:

the DINOv2 CLS embedding;
the predicted side logits.

The side logits are detached before being concatenated with the DINOv2 representation:

python

pattern_input = torch.cat(
    [
        features,
        sides_logits.detach().flatten(start_dim=1)
    ],
    dim=1
)

The pattern head architecture is:

text

LayerNorm
Linear(hidden_size + 12 → 512)
GELU
Dropout(0.2)
Linear(512 → 60)

This creates a multi-task architecture where the model learns both the individual side characteristics and the global configuration of the piece.

Labels

Side Labels

The side classification uses three labels:

text

SMOOTH
OUTER
INNER

Pattern Labels

The pattern is represented by four characters.

Each character corresponds to one side:

TOP RIGHT BOTTOM LEFT

The complete set configured by the model is:

text

SSOI
SSIO
OSSI
...

Pattern Encoding

For example:

SOIS

means:

text

TOP    = SMOOTH
RIGHT  = OUTER
BOTTOM = INNER
LEFT   = SMOOTH

The conversion is performed by taking the first character of each side label.

python

def build_pattern_label(sides_labels):
    return "".join([side[0] for side in sides_labels])

Input Processing

The model expects a cropped puzzle piece.

During training, the original piece was processed approximately as follows:

text

Original puzzle image
        │
        ▼
Piece bounding box
        │
        ▼
20% expanded bounding box
        │
        ▼
Crop        
        │
        ▼
Resize to 224 × 224
        │
        ▼
DINOv2 image processor

The expanded crop was used to preserve contextual pixels around the puzzle piece.

The final image is resized using the greatest dimension while preserving the aspect ratio and padding the remaining area with transparency.

Training

Training was performed in two stages using the PuzzleMap dataset.

The dataset contains annotated jigsaw puzzle pieces used to train the model to classify the four sides of each piece and its complete side-pattern configuration.

Stage 1 — Classification Heads

The DINOv2 backbone was frozen.

Only the classification heads were trained.

Configuration:

text

Optimizer: AdamW
Learning rate: 1e-3
Batch size: 64
Maximum epochs: 30
Backbone: frozen
Pattern loss weight: 0.50
Early stopping patience: 4

The total loss is:

loss = sides_loss + pattern_weight × pattern_loss

For the released training configuration:

pattern_weight = 0.50

Both losses use cross entropy.

Stage 2 — DINOv2 Fine-Tuning

The complete model was then fine-tuned, including the DINOv2 backbone.

Configuration:

text

Optimizer: AdamW
Learning rate: 5e-6
Batch size: 16
Maximum epochs: 40
Backbone: trainable
Pattern loss weight: 0.50
Early stopping patience: 4

The best model was selected according to validation loss.

Inference

Installation

Install the required packages:

pip install torch torchvision transformers pillow

Loading the Model

The model can be loaded directly using AutoModel.

Because this is a custom Transformers architecture, trust_remote_code=True is required.

python

from transformers import AutoModel, AutoImageProcessor
from PIL import Image
import torch

MODEL_ID = "pablo-moreira/puzzle-piece-sides-classifier"

processor = AutoImageProcessor.from_pretrained(
    MODEL_ID,
    trust_remote_code=True
)

model = AutoModel.from_pretrained(
    MODEL_ID,
    trust_remote_code=True
)

model.eval()

Running Inference

Load a cropped puzzle piece:

image = Image.open("puzzle_piece.png").convert("RGBA")

Process the image:

python

inputs = processor(
    images=image,
    return_tensors="pt"
)

Run the model:

python

with torch.no_grad():
    outputs = model(**inputs)

The model returns two sets of logits:

python

outputs.sides_logits
outputs.pattern_logits

Their shapes are:

text

sides_logits:
(batch_size, 4, 3)

pattern_logits:
(batch_size, 60)

For a single image:

python

print(outputs.sides_logits.shape)
# torch.Size([1, 4, 3])

print(outputs.pattern_logits.shape)
# torch.Size([1, 60])

Decoding Side Predictions

The processor provides helper methods to convert numeric predictions back to human-readable labels.

First, obtain the predicted side IDs:

sides_ids = outputs.predicted_sides_labels()[0]

Decode them:

python

sides = processor.decode_sides_labels(
    [sides_ids.tolist()]
)[0]

print(sides)

Example:

['SMOOTH', 'OUTER', 'INNER', 'SMOOTH']

The order is always:

text

[
    TOP,
    RIGHT,
    BOTTOM,
    LEFT
]

You can also obtain the probabilities:

python

sides_probabilities = outputs.sides_probabilities()[0]

print(sides_probabilities.shape)
# torch.Size([4, 3])

For each side:

python

for side_index, probabilities in enumerate(sides_probabilities):
    print(
        side_index,
        probabilities.tolist()
    )

The class order is:

text

→ SMOOTH
→ OUTER
→ INNER

Decoding Pattern Predictions

Obtain the predicted pattern ID:

pattern_id = outputs.predicted_pattern_labels()[0].item()

Decode it using the processor:

python

pattern = processor.decode_pattern_labels(
    [pattern_id]
)[0]

print(pattern)

Example:

SOIS

The pattern probability distribution can be obtained with:

python

pattern_probabilities = outputs.pattern_probabilities()[0]

print(pattern_probabilities.shape)
# torch.Size([60])

Complete Inference Example

The following example shows how to run inference with the trained model using a sample puzzle-piece image from the PuzzleMap dataset.

The example downloads the image directly from Hugging Face, loads the trained model, processes the image, and displays the predicted side classification for each side of the puzzle piece.

python

from transformers import AutoModel, AutoImageProcessor
from PIL import Image
import torch
from io import BytesIO
import requests

MODEL_ID = "pablo-moreira/puzzle-piece-sides-classifier"

IMAGES = [
    "https://huggingface.co/datasets/pablo-moreira/puzzle-map/resolve/main/samples/puzzle-piece-sides-classifier/120_avengers_20260709_105353_d1ee901639e04ba7972a6505cb07d541_OOII.png",
    "https://huggingface.co/datasets/pablo-moreira/puzzle-map/resolve/main/samples/puzzle-piece-sides-classifier/120_avengers_20260709_110848_3b131fdfeb134505b994c6c7cab791ad_IIOO.png",
    "https://huggingface.co/datasets/pablo-moreira/puzzle-map/resolve/main/samples/puzzle-piece-sides-classifier/120_avengers_20260709_112115_5a0ca30002bf48dfb97b9184f038cf30_IISI.png",
    "https://huggingface.co/datasets/pablo-moreira/puzzle-map/resolve/main/samples/puzzle-piece-sides-classifier/camera_2c545dac37064e56aea6139df0951608_SOIS.png",
    "https://huggingface.co/datasets/pablo-moreira/puzzle-map/resolve/main/samples/puzzle-piece-sides-classifier/puzzle-focus_d378d8e8-20240108_211624.redimensionado_ISSO.png"
]


# Load processor
processor = AutoImageProcessor.from_pretrained(
    MODEL_ID,
    trust_remote_code=True
)

# Load model
model = AutoModel.from_pretrained(
    MODEL_ID,
    trust_remote_code=True
)

model.eval()

# Load image
response = requests.get(IMAGES[0])
response.raise_for_status()

image = Image.open(BytesIO(response.content)).convert("RGB")

# Prepare input
inputs = processor(
    images=image,
    return_tensors="pt"
)

# Inference
with torch.no_grad():
    outputs = model(**inputs)

# --------------------------------------------------
# Sides
# --------------------------------------------------

sides_ids = outputs.predicted_sides_labels()[0].tolist()

sides = processor.decode_sides_labels(
    [sides_ids]
)[0]

print("Sides:")
print(f"TOP:    {sides[0]}")
print(f"RIGHT:  {sides[1]}")
print(f"BOTTOM: {sides[2]}")
print(f"LEFT:   {sides[3]}")

# --------------------------------------------------
# Pattern
# --------------------------------------------------

pattern_id = outputs.predicted_pattern_labels()[0].item()

pattern = processor.decode_pattern_labels(
    [pattern_id]
)[0]

print()
print("Pattern:")
print(pattern)

Example output:

text

Sides:
TOP:    INNER
RIGHT:  SMOOTH
BOTTOM: SMOOTH
LEFT:   OUTER

Pattern:
ISSO

Converting a Pattern to Sides

The processor also provides a helper for converting a pattern string into the corresponding side labels.

python

sides = processor.convert_pattern_to_sides("SOIS")

print(sides)

Result:

python

[
    "SMOOTH",
    "OUTER",
    "INNER",
    "SMOOTH"
]

The order is:

TOP → RIGHT → BOTTOM → LEFT

This can be useful when only the pattern prediction is required.

Accessing Raw Logits

The model output is a PuzzlePieceSidesClassifierOutput.

It exposes:

python

outputs.sides_logits
outputs.pattern_logits

The output object also provides convenience methods:

python

outputs.sides_probabilities()
outputs.pattern_probabilities()

outputs.predicted_sides_labels()
outputs.predicted_pattern_labels()

For example:

python

sides_probabilities = outputs.sides_probabilities()
pattern_probabilities = outputs.pattern_probabilities()

sides_predictions = outputs.predicted_sides_labels()
pattern_predictions = outputs.predicted_pattern_labels()

Limitations

The model has several important limitations.

Cropping Quality

The model expects an image containing an individual puzzle piece.

Performance may degrade if:

the bounding box is inaccurate;
multiple pieces are present;
a significant part of the piece is missing;
the piece is heavily occluded.

Rotation

The training pipeline explicitly generates the four principal orientations:

text

0°
90°
180°
270°

plus a small random rotation perturbation.

Very large arbitrary rotations or unusual perspective distortions may not be represented adequately by the training distribution.

Pattern Classes

The model predicts a predefined set of pattern configurations.

It should not be assumed that an arbitrary four-side combination is represented by a unique class.

Model Files

The repository contains the Hugging Face model artifacts required to load the model.

The custom architecture is exposed through:

puzzle_piece_sides_classifier.py

and the custom processor through:

puzzle_piece_sides_classifier_processor.py

The model weights are stored using SafeTensors.

The configuration and processor configuration are saved using the standard Hugging Face save_pretrained() mechanism.

The repository can therefore be loaded using:

python

AutoModel.from_pretrained(
    "pablo-moreira/puzzle-piece-sides-classifier",
    trust_remote_code=True
)

and:

python

AutoImageProcessor.from_pretrained(
    "pablo-moreira/puzzle-piece-sides-classifier",
    trust_remote_code=True
)

Version Comparison

The following table compares the results from versions of the Puzzle Piece Sides and Pattern Classifier.

Metric	v11	v12	v13	v14	v15	v16	v17	v18	v19	v20	v21
Dataset	518	518	2072	2080	3648	4664	6824	7160	7784	7784	9472
Training examples	414	414	1657	1664	2918	3731	5459	5728	6227	6227	7577
Validation examples	104	104	415	416	730	933	1365	1432	1557	1557	1895
Pattern weight	0.35	0.50	0.5	0.5	0.5	0.5	0.5	0.5	0.5	0.5	0.5
Best epoch	14	11	9	9	8	6	6	16	10	14	8
Train loss	0.0419	0.1039	0.0134	0.0060	0.0041	0.0080	0.0136	0.0049	0.0011	0.0106	0.0235
Train sides loss	0.0178	0.0455	0.0031	0.0007	0.0006	0.0009	0.0034	0.0011	0.0001	0.0044	0.0089
Train pattern loss	0.0688	0.1168	0.0206	0.0106	0.0069	0.0142	0.0204	0.0076	0.0019	0.0124	0.0291
Train sides accuracy	0.9982	0.9886	0.9995	1.0000	1.0000	1.0000	0.9992	0.9996	1.0000	0.9991	0.9984
Train pattern accuracy	0.9952	0.9808	0.9982	1.0000	0.9993	0.9997	0.9965	0.9988	1.0000	0.9979	0.9933
Validation loss	0.3734	0.3613	0.0634	0.0625	0.0285	0.0246	0.0212	0.0148	0.0052	0.0038	0.0011
Validation sides loss	0.1471	0.1156	0.0304	0.0299	0.0092	0.0084	0.0085	0.0078	0.0022	0.0019	0.0005
Validation pattern loss	0.6465	0.4914	0.0661	0.0652	0.0387	0.0325	0.0253	0.0139	0.0059	0.0038	0.0012
Validation sides accuracy	0.9815	0.9630	0.9939	0.9976	0.9986	0.9989	0.9987	0.9983	0.9990	0.9992	0.9999
Validation pattern accuracy	0.8426	0.9167	0.9904	0.9928	0.9918	0.9979	0.9978	0.9951	0.9987	0.9987	1.0000
Validation sides F1											0.9998
Validation pattern F1											1.0000

Version 21

Dataset: #9472
Epoch 08/40
Added/updated pattern F1 metric
Added/updated sides F1 metric

Version 20

Dataset: #7784
Fixed duplicated pattern class 0000
Epoch 14/40

Version 19

Dataset: #7784
Epoch 10/40

Version 18

Dataset: #7160
Fixed rotation
Fixed resize
Transparent background handling
Epoch 16/40

Version 17

Correct DINO crop
Data augmentation
Dataset error fixed
Dataset: #6824
Epoch 06/40

Version 16

Correct DINO crop
Data augmentation
Dataset error fixed
Dataset: #4664
Epoch 06/40

Version 15

Correct DINO crop
Data augmentation
Dataset error fixed
Dataset: #3648
Epoch 08/40

Version 14

Correct DINO crop
Data augmentation
Dataset error fixed
Dataset: #2080
Epoch 09/40

Version 13

Correct DINO crop
Data augmentation
Epoch 09/40

Version 12

Pattern loss weight: 0.50
Epoch 11/40

Version 11

Pattern loss weight: 0.35
Epoch 14/40

Relation to PuzzleMap

This model is one component of the PuzzleMap computer-vision pipeline.

The broader project uses computer vision and machine learning to:

detect puzzle pieces;
classify puzzle-piece properties;
determine the side configuration of each piece;
estimate piece orientation;
identify similar pieces;
assist in assembling jigsaw puzzles.

The sides classifier provides structured information that can be used by subsequent puzzle-solving components.

Citation

If you use this model in your project, please reference the PuzzleMap project and this model repository:

text

Pablo Moreira.
Puzzle Piece Sides Classifier.
PuzzleMap project

License

This model is released under the terms specified by the repository license.

The underlying facebook/dinov2-base model is subject to its own license and terms of use.

Users are responsible for verifying the licensing requirements of the underlying datasets, pretrained models and other dependencies used in their applications.