Teaching PaliGemma 2 to Spot Construction-Site Hazards

TL;DR

Use google/paligemma2-3b-pt-448 as the practical fine-tuning checkpoint for construction-site object detection.
Convert COCO-style boxes into PaliGemma’s text format: <loc0000><loc0000><loc0000><loc0000> label.
Train with QLoRA through transformers, peft, and bitsandbytes, then run the same parser for ground truth and predictions.

Abstract

Construction safety detection is a good fit for a task-adapted vision-language model (VLM): the objects are domain-specific, the prompts are natural language, and the output can stay inspectable as text. This tutorial turns a notebook-style PaliGemma 2 experiment into a reproducible transformers 5 workflow for detecting hardhats, masks, safety vests, trucks, equipment, and related construction-site objects. You will load a public dataset, serialize its bounding boxes into PaliGemma detection labels, fine-tune a 3B PaliGemma 2 checkpoint with QLoRA, and run inference with a small parser that turns generated location tokens back into boxes.

Requirements

Python 3.10+
A CUDA GPU for training
Access to the gated PaliGemma 2 checkpoints on Hugging Face
transformers>=5.12.1
datasets
peft
bitsandbytes
accelerate
pillow
torch

pip install -U "transformers>=5.12.1" datasets peft bitsandbytes accelerate pillow torch

Why PaliGemma 2

There are newer general-purpose VLM families, but PaliGemma 2 remains the cleanest choice for this particular tutorial because object localization is part of the model interface. It can answer detect ... prompts with structured <loc> tokens, and the Hugging Face transformers implementation exposes a direct PaliGemmaForConditionalGeneration path for fine-tuning.

For this post, use:

google/paligemma2-3b-pt-448 for fine-tuning. It is the runnable default: 3B parameters, 448px input resolution, and intended for downstream adaptation.
google/paligemma2-3b-mix-224 for quick out-of-the-box detection experiments before fine-tuning.
google/paligemma2-10b-pt-448 or google/paligemma2-28b-pt-448 only if you have the memory budget and want to scale the same recipe.

The 448px pretrained checkpoint is a better training default than the 224px checkpoint for small safety objects, while the 3B size keeps the tutorial practical.

The Task

The dataset is keremberke/construction-safety-object-detection, which contains images and COCO-style object annotations. The labels include personal protective equipment (PPE), missing PPE, workers, vehicles, and equipment:

CATEGORY_LABELS = [
    "barricade",
    "dumpster",
    "excavators",
    "gloves",
    "hardhat",
    "mask",
    "no-hardhat",
    "no-mask",
    "no-safety vest",
    "person",
    "safety net",
    "safety shoes",
    "safety vest",
    "dump truck",
    "mini-van",
    "truck",
    "wheel loader",
]

Instead of training a classical detector head, you train the model to generate text like this:

<loc0122><loc0308><loc0469><loc0615> hardhat ; <loc0204><loc0118><loc0930><loc0712> person

Each object is a sequence of four normalized location tokens followed by the class label. The order is:

y_min, x_min, y_max, x_max

Load the Dataset

from datasets import DatasetDict, load_dataset

DATASET_ID = "keremberke/construction-safety-object-detection"

dataset: DatasetDict = load_dataset(DATASET_ID, name="full")
print(dataset)
print(dataset["train"][0].keys())

Each row contains an image plus an objects field with bbox and category arrays. The bounding boxes are in COCO format:

x_min, y_min, width, height

PaliGemma expects normalized location tokens, so the first real step is serialization.

Convert COCO Boxes to PaliGemma Labels

LOC_BINS = 1024
MAX_LOC = LOC_BINS - 1


def coco_to_xyxy(box: list[float]) -> tuple[float, float, float, float]:
    """Convert a COCO [x, y, width, height] box into [x1, y1, x2, y2]."""
    x, y, width, height = box
    return x, y, x + width, y + height


def to_loc_token(value: float, max_value: int) -> str:
    """Scale an image coordinate into PaliGemma's <loc0000> token range."""
    scaled = round(value * LOC_BINS / max_value)
    clipped = max(0, min(MAX_LOC, scaled))
    return f"<loc{clipped:04d}>"


def format_detection_target(example: dict, category_labels: list[str]) -> str:
    """Serialize all boxes in one example as a PaliGemma detection target."""
    width = example["width"]
    height = example["height"]
    boxes = example["objects"]["bbox"]
    categories = example["objects"]["category"]

    detections = []
    for box, category in zip(boxes, categories, strict=True):
        x1, y1, x2, y2 = coco_to_xyxy(box)
        label = category_labels[int(category)]
        loc_tokens = [
            to_loc_token(y1, height),
            to_loc_token(x1, width),
            to_loc_token(y2, height),
            to_loc_token(x2, width),
        ]
        detections.append("".join(loc_tokens) + f" {label}")

    return " ; ".join(detections)


def add_paligemma_target(example: dict) -> dict:
    return {"target": format_detection_target(example, CATEGORY_LABELS)}


prepared = dataset.map(add_paligemma_target)
print(prepared["train"][0]["target"])

This function is intentionally small. If the generated target looks wrong, you only have three places to inspect: the source box format, the coordinate scaling, and the final label string.

Build Detection Prompts

PaliGemma detection prompts work best when you tell the model what to look for. For this dataset, use the labels already present in each image during training:

def labels_in_example(example: dict) -> list[str]:
    label_ids = [int(category) for category in example["objects"]["category"]]
    labels = [CATEGORY_LABELS[label_id] for label_id in label_ids]
    return sorted(set(labels))


def detection_prompt(example: dict) -> str:
    labels = labels_in_example(example)
    return "detect " + " ; ".join(labels)

For open-ended inference, you can pass a broader prompt:

ALL_OBJECTS_PROMPT = "detect " + " ; ".join(CATEGORY_LABELS)

That prompt asks the fine-tuned model to look for every construction-safety class the dataset knows about.

Load PaliGemma 2 with QLoRA

import torch
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from transformers import AutoProcessor, BitsAndBytesConfig, PaliGemmaForConditionalGeneration

MODEL_ID = "google/paligemma2-3b-pt-448"

processor = AutoProcessor.from_pretrained(MODEL_ID)

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)

model = PaliGemmaForConditionalGeneration.from_pretrained(
    MODEL_ID,
    quantization_config=quantization_config,
    device_map="auto",
    torch_dtype=torch.bfloat16,
)

model = prepare_model_for_kbit_training(model)

lora_config = LoraConfig(
    r=8,
    lora_alpha=16,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules=[
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
        "gate_proj",
        "up_proj",
        "down_proj",
    ],
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

The LoRA target modules adapt the language decoder. That keeps the update small while preserving the PaliGemma vision encoder and multimodal projection behavior.

Write the Data Collator

The processor can take a prompt, image, and suffix. The suffix is the expected answer, so this is where the detection target becomes the training label.

def collate_fn(examples: list[dict]) -> dict[str, torch.Tensor]:
    texts = [detection_prompt(example) for example in examples]
    images = [example["image"].convert("RGB") for example in examples]
    suffixes = [example["target"] for example in examples]

    return processor(
        images=images,
        text=texts,
        suffix=suffixes,
        return_tensors="pt",
        padding="longest",
    )

Keep the batch on CPU here. Trainer handles device placement, and avoiding manual dtype conversion prevents accidental casting of token IDs.

Fine-Tune

from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir="paligemma2-construction-safety-qlora",
    num_train_epochs=2,
    per_device_train_batch_size=1,
    per_device_eval_batch_size=1,
    gradient_accumulation_steps=4,
    learning_rate=2e-5,
    weight_decay=1e-6,
    warmup_steps=20,
    logging_steps=25,
    save_steps=500,
    save_total_limit=2,
    eval_strategy="steps",
    eval_steps=500,
    remove_unused_columns=False,
    bf16=True,
    optim="paged_adamw_8bit",
    report_to=["tensorboard"],
    dataloader_pin_memory=False,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=prepared["train"],
    eval_dataset=prepared["validation"],
    data_collator=collate_fn,
    processing_class=processor,
)

trainer.train()

Start with two epochs to prove the pipeline. Then increase training time only after you inspect generated boxes on held-out images.

Save the Adapter

ADAPTER_ID = "your-username/paligemma2-construction-safety-qlora"

trainer.model.push_to_hub(ADAPTER_ID)
processor.push_to_hub(ADAPTER_ID)

If you do not want to publish the adapter, save locally:

trainer.model.save_pretrained("paligemma2-construction-safety-qlora")
processor.save_pretrained("paligemma2-construction-safety-qlora")

Parse Generated Boxes

The model returns text, so inference needs a parser. Keep this parser separate from training so you can reuse it for ground truth labels and generated predictions.

import re
from dataclasses import dataclass
from PIL import Image, ImageDraw

LOCATION_PATTERN = re.compile(
    r"<loc(\d{4})><loc(\d{4})><loc(\d{4})><loc(\d{4})>\s*([^;]+)"
)


@dataclass(frozen=True)
class Detection:
    label: str
    box: tuple[float, float, float, float]


def parse_detections(text: str) -> list[Detection]:
    detections = []
    for y1, x1, y2, x2, label in LOCATION_PATTERN.findall(text):
        detections.append(
            Detection(
                label=label.strip(),
                box=(
                    int(x1) / LOC_BINS,
                    int(y1) / LOC_BINS,
                    int(x2) / LOC_BINS,
                    int(y2) / LOC_BINS,
                ),
            )
        )
    return detections


def draw_detections(image: Image.Image, detections: list[Detection]) -> Image.Image:
    canvas = image.convert("RGB").copy()
    draw = ImageDraw.Draw(canvas)
    width, height = canvas.size

    for detection in detections:
        x1, y1, x2, y2 = detection.box
        pixel_box = (x1 * width, y1 * height, x2 * width, y2 * height)
        draw.rectangle(pixel_box, outline="red", width=3)
        draw.text((pixel_box[0], pixel_box[1]), detection.label, fill="red")

    return canvas

Notice that the parser stores boxes as normalized (x1, y1, x2, y2) tuples. That is easier to draw and compare than keeping PaliGemma’s generated (y1, x1, y2, x2) token order everywhere.

Run Inference

from peft import PeftModel

base_model = PaliGemmaForConditionalGeneration.from_pretrained(
    MODEL_ID,
    quantization_config=quantization_config,
    device_map="auto",
    torch_dtype=torch.bfloat16,
)
finetuned_model = PeftModel.from_pretrained(base_model, ADAPTER_ID)
finetuned_model.eval()

Now generate detections for a held-out image:

def generate_detections(image: Image.Image, prompt: str, max_new_tokens: int = 256) -> str:
    inputs = processor(
        images=image.convert("RGB"),
        text=prompt,
        return_tensors="pt",
    ).to(finetuned_model.device)

    with torch.inference_mode():
        output = finetuned_model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
        )

    return processor.decode(output[0], skip_special_tokens=True)


example = prepared["test"][9]
image = example["image"]

generated_text = generate_detections(image, ALL_OBJECTS_PROMPT)
predicted_detections = parse_detections(generated_text)

draw_detections(image, predicted_detections)

To inspect the ground truth with the same drawing code:

ground_truth = parse_detections(example["target"])
draw_detections(image, ground_truth)

Using the same parser for labels and predictions makes debugging much simpler. If ground truth boxes draw correctly but predictions do not, the issue is model behavior. If both draw incorrectly, the issue is the formatter or parser.

Practical Checks

Before you run a long training job, check these pieces:

The first converted target contains four <loc> tokens per object.
Ground truth boxes draw in the right places on several images.
The prompt contains labels that actually appear in the dataset.
remove_unused_columns=False is set, because the collator needs image, objects, width, height, and target.
Token IDs are not manually cast to bfloat16 in the collator.

For training quality, look beyond the loss. Decode validation examples during training and inspect whether the model learns label names, box count, and rough location before worrying about precise coordinates.

Failure Modes and Caveats

The model may hallucinate objects when the prompt asks for every class at once. If that happens, use narrower prompts such as detect hardhat ; no-hardhat ; person and compare outputs.

Small PPE objects are hard at 224px. The 448px checkpoint improves the input resolution without jumping to the 10B or 28B model sizes.

Generated text is not the same as a calibrated detector score. If this workflow is used for safety review, add a separate evaluation step with mAP, per-class recall, and human inspection. Treat the model as an assistive signal, not a safety authority.

Dataset labels also matter. Classes like no-hardhat and hardhat are visually close, and labeling ambiguity can teach the model inconsistent behavior. Inspect those classes early.

Summary

PaliGemma 2 is still a strong fit for construction-site object detection because it gives you a text-native localization format and a direct fine-tuning path in transformers. The clean version of the workflow is: convert COCO boxes into <loc> targets, train google/paligemma2-3b-pt-448 with QLoRA, and parse generated text back into drawable boxes. Once that baseline works, you can scale resolution, model size, training duration, or evaluation rigor without changing the core data contract.

Next Steps

Add a validation callback that decodes a fixed set of images after each evaluation step.
Compute mAP from parsed predictions so the tutorial has quantitative feedback.
Try class-specific prompts for the most safety-critical labels, especially person, hardhat, no-hardhat, safety vest, and no-safety vest.
Scale to google/paligemma2-10b-pt-448 if the 3B model underfits and your GPU memory allows it.