puzzle-piece-detection

Puzzle Piece Detection

An object detection model for detecting individual jigsaw puzzle pieces in images containing one or more pieces.

The model is based on DETR ResNet-50 and is fine-tuned to detect a single object category:

puzzle-piece

This model is part of the PuzzleMap project and provides the piece localization stage used by subsequent puzzle-analysis components.

Model Details

The model uses the standard DETR architecture with a ResNet-50 backbone and a custom classification configuration containing a single foreground class.

---

Intended Use

The model is intended to detect individual jigsaw puzzle pieces in photographs or images containing a collection of pieces.

A typical input can contain:

+---------------------------------------+
|                                       |
|      ┌───────────┐                    |
|      │           │        ┌───────┐   |
|      │   piece   │        │ piece │   |
|      │           │        │       │   |
|      └───────────┘        └───────┘   |
|                                       |
|             ┌────────────┐            |
|             │    piece   │            |
|             └────────────┘            |
|                                       |
+---------------------------------------+

The model produces bounding boxes identifying the location of each detected puzzle piece.

For each detection, the model provides:

bounding box
class label
confidence score

The only foreground class is:

puzzle-piece

---

Architecture

The model uses facebook/detr-resnet-50 as its base architecture.

The architecture can be summarized as:

Input Image
     │
     ▼
Image Processor
     │
     ▼
ResNet-50 Backbone
     │
     ▼
DETR Transformer
     │
     ▼
Object Queries
     │
     ├──────────────────────┐
     │                      │
     ▼                      ▼
Class Prediction       Bounding Box
     │                      │
     ▼                      ▼
puzzle-piece            [x, y, w, h]

DETR performs object detection using a fixed set of object queries and predicts both the object class and its bounding box.

The original DETR classification configuration is adapted to the PuzzleMap task so that the foreground object is represented by:

0 → puzzle-piece

DETR also internally represents the special no-object state for queries that do not correspond to a detected puzzle piece.

---

Labels

The model contains a single foreground label:

0 = puzzle-piece

The corresponding configuration is:

id2label = {
    0: "puzzle-piece"
}

label2id = {
    "puzzle-piece": 0
}

The model therefore does not attempt to distinguish different types of puzzle pieces.

Its purpose is exclusively to answer:

Where are the puzzle pieces in this image?

---

Input Processing

The training images are loaded as RGB images:

image = Image.open(path).convert("RGB")

Annotations are stored as bounding boxes using the format:

[x, y, width, height]

where:

x      = left coordinate
y      = top coordinate
width  = bounding-box width
height = bounding-box height

The annotations are converted to the format expected by DetrImageProcessor:

{
    "bbox": [x, y, width, height],
    "category_id": 0,
    "area": width * height,
    "iscrowd": 0
}

The processor is initialized from the base DETR model:

processor = DetrImageProcessor.from_pretrained(
    "facebook/detr-resnet-50"
)

The same processor configuration is saved together with the released model.

---

Dataset

Training uses the PuzzleMap dataset.

The training annotations are obtained from the puzzle-piece annotations in the dataset.

Only pieces marked as valid are included as detection targets.

For each source image, all valid puzzle-piece bounding boxes are collected. This allows a single image to contain multiple detection targets.

The dataset is divided into:

Training set
Validation set

using a fixed random seed.

The validation split is used to monitor the model during training and to determine when training should stop.

---

Training

The model is fine-tuned from:

facebook/detr-resnet-50

The classification head is adapted to the single PuzzleMap object category:

puzzle-piece

Training Configuration

The model was trained using:

Optimizer: AdamW
Learning rate: 1e-4
Backbone learning rate: 1e-5
Weight decay: 1e-4
Maximum epochs: 30
Batch size: 2
Gradient accumulation: 4 batches
Gradient clipping: 0.1
Early stopping patience: 5

Gradient accumulation is used to increase the effective batch size while keeping the individual GPU batches small.

The optimizer uses separate parameter groups so that the ResNet backbone is fine-tuned with a lower learning rate than the remaining model parameters.

Conceptually:

DETR parameters
      │
      ├── Backbone
      │      └── learning rate: 1e-5
      │
      └── Other parameters
             └── learning rate: 1e-4

---

Loss

Training uses the standard loss formulation provided by Hugging Face's DetrForObjectDetection.

The DETR loss combines multiple components related to:

• object classification;
• bounding-box regression;
• generalized IoU;
• cardinality.

The model therefore learns both what is present and where it is located.

The training objective can be represented conceptually as:

DETR Loss
   │
   ├── Classification Loss
   ├── Bounding Box Loss
   ├── Generalized IoU Loss
   └── Cardinality Error

The loss is computed by the underlying DetrForObjectDetection implementation.

---

Training Framework

The training loop is implemented using PyTorch Lightning.

The Lightning module wraps the Hugging Face model and handles:

• training steps;
• validation steps;
• optimizer configuration;
• metric logging;
• gradient accumulation;
• gradient clipping;
• early stopping.

Validation loss is monitored during training.

Training stops early when the monitored validation loss does not improve according to the configured patience.

---

Inference

Installation

Install the required packages:

pip install torch torchvision transformers pillow

Loading the Model

The released model can be loaded using the standard Hugging Face Transformers API.

No custom remote model code is required.

from transformers import (
    DetrForObjectDetection,
    DetrImageProcessor,
)

MODEL_ID = "pablo-moreira/puzzle-piece-detection"

processor = DetrImageProcessor.from_pretrained(
    MODEL_ID
)

model = DetrForObjectDetection.from_pretrained(
    MODEL_ID
)

model.eval()

---

Running Inference

Load an image:

from PIL import Image

image = Image.open(
    "puzzle_image.png"
).convert("RGB")

Prepare the image:

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

Run the model:

import torch

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

The raw model output contains the classification logits and predicted bounding boxes:

outputs.logits
outputs.pred_boxes

The tensors correspond to the DETR object queries.

The raw bounding boxes are represented using normalized coordinates internally.

---

Post-Processing Predictions

The processor provides a helper for converting the raw DETR outputs into image-space bounding boxes.

target_sizes = [
    (image.height, image.width)
]

results = processor.post_process_object_detection(
    outputs,
    target_sizes=target_sizes
)[0]

The resulting dictionary contains:

results["scores"]
results["labels"]
results["boxes"]

The bounding boxes are returned in image coordinates using:

[x_min, y_min, x_max, y_max]

For example:

for score, label, box in zip(
    results["scores"],
    results["labels"],
    results["boxes"],
):
    print(
        score.item(),
        label.item(),
        box.tolist()
    )

The label ID can be converted to the corresponding class name using the model configuration:

label_id = label.item()

label_name = model.config.id2label[
    label_id
]

For this model, the foreground class is:

puzzle-piece

---

Complete Inference Example

The following example shows how to run inference with the trained model using sample images from the PuzzleMap dataset.

The example downloads the first image from Hugging Face, loads the trained model, detects the puzzle pieces, and displays the resulting bounding boxes.

from transformers import (
    DetrForObjectDetection,
    DetrImageProcessor,
)

from PIL import Image, ImageDraw
import torch
from io import BytesIO
import requests


MODEL_ID = "pablo-moreira/puzzle-piece-detection"

IMAGES = [
    "https://huggingface.co/datasets/pablo-moreira/puzzle-map/resolve/main/pieces/120_avengers_20260709_102651_6a038b96ef534b848f91fec292004758.jpg",
    "https://huggingface.co/datasets/pablo-moreira/puzzle-map/resolve/main/pieces/100_dc_20260711_121203_1.jpg",
    "https://huggingface.co/datasets/pablo-moreira/puzzle-map/resolve/main/pieces/100_dc_20260708_214724_0dfe5828a0a8493f97db5124b80ad584.jpg",
]


# Load processor
processor = DetrImageProcessor.from_pretrained(
    MODEL_ID
)

# Load model
model = DetrForObjectDetection.from_pretrained(
    MODEL_ID
)

model.eval()


# Load first 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)


# Post-process detections
results = processor.post_process_object_detection(
    outputs,
    target_sizes=[
        (image.height, image.width)
    ]
)[0]


# Draw detections
draw = ImageDraw.Draw(image)

for score, label, box in zip(
    results["scores"],
    results["labels"],
    results["boxes"],
):

    label_id = label.item()

    label_name = model.config.id2label[
        label_id
    ]

    x1, y1, x2, y2 = box.tolist()

    draw.rectangle(
        (x1, y1, x2, y2),
        outline="red",
        width=3,
    )

    draw.text(
        (x1, y1),
        label_name,
        fill="red",
    )


display(image)

The IMAGES list contains additional sample images that can be used to test the model:

1. 120_avengers_20260709_102651_6a038b96ef534b848f91fec292004758.jpg
2. 100_dc_20260711_121203_1.jpg
3. 100_dc_20260708_214724_0dfe5828a0a8493f97db5124b80ad584.jpg

To test another image, simply change:

response = requests.get(IMAGES[0])

to, for example:

response = requests.get(IMAGES[1])

or:

response = requests.get(IMAGES[2])

The resulting image contains a bounding box around each detected puzzle piece.

---

Output

For each detected object, the post-processed output provides:

results["scores"]

Confidence values for the detections.

results["labels"]

Predicted class IDs.

results["boxes"]

Predicted bounding boxes in image coordinates.

The bounding-box format is:

[x_min, y_min, x_max, y_max]

A downstream application can use these coordinates to crop individual puzzle pieces for additional analysis.

---

Example Pipeline

The detector is designed to be used as an early stage of the PuzzleMap computer-vision pipeline.

A typical workflow is:

Puzzle Image
     │
     ▼
Puzzle Piece Detection
     │
     ▼
Bounding Boxes
     │
     ▼
Individual Piece Cropping
     │
     ├───────────────┐
     │               │
     ▼               ▼
Side Classification  Other Piece Analysis
     │
     ▼
Piece Properties
     │
     ▼
Puzzle Assembly

The detection model is therefore independent of the later classification and matching models.

---

Limitations

Detection Quality

The model is trained specifically to detect jigsaw puzzle pieces.

Performance may degrade when images contain:

• objects that do not resemble puzzle pieces;
• heavily occluded pieces;
• pieces that are extremely small in the image;
• severe motion blur;
• strong perspective distortion;
• very poor lighting;
• pieces that overlap significantly;
• incomplete or damaged puzzle pieces.

Domain Specificity

The model is specialized for jigsaw puzzles.

It should not be considered a general-purpose object detector.

The single foreground class means that the model does not distinguish between:

• border pieces;
• corner pieces;
• edge pieces;
• internal pieces;
• different puzzle manufacturers;
• different puzzle themes.

All detected pieces are represented as:

puzzle-piece

Further classification can be performed by other PuzzleMap models.

Image Composition

The model was trained using images from the PuzzleMap dataset. Images with characteristics substantially different from the training distribution may produce different detection behavior.

---

Model Files

The repository contains the standard Hugging Face model artifacts generated by:

model.save_pretrained(...)

and:

processor.save_pretrained(...)

The model can therefore be loaded directly with:

from transformers import DetrForObjectDetection

model = DetrForObjectDetection.from_pretrained(
    "pablo-moreira/puzzle-piece-detection"
)

The corresponding image processor can be loaded with:

from transformers import DetrImageProcessor

processor = DetrImageProcessor.from_pretrained(
    "pablo-moreira/puzzle-piece-detection"
)

The model does not require trust_remote_code=True.

---

Base Model

This model is fine-tuned from:

facebook/detr-resnet-50

The base architecture is DETR (DEtection TRansformer) with a ResNet-50 visual backbone.

The pretrained backbone and architecture originate from the original DETR model and its associated pretrained weights. Users should review the licensing and usage terms of the base model in addition to the license of this repository.

---

Relation to PuzzleMap

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

The broader project uses computer vision and machine learning to assist with the analysis and assembly of jigsaw puzzles.

The detection stage provides the bounding boxes required by subsequent components, which can then:

1. crop individual puzzle pieces;
2. classify the four sides of each piece;
3. determine piece orientation;
4. identify similar pieces;
5. assist in reconstructing the puzzle.

The detection model therefore serves as the localization component of the PuzzleMap pipeline.

---

Version Comparison

The following table compares the results from versions.

Version 6

• Uses the puzzle-map dataset with 6,946 examples.
• 66 duplicate items were removed from the dataset.
• 66 new puzzle-focus items were added to the dataset.
• 6 items had their mirroring corrected.
• The resulting dataset size remains 6,946 examples after removing duplicates and adding new items.
• The version was trained and evaluated using the same stratified train/validation split introduced in version 5, allowing direct comparison with the previous version.

Version 5

• Uses the puzzle-map dataset with 6946 examples
• Some examples from the dataset with the "puzzle-pieces" prefix had their bounding boxes corrected.
• Examples without bounding boxes will be used to train the "no-object" class.
• 504 hard negatives examples without puzzle pieces were added from the COCO dataset.
• Version 4 of the model was used to calculate the most difficult examples without a piece.
• The training and validation split uses stratification based on the number of puzzle pieces in each image, ensuring a more representative distribution of images with zero, one, or multiple puzzle pieces across both splits.
• Rare piece-count categories with very few examples were grouped into a shared stratum before splitting to allow stratified sampling.
• This version uses a validation split that is more representative of the overall dataset, making the evaluation more reliable when comparing future model versions.
• Previous experiments showed that increasing no-object examples changes the detector's operating behavior, generally favoring fewer false positives at the cost of a small increase in false negatives.
• Version 4 was also evaluated using the validation dataset from this release, allowing a direct comparison between the previous model and the new version using the same validation set.

The best checkpoint from Version 4 was also evaluated using the validation split introduced in this version to enable a direct comparison using the same evaluation dataset.

Although the Version 4 checkpoint achieved lower validation losses and slightly higher recall, the new version produced fewer false positives and achieved a better F1 score and cardinality error on the same validation dataset.

The comparison indicates that the additional negative examples and the stratified dataset split improved the model's ability to estimate the number of puzzle pieces and reduced false positive detections, with a small trade-off in recall.

Citation

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

Pablo Moreira.
Puzzle Piece Detection.
PuzzleMap project.

---

License

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

The underlying facebook/detr-resnet-50 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.