pablo-moreira/puzzle-piece-rotation-regressor

Puzzle Piece Rotation Regressor

A computer-vision regression model for estimating the rotation angle of an individual jigsaw puzzle piece.

The model is based on DINOv2 Base and uses a regression head to predict the orientation of the puzzle piece as a normalized two-dimensional vector:

[sin(angle), cos(angle)]

The predicted vector is converted back into an angle in degrees during inference.

This model is part of the PuzzleMap project.

Model Details

Property	Value
Model version	2
Architecture	DINOv2 Base + regression head
Backbone	facebook/dinov2-base
Framework	PyTorch
Integration	Hugging Face Transformers
Output dimension	2
Output representation	[sin(angle), cos(angle)]
Angle range	-90° to 90°
Input size	224 × 224
Color format	RGBA/RGB converted by the image processor
Loss function	Smooth L1 Loss
Training strategy	Two-stage training
Augmentation	Brightness, contrast and autocontrast

---

Intended Use

The model is intended to estimate the orientation of an individual jigsaw puzzle piece after it has been detected and cropped from a puzzle image.

The model predicts the angle required to represent the orientation of the piece.

For example:

Input:
    Cropped puzzle piece

Prediction:
    27.4°

The model does not directly perform puzzle-piece detection. The input should already contain a single puzzle piece.

A typical PuzzleMap pipeline can therefore be represented as:

Puzzle Image
     │
     ▼
Piece Detection
     │
     ▼
Piece Bounding Box
     │
     ▼
Puzzle Piece Crop
     │
     ▼
Rotation Regressor
     │
     ▼
Estimated Rotation

---

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
     │
     ▼
Regression Head
     │
     ▼
Linear(hidden_size → 512)
     │
     ▼
GELU
     │
     ▼
Dropout(0.2)
     │
     ▼
Linear(512 → 2)
     │
     ▼
L2 Normalization
     │
     ▼
[sin(angle), cos(angle)]

Regression Head

The regression head receives the DINOv2 CLS embedding.

Its architecture is:

Linear(hidden_size → 512)
GELU
Dropout(0.2)
Linear(512 → 2)
L2 Normalization

The final normalization is applied using:

pred = F.normalize(pred, dim=-1)

This constrains the output to lie approximately on the unit circle.

The two output values represent:

output[0] = sin(angle)
output[1] = cos(angle)

---

Angle Representation

Instead of directly predicting an angle in degrees, the model predicts its sine and cosine components.

An angle is converted to the regression target using:

def angle_to_sincos(angle):
    angle = max(-90, min(90, angle))

    rad = np.deg2rad(angle)

    return [
        np.sin(rad),
        np.cos(rad)
    ]

For example:

Angle: 0°

sin(0°) = 0
cos(0°) = 1

Target:
[0.0, 1.0]

For:

Angle: 45°

the target is approximately:

[0.7071, 0.7071]

The inverse conversion is performed with:

def sincos_to_angle(pred):
    sin, cos = pred

    return np.degrees(
        np.arctan2(sin, cos)
    )

This representation avoids directly regressing a periodic angular quantity and provides a continuous representation suitable for neural-network regression.

---

Labels

The training labels are the annotated rotation angles of the puzzle pieces.

Only annotations satisfying both conditions were included:

valid == True
rotation != None

The angle is represented internally as:

[sin(angle), cos(angle)]

The processor supports angles in the range:

-90° ≤ angle ≤ 90°

Angles outside this range are clamped during target conversion:

angle = max(-90, min(90, angle))

---

Input Processing

The model expects an image containing an individual puzzle piece.

During training, the annotated bounding box was used to extract the piece:

Original puzzle image
        │
        ▼
Piece bounding box
        │
        ▼
Crop
        │
        ▼
Color / contrast augmentation
        │
        ▼
Resize to 224 × 224
        │
        ▼
DINOv2 image processor

The crop itself is obtained directly from the annotated bounding box:

image.crop(bbox)

The resulting image is converted to RGBA and resized while preserving its aspect ratio.

The greatest dimension is scaled to 224 pixels.

The remaining area is placed on a transparent 224 × 224 canvas:

Original crop
     │
     ▼
Preserve aspect ratio
     │
     ▼
Greatest dimension = 224
     │
     ▼
Transparent 224 × 224 canvas

This avoids geometric distortion caused by independently resizing width and height.

---

Data Augmentation

The training dataset applies lightweight photometric augmentation:

transforms.ColorJitter(
    brightness=0.1,
    contrast=0.1
)

transforms.RandomAutocontrast(
    p=0.2
)

These transformations are intended to improve robustness to differences in lighting, exposure and image contrast.

No geometric augmentation is applied directly to the training image because the annotated rotation is the regression target.

---

Training

Training was performed using the PuzzleMap dataset.

The training dataset was expanded in this model version, increasing the number of available examples from approximately 108 to 973 annotated samples.

The data was divided into training and validation sets using:

train_test_split(
    pieces,
    test_size=0.2,
    random_state=42
)

Therefore:

80% → Training
20% → Validation

The split uses a fixed random seed of 42.

Loss Function

The model uses Smooth L1 Loss:

torch.nn.SmoothL1Loss(beta=0.1)

The loss is calculated between the predicted normalized vector and the target [sin(angle), cos(angle)] vector:

loss = SmoothL1Loss(
    predicted_sin_cos,
    target_sin_cos
)

This means the training objective is not directly measured in degrees.

The angular error is instead evaluated separately by converting both vectors back into angles.

---

Two-Stage Training

Training was performed in two stages.

Stage 1 — Regression Head

During the first stage, the DINOv2 backbone was frozen.

Only the regression head was trained.

Configuration:

Optimizer: AdamW
Learning rate: 1e-4
Batch size: 32
Maximum epochs: 30
Backbone: frozen
Early stopping patience: 4
Loss: Smooth L1 Loss

The best model was selected according to validation loss.

For model version 2, training stopped after:

Epoch 4/30

with:

Training loss:   0.0533
Validation loss: 0.0431

Stage 2 — DINOv2 Fine-Tuning

The best model from the first stage was then loaded and the DINOv2 backbone was unfrozen.

The complete model was fine-tuned using a substantially smaller learning rate:

Optimizer: AdamW
Learning rate: 5e-6
Batch size: 32
Maximum epochs: 40
Backbone: trainable
Early stopping patience: 4
Loss: Smooth L1 Loss

For model version 2, the best validation result was obtained at:

Epoch 25/40

with:

Training loss:   0.0015
Validation loss: 0.0142

The best checkpoint according to validation loss was saved as the released model.

---

Evaluation

The training notebook evaluates the model using the absolute difference between the real and predicted angles:

error = abs(real_angle - pred_angle)

A prediction is considered successful when:

error < 10

Therefore, the evaluation criterion used during the visual validation step is:

Absolute angular error < 10°

The notebook also generates a visual comparison between the real and predicted orientations:

┌─────────────────┐    ┌─────────────────┐
│      REAL       │    │      PRED       │
│                 │    │                 │
│   rotated by    │    │   rotated by    │
│   real angle    │    │   predicted     │
│                 │    │   angle         │
└─────────────────┘    └─────────────────┘

ERROR: 4.32°
RESULTADO: True

This visual inspection is useful for identifying cases where the numerical error alone may not adequately describe the quality of the prediction.

---

Inference

Installation

Install the required packages:

pip install torch torchvision transformers pillow

Loading the Model

The model can be loaded directly from Hugging Face using AutoModel.

Because this repository contains 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-rotation-regressor"

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 an image containing an individual 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["pixel_values"]
    )

The model returns a tensor containing:

[sin(angle), cos(angle)]

For a batch of images, the output shape is:

(batch_size, 2)

For example:

torch.Size([1, 2])

The predicted angle can be obtained using the processor:

angle = processor.sincos_to_angle(
    outputs[0].cpu().numpy()
)

print(angle)

Example:

27.43

meaning an estimated rotation of approximately:

27.43°

---

Complete Inference Example

The following example loads the model and processor, downloads the same sample puzzle-piece images used by the Puzzle Piece Sides Classifier, and performs inference using the first image from the list.

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

MODEL_ID = "pablo-moreira/puzzle-piece-rotation-regressor"

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("RGBA")

# --------------------------------------------------
# Prepare input
# --------------------------------------------------

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

# --------------------------------------------------
# Inference
# --------------------------------------------------

with torch.no_grad():
    prediction = model(
        inputs["pixel_values"]
    )

# --------------------------------------------------
# Convert sin/cos to angle
# --------------------------------------------------

predicted_vector = prediction[0].cpu().numpy()

predicted_angle = processor.sincos_to_angle(
    predicted_vector
)

print(
    f"Predicted rotation: {predicted_angle:.2f}°"
)

Example output:

Predicted rotation: -0.84°

Using the Processor Directly

The custom processor provides two utility methods for converting between angle and sine/cosine representations.

Angle → Sin/Cos

target = processor.angle_to_sincos(45)

print(target)

Result:

[0.70710678, 0.70710678]

Sin/Cos → Angle

angle = processor.sincos_to_angle(
    [0.70710678, 0.70710678]
)

print(angle)

Result:

45.0

These utilities use the same representation used during training.

---

Model Output

The model directly returns the normalized two-dimensional prediction:

prediction = model(
    inputs["pixel_values"]
)

Its output represents:

prediction[:, 0] → sin(angle)
prediction[:, 1] → cos(angle)

For example:

Prediction:
[0.4617, 0.8870]

which corresponds approximately to:

angle = atan2(0.4617, 0.8870)
      ≈ 27.5°

The model normalizes its output internally, so the vector is constrained to have approximately unit magnitude.

---

Limitations

The model has several important limitations.

Piece Detection and Cropping

The model does not detect puzzle pieces.

It expects an input image containing an individual piece.

Performance can degrade when:

• the bounding box is inaccurate;
• significant parts of the piece are cropped;
• multiple pieces are present;
• the piece is heavily occluded;
• the image contains excessive background.

Rotation Range

The processor clamps target angles to:

-90° to 90°

Therefore, this model should not be interpreted as a general-purpose 0°–360° rotation estimator.

It is specifically designed for the rotation representation used by the PuzzleMap dataset and pipeline.

Symmetry

Some puzzle pieces may have visual structures that make their orientation ambiguous.

For example, pieces with approximately symmetric shapes or textures can produce visually similar appearances under different rotations.

The model cannot resolve information that is not visually distinguishable in the input.

Dataset Distribution

Performance depends on how closely an input image resembles the images used during training.

Changes in:

• puzzle types;
• camera characteristics;
• lighting;
• image quality;
• piece size;
• background;
• segmentation/cropping quality;

may affect the prediction accuracy.

---

Model Files

The repository contains the Hugging Face artifacts required to load the custom architecture.

The custom model implementation is provided through:

puzzle_piece_rotation_regressor.py

The custom image processor is provided through:

puzzle_piece_rotation_regressor_processor.py

The model configuration defines:

model_type:
puzzle_piece_rotation_regressor

and maps the custom classes through Hugging Face's auto_map mechanism.

The model weights are stored using SafeTensors.

The model can therefore be loaded using:

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

and the processor using:

AutoImageProcessor.from_pretrained(
    "pablo-moreira/puzzle-piece-rotation-regressor",
    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. classify the four sides of each piece;
4. estimate the piece rotation;
5. identify similar pieces;
6. assist in assembling jigsaw puzzles.

The rotation regressor provides the estimated orientation required by subsequent puzzle-solving components.

A simplified pipeline is:

                    Puzzle Image
                         │
                         ▼
                  Piece Detection
                         │
                         ▼
                  Piece Bounding Box
                         │
              ┌──────────┴──────────┐
              │                     │
              ▼                     ▼
       Side Classification    Rotation Regression
              │                     │
              ▼                     ▼
       Piece Structure        Piece Orientation
              │                     │
              └──────────┬──────────┘
                         ▼
                  Puzzle Solving

---

Dataset

The model was trained using the PuzzleMap dataset.

The training examples are based on annotated puzzle pieces stored in the pieces dataset structure.

Each training sample uses:

Image
Bounding Box
Rotation
Validity

Only pieces with:

valid = true

and a defined:

rotation

were included in the training set.

The model version 2 incorporates additional examples compared with the previous version, increasing the source dataset from approximately 108 to 973 available annotated examples.

---

Reproducibility

The train/validation split uses:

random_state=42

with:

Validation size: 20%
Training size: 80%

The main training configuration is:

Backbone:
    facebook/dinov2-base

Input:
    224 × 224

Batch size:
    32

Loss:
    SmoothL1Loss(beta=0.1)

Stage 1:
    LR = 1e-4
    Backbone frozen
    Maximum epochs = 30

Stage 2:
    LR = 5e-6
    Backbone trainable
    Maximum epochs = 40

Early stopping:
    Patience = 4

---

Version History

Version 2

The second version of the model introduces a larger training set and follows the standardized Hugging Face Hub repository structure.

Main changes:

• Increased the number of available training examples from approximately 108 to 973.
• Standardized the model repository for Hugging Face Hub.
• Continued using facebook/dinov2-base as the visual backbone.
• Maintained the two-stage training strategy.
• Uses normalized [sin(angle), cos(angle)] regression.
• Uses Smooth L1 Loss with beta=0.1.

The released model corresponds to the best validation-loss checkpoint from the second fine-tuning stage.

---

Citation

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

Pablo Moreira.
Puzzle Piece Rotation Regressor.
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.

The training dataset may contain data originating from third-party datasets. Users are responsible for verifying the licensing requirements of the underlying datasets, pretrained models and other dependencies before using this model in their own applications.