pablo-moreira/puzzle-piece-sides-classifier

Puzzle Piece Sides Classifier

A multi-task image classification model for identifying the four sides of a jigsaw puzzle piece and its complete side-pattern configuration.

The model is based on DINOv2 Base and uses two classification heads:

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

This model is part of the PuzzleMap project.

Model Details

Property	Value
Model version	20.0.0
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:

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:

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:

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:

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:

0 = TOP
1 = RIGHT
2 = BOTTOM
3 = LEFT

Each position has three classes:

0 = SMOOTH
1 = OUTER
2 = INNER

The head architecture is:

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:

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

The pattern head architecture is:

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:

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:

SSOI
SSIO
OSSI
...

Pattern Encoding

For example:

SOIS

means:

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

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

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:

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:

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:

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.

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:

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

Run the model:

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

The model returns two sets of logits:

outputs.sides_logits
outputs.pattern_logits

Their shapes are:

sides_logits:
(batch_size, 4, 3)

pattern_logits:
(batch_size, 60)

For a single image:

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:

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

print(sides)

Example:

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

The order is always:

[
    TOP,
    RIGHT,
    BOTTOM,
    LEFT
]

You can also obtain the probabilities:

sides_probabilities = outputs.sides_probabilities()[0]

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

For each side:

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

The class order is:

0 → SMOOTH
1 → OUTER
2 → INNER

---

Decoding Pattern Predictions

Obtain the predicted pattern ID:

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

Decode it using the processor:

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

print(pattern)

Example:

SOIS

The pattern probability distribution can be obtained with:

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.

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:

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.

sides = processor.convert_pattern_to_sides("SOIS")

print(sides)

Result:

[
    "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:

outputs.sides_logits
outputs.pattern_logits

The output object also provides convenience methods:

outputs.sides_probabilities()
outputs.pattern_probabilities()

outputs.predicted_sides_labels()
outputs.predicted_pattern_labels()

For example:

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:

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:

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

and:

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

---

Relation to PuzzleMap

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

The broader project uses computer vision and machine learning to:

1. detect puzzle pieces;
2. classify puzzle-piece properties;
3. determine the side configuration of each piece;
4. estimate piece orientation;
5. identify similar pieces;
6. 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:

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.