Fine-Tuning PaliGemma 2 for Object Detection

In this article, we will be fine-tuning the PaliGemma 2 VLM for object detection. Nowadays, VLMs are great at OCR, image captioning, and video understanding out of the box. Along with that, they are also catching up with object detection. However, an extremely custom use case for object detection is still a struggle for many VLMs. That’s why we will tackle one of the real-world use cases of object detection with the PaliGemma 2 VLM here.

Here, we explore one such difficult use case: wheat head detection for agricultural applications.

Paligemma 2 is an open VLM (Vision Language Model) from Google based on the Gemma 2 LLM and SigLIP vision encoder. The model family includes both pretrained and task-mixed checkpoints, with the latter performing reasonably well on downstream tasks such as OCR, segmentation, and generic object detection.

However, wheat head detection presents a far more difficult scenario:

• The objects are extremely small.
• Densely packed.
• Visually repetitive.
• And often partially occluded.

Not to mention the extreme requirement of tens of gigabytes of VRAM for fine-tuning.

This makes it an interesting benchmark for testing the practical limits of VLM-based detection systems.

The goal of this article is not to achieve state-of-the-art results or claim production readiness. Instead, we will experiment with fine-tuning PaliGemma 2 on a highly specialized detection task and analyze:

• how far the model can be pushed,
• where it struggles,
• and whether VLMs are becoming viable alternatives to traditional detectors for niche applications.

It is also worth noting that PaliGemma 2 is no longer a cutting-edge model by 2026 standards, having originally been released in late 2024. Nevertheless, it remains one of the most accessible open VLMs for experimentation and research on consumer hardware.

The Global Wheat Detection Dataset

The Global Wheat Detection competition started in 2020. Since then, there have been multiple submissions and benchmarks for the same in the subsequent years as the competition kept running. The most prominent have been models like YOLO, EfficientDet, and FasterRCNN, which perform well when fine-tuned properly on the dataset. However, VLM benchmarks and training pipelines remain sparse on the same.

We will use the 2021 version of the dataset from Hugging Face in this article. It contains around 3.6K training, 1.48K validation, and 1.38K test samples across 300K annotations.

The following figure shows what the dataset rows look like.

Let’s take a look at one of the annotated images.

The images are densely packed with wheat head images. This is perhaps one of the best datasets to fine-tune VLMs on for object detection to understand how well they can perform on small objects and dense scenes.

Project Directory Structure

We are using the following directory structure.

├── paligemma2-pt-448-wheat
│   ├── checkpoint-1000
│   ├── checkpoint-1500
│   ├── checkpoint-2000
│   ├── checkpoint-2500
│   ├── checkpoint-3000
│   ├── checkpoint-500
│   ├── runs
│   ├── adapter_config.json
│   ├── adapter_model.safetensors
│   ├── processor_config.json
│   ├── README.md
│   ├── tokenizer_config.json
│   ├── tokenizer.json
│   └── training_args.bin
├── test_data
│   ├── image_1.jpg
│   ├── image_2.jpg
│   ├── image_3.jpg
│   └── image_4.jpg
├── paligemma2_3b_pt_448_wheat.ipynb
└── paligemma2_od_inference.py

• We have one Jupyter Notebook for fine-tuning and a Python script for inference.
• The paligemma2-pt-448-wheat contains the trained checkpoints, and the inference data resides in test_data.

The Jupyter Notebook, Gradio inference script, and trained adapters are available for download as a zip file.

Download Code

We have everything ready now, let’s jump into the coding part.

Training PaliGeamm 2 for Wheat Head Detection

All the fine-tuning code is present in the paligemma2_3b_pt_448_wheat.ipynb Jupyter Notebook.

We will follow the important parts of the notebook first, and then jump into the inference.

Note: The fine-tuning process requires upwards of 75GB VRAM. Although we are training an adapter here, we are using LoRA instead of QLoRA, which means loading the backbone in FP16/BF16 format instead of INT4. You can reduce the batch size from 2 to 1 to bring down the VRAM requirement to around 60GB.

Installing Dependencies and Importing Libraries

The first code cell installs all the necessary libraries and frameworks.

!pip install \
transformers==5.8.0 \
datasets==4.8.5 \
trl==1.4.0 \
supervision==0.28.0 \
albumentations==2.0.8 \
tensorboard==2.20.0 \
peft==0.19.1 \
torchao==0.17.0

The second code block imports all the necessary modules.

import json
import requests
import supervision as sv
import numpy as np
import torch
import re
import albumentations as A
import os

from PIL import Image
from io import BytesIO
from datasets import load_dataset
from transformers import (
    PaliGemmaProcessor,
    PaliGemmaForConditionalGeneration,
)
from peft import LoraConfig, get_peft_model
from trl import SFTConfig, SFTTrainer
from functools import partial

Preparing the Global Wheat Detection Dataset for PaliGemma 2 Training

Let’s load the dataset and manage the splits.

dataset = load_dataset('sovitrath/GlobalWheatHeadDataset2021')

print(dataset)

DatasetDict({
    train: Dataset({
        features: ['image', 'domain', 'country', 'location', 'development_stage', 'objects'],
        num_rows: 3655
    })
    validation: Dataset({
        features: ['image', 'domain', 'country', 'location', 'development_stage', 'objects'],
        num_rows: 1476
    })
    test: Dataset({
        features: ['image', 'domain', 'country', 'location', 'development_stage', 'objects'],
        num_rows: 1381
    })
})

The following code blocks show how the annotations are structured.

sample_annots = dataset['train']['objects'][0]['boxes']
print(sample_annots)

[[936, 4, 41, 78], [966, 0, 57, 114], [665, 9, 130, 70], [877, 0, 81, 20], [742, 0, 110, 132], [648, 45, 252, 128], [652, 86, 218, 101], [958, 141, 52, 76], [786, 104, 168, 176], [986, 173, 37, 198], [714, 232, 290, 208], [606, 437, 76, 64], [692, 498, 114, 61], [816, 376, 62, 160], [894, 481, 74, 208], [937, 464, 76, 120], [834, 629, 65, 40], [608, 614, 194, 101], [606, 725, 174, 246], [961, 714, 62, 62], [981, 705, 42, 26], [994, 969, 29, 54], [466, 789, 404, 64], [165, 484, 117, 61], [5, 469, 108, 258], [0, 596, 236, 173], [253, 589, 237, 98], [68, 696, 280, 221], [205, 688, 241, 300], [394, 725, 117, 290], [165, 993, 109, 30], [194, 893, 84, 112], [116, 910, 104, 113], [0, 845, 104, 178], [386, 341, 282, 101], [322, 378, 138, 77], [376, 441, 122, 60], [337, 232, 70, 161], [198, 244, 206, 109], [204, 116, 40, 58], [93, 66, 140, 266], [174, 84, 161, 309], [105, 264, 206, 157], [0, 305, 185, 160], [0, 0, 142, 169], [130, 0, 56, 49], [0, 634, 24, 25]]

We have the bounding box annotations as list of lists in format: [x1, y1, width, height].

The following code block prepares one sample for visualization.

# Prepare sample parsed labels.
parsed_labels = []
for sample_annot in sample_annots:
    sample_annot = [
        sample_annot[0],
        sample_annot[1],
        sample_annot[0] + sample_annot[2],
        sample_annot[1] + sample_annot[3]
    ]
    parsed_labels.append(('wheat_head', sample_annot))

def get_annotated_image(image, parsed_labels):
    if not parsed_labels:
        return image

    xyxys = []
    labels = []

    for label, bbox in parsed_labels:
        xyxys.append(bbox)
        labels.append(label)

    detections = sv.Detections(xyxy=np.array(xyxys))

    bounding_box_annotator = sv.BoxAnnotator(color_lookup=sv.ColorLookup.INDEX)
    label_annotator = sv.LabelAnnotator(color_lookup=sv.ColorLookup.INDEX)

    annotated_image = bounding_box_annotator.annotate(
        scene=image, detections=detections
    )
    annotated_image = label_annotator.annotate(
        scene=annotated_image, detections=detections, labels=labels
    )

    return annotated_image

annotated_image = get_annotated_image(sample_image, parsed_labels)
annotated_image

We are using the Supervision library for easier visualization. One thing to note is that we are adding the wheat_head label manually to each bounding box annotation. We will consistently use the same label across training and inference.

We will use the test dataset for evaluation during the training phase.

train_dataset = dataset['train']
val_dataset = dataset['test']
train_dataset, val_dataset

Loading the PaliGemma 2 Model

Let’s load the PaliGemma 2 Pretrained model. The pretrained model is the base model that has not been fine-tuned for downstream tasks like OCR, object detection, or image segmentation. This gives us a bit of freedom to define our own prompt structure when preparing the labels.

model_id = 'google/paligemma2-3b-pt-448'

model = PaliGemmaForConditionalGeneration.from_pretrained(
    model_id, torch_dtype=torch.bfloat16, device_map='auto'
).eval()
processor = PaliGemmaProcessor.from_pretrained(model_id, use_fast=True)

Now, we will load one sample from the validation dataset and run inference using the above-loaded model. This will give us an idea of how the model performs before fine-tuning.

image = val_dataset[13]['image']
caption = 'wheat_head'

prompt = f"detect {caption}"
model_inputs = processor(text=prompt, images=image, return_tensors='pt').to(torch.bfloat16).to(model.device)
input_len = model_inputs['input_ids'].shape[-1]

with torch.inference_mode():
    generation = model.generate(**model_inputs, max_new_tokens=100, do_sample=False)
    generation = generation[0][input_len:]
    output = processor.decode(generation, skip_special_tokens=True)
    print(output)

# https://github.com/ariG23498/gemma3-object-detection/blob/main/utils.py#L17
def parse_paligemma_labels(label, width, height):
    predictions = label.strip().split(";")
    results = []

    for pred in predictions:
        pred = pred.strip()
        if not pred:
            continue

        loc_pattern = r""
        locations = [int(loc) for loc in re.findall(loc_pattern, pred)]

        if len(locations) != 4:
            continue

        category = pred.split(">")[-1].strip()

        y1_norm, x1_norm, y2_norm, x2_norm = locations
        x1 = (x1_norm / 1024) * width
        y1 = (y1_norm / 1024) * height
        x2 = (x2_norm / 1024) * width
        y2 = (y2_norm / 1024) * height

        results.append((category, [x1, y1, x2, y2]))

    return results

width, height = image.size
parsed_labels = parse_paligemma_labels(output, width, height)
parsed_labels

annotated_image = get_annotated_image(image, parsed_labels)

annotated_image

We are using the following prompt format for the model:

Detection input format: detect [CAPTION]

For consistency, we will use the same format when creating the labels for training.

The following image shows the results.

As we can see, the model is detecting the entire area as a wheat head. Hopefully, the performance will improve after fine-tuning.

Preparing the LoRA Adapter and Fine-Tuning

The following code block prepares the LoRA.

target_modules = [
    'q_proj',
    'v_proj',
    'fc1',
    'fc2',
    'linear',
    'gate_proj',
    'up_proj',
    'down_proj'
]

# Configure LoRA
peft_config = LoraConfig(
    lora_alpha=16,
    lora_dropout=0.05,
    r=8,
    bias='none',
    target_modules=target_modules,
    task_type='CAUSAL_LM',
)

# Apply PEFT model adaptation
peft_model = get_peft_model(model, peft_config)

# Print trainable parameters
peft_model.print_trainable_parameters()

We are using a rank of 8 and an alpha of 16. This particular configuration generates an adapter consisting of around 12.16 million trainable parameters.

Next, we define the Supervised Fine-Tuning configuration.

training_args = SFTConfig(
    output_dir='paligemma2-pt-448-wheat',
    per_device_train_batch_size=2,
    per_device_eval_batch_size=2,
    gradient_accumulation_steps=4,
    gradient_checkpointing=False,
    learning_rate=1e-05,
    # num_train_epochs=1,
    max_steps=3000,
    logging_steps=10,
    eval_steps=500,
    eval_strategy='steps',
    save_steps=500,
    bf16=True,
    report_to=['tensorboard'],
    dataset_kwargs={'skip_prepare_dataset': True},
    remove_unused_columns=False,
    # push_to_hub=True,
    dataloader_pin_memory=False,
    label_names=['labels'],
)

The current batch size is 2 with 4 gradient accumulation steps. You can reduce the batch size according to the VRAM at your disposal.

We are training for 3000 steps while saving and logging every 500 steps.

The next code block performs some of the most important aspects before training. These include:

• Converting the dataset to prompt strings with proper formatting.
• Resizing the images to adjust VRAM.
• And collating the samples for batches.

def coco_to_xyxy(coco_bbox):
    x, y, width, height = coco_bbox
    x1, y1 = x, y
    x2, y2 = x + width, y + height
    return [x1, y1, x2, y2]

def convert_to_detection_string(bboxs, image_width, image_height, category):
    def format_location(value, max_value):
        return f""

    detection_strings = []
    for bbox in bboxs:
        x1, y1, x2, y2 = coco_to_xyxy(bbox)
        if x2 > x1 and y2 > y1:
            locs = [
                format_location(y1, image_height),
                format_location(x1, image_width),
                format_location(y2, image_height),
                format_location(x2, image_width),
            ]
            detection_string = ''.join(locs) + f" {category}"
            detection_strings.append(detection_string)
    return ' ; '.join(detection_strings)

def format_objects(example):
    height = example['height']
    width = example['width']
    bboxs = example['bbox']
    # category = example['caption'][0]
    category = 'wheat_head'
    formatted_objects = convert_to_detection_string(bboxs, width, height, category)
    return {'label_for_paligemma': formatted_objects}

resize_size = 448

augmentations = A.Compose([
    A.Resize(height=resize_size, width=resize_size),
], bbox_params=A.BboxParams(
    format='coco', 
    label_fields=['category_ids'], 
    filter_invalid_bboxes=True,
))

# Create a data collator to encode text and image pairs
def collate_fn(examples, transform=None):
    images = []
    prompts = []
    suffixes = []
    for sample in examples:
        # print(sample)
        if transform:
            bboxes = sample['objects']['boxes']
            # Unnormalized xmin and ymin should be less than image width and width.
            bboxes = [b for b in bboxes if b[0] < sample['image'].size[0] and b[1] < sample['image'].size[1]]
            category_ids = ['wheat_head'] * len(bboxes)
            # print(category_ids)
            transformed = transform(
                image=np.array(sample['image']), 
                bboxes=bboxes, 
                category_ids=category_ids
            )
            sample['image'] = transformed['image']
            sample['bbox'] = transformed['bboxes']
            sample['caption'] = transformed['category_ids']
            sample['height'] = sample['image'].shape[0]
            sample['width'] = sample['image'].shape[1]
            sample['label_for_paligemma'] = format_objects(sample)['label_for_paligemma']
        images.append([sample['image']])
        prompts.append(f"Detect {sample['caption']}.")
        suffixes.append(sample['label_for_paligemma'])
    batch = processor(images=images, text=prompts, suffix=suffixes, return_tensors='pt', padding=True)

    # The labels are the input_ids, and we mask the padding tokens in the loss computation
    labels = batch['input_ids'].clone()  # Clone input IDs for labels
    image_token_id = processor.tokenizer.convert_tokens_to_ids('')
    # Mask tokens for not being used in the loss computation
    labels[labels == processor.tokenizer.pad_token_id] = -100
    labels[labels == image_token_id] = -100
    batch['labels'] = labels

    batch['pixel_values'] = batch['pixel_values'].to(model.device)
    return batch

train_collate_fn = partial(
    collate_fn, transform=augmentations
)

In the collate_fn function, we have the following line:

bboxes = [b for b in bboxes if b[0] < sample['image'].size[0] and b[1] < sample['image'].size[1]]

This filters out any samples that have the xmax and ymax coordinates smaller than equal to the xmin and ymin coordinates respectively. Otherwise, Albumentation throws an error.

Finally, we start the fine-tuning process.

trainer = SFTTrainer(
    model=peft_model,
    args=training_args,
    data_collator=train_collate_fn,
    train_dataset=train_dataset,
    eval_dataset=val_dataset,
)

trainer.train()

The following are the training logs.

We will get more insights as to how the model performs when running inference.

Before that, let’s save the model and processor.

processor.save_pretrained(training_args.output_dir)
trainer.save_model(training_args.output_dir)

We are done with the fine-tuning of the PaliGemma 2 model for wheat head detection here. The next section covers the inference using a Gradio application.

Inference Using the Fine-Tuned PaliGemma 2 Model

The inference code is present in paligemma2_od_inference.py file. We will not cover the code here. Rather, we will focus on the inference experiments and results. If you do not have Gradio installed, you can install it via:

pip install gradio

Let’s execute the application.

python paligemma2_od_inference.py

The following video shows one inference experiment after loading the model.

The model seems to be working here. However, upon closer inspection, we can see that it misses a lot of the wheat heads as well. This is a mixed result at the moment, with a lot of room for improvement.

The following figure shows three more results where the model detects wheat heads inconsistently.

This clearly shows the limitations of our current approach.

Key Takeaways

Although we fine-tuned the PaliGemma 2 model for object detection in this article, we faced hurdles along the way:

• The GPU VRAM requirements are more than modest. With a batch size of 2, we had to rely on an 80GB VRAM GPU for fine-tuning.
• Although the model was trained for 3000 steps, the results were not ideal. Out of 4 inference experiments, the model was able to detect the wheat heads only sporadically.

More recent VLMs like Qwen3-VL might perform much better on dense object detection after fine-tuning. However, they are also subject to experiment.

Summary and Conclusion

In this article, we tried fine-tuning the PaliGemma 2 model for custom object detection. We discussed the dataset preparation and the training steps. However, the results were not up to the mark. We also discussed the limitations. In future articles, we will try to explore how to mitigate such issues and improve the results, along with experimenting with newer VLMs.

If you have any questions, thoughts, or suggestions, please leave them in the comment section. I will surely address them.

You can contact me using the Contact section. You can also find me on LinkedIn, and X.