Fine-Tuning OWL-ViT on a Custom Object Detection Dataset

Fine-tuning OWL-ViT on a custom dataset requires building a complete detection pipeline because OWL-ViT is fundamentally different from traditional detectors like YOLO or Faster-RCNN. Instead of predicting object classes directly, OWL-ViT predicts bounding boxes together with similarity scores between image regions and text descriptions. During training, the model learns how visual patterns correspond to textual class descriptions, allowing it to detect objects by comparing image features with language features. This guide explains the full pipeline for fine-tuning OWL-ViT on a custom COCO-format dataset, including dataset preparation, model training, loss computation, and inference.

Dataset Acquisition

The dataset is downloaded in COCO format using Roboflow. COCO format provides a structured representation of object detection data where each image is associated with bounding boxes and class labels. This structured format allows consistent training across different datasets and ensures compatibility with transformer-based detection models.

!pip install roboflow

from roboflow import Roboflow

rf = Roboflow(api_key="YOUR_KEY")
project = rf.workspace("workspace-name").project("project-name")
version = project.version(6)

dataset = version.download("coco")

This step retrieves the dataset and organizes it into training and validation directories containing images and annotation JSON files. Each annotation contains bounding boxes in pixel coordinates and category identifiers. The dataset structure allows the training pipeline to map images to annotations and class names.

Bounding Box Utility Functions

OWL-ViT predicts bounding boxes in normalized center-coordinate format (cx, cy, w, h), while COCO annotations use pixel-based (x, y, width, height) format. The following functions convert between formats and compute geometric relationships between bounding boxes. These conversions are necessary because training loss calculations require comparing predicted boxes with ground-truth boxes in the same coordinate system.

def box_cxcywh_to_xyxy(boxes: torch.Tensor):

    cx, cy, w, h = boxes.unbind(-1)

    x1 = cx - 0.5 * w
    y1 = cy - 0.5 * h

    x2 = cx + 0.5 * w
    y2 = cy + 0.5 * h

    return torch.stack([x1, y1, x2, y2], dim=-1)

This function converts center-based normalized coordinates into corner-based coordinates, which are required for calculating distances and overlaps between predicted and real bounding boxes. Matching bounding boxes in normalized space ensures that the model learns consistent geometric relationships regardless of image resolution.

Generalized IoU Calculation

Intersection over Union measures how well predicted boxes overlap with real boxes. OWL-ViT training uses Generalized IoU, which improves training stability when boxes do not overlap.

def generalized_box_iou(boxes1, boxes2):

    lt = torch.max(boxes1[:, None, :2], boxes2[None, :, :2])
    rb = torch.min(boxes1[:, None, 2:], boxes2[None, :, 2:])

    wh = (rb - lt).clamp(min=0)

    inter = wh[..., 0] * wh[..., 1]

    area1 = box_area_xyxy(boxes1)[:, None]
    area2 = box_area_xyxy(boxes2)[None, :]

    union = area1 + area2 - inter + 1e-6

    iou = inter / union

This function measures spatial similarity between predicted and ground-truth boxes. Unlike basic IoU, Generalized IoU penalizes predictions that are far away from ground truth. This improves convergence during transformer-based detection training where predictions are initially random.

COCO Dataset Loader

The dataset loader reads images and annotations and converts them into training samples that OWL-ViT can process.

class CocoDetDataset(Dataset):

    def __init__(self, images_dir, ann_path, category_id_to_idx):

        self.images_dir = images_dir
        self.ann = json.load(open(ann_path))

        self.category_id_to_idx = category_id_to_idx

        self.images = {img["id"]: img for img in self.ann["images"]}

        self.image_ids = list(self.images.keys())

This class organizes annotations by image and allows the training loop to retrieve samples efficiently. Each training sample includes the image, bounding boxes, and labels.

def __getitem__(self, i):

    image_id = self.image_ids[i]

    info = self.images[image_id]

    path = os.path.join(self.images_dir, info["file_name"])

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

This portion loads the image and retrieves annotation metadata.

x1 = x / W
y1 = y / H

x2 = (x + w) / W
y2 = (y + h) / H

Bounding boxes are normalized into the range [0,1]. OWL-ViT predicts normalized boxes, so training requires ground truth boxes in the same format.

OWL-ViT Collator

OWL-ViT requires both images and text descriptions. The collator prepares batches containing both.

class OwlViTCollator:

    def __call__(self, batch):

        images, targets = zip(*batch)

        texts = [self.class_texts for _ in images]

        enc = self.processor(
            text=texts,
            images=list(images),
            return_tensors="pt",
            padding=True
        )

This class converts raw images and class names into tensors suitable for the transformer model. Each image receives the same list of class prompts, allowing OWL-ViT to compare each region against all classes.

Hungarian Matching and Detection Loss

OWL-ViT produces hundreds of candidate detections per image. Hungarian matching determines which predictions correspond to real objects.

class OwlViTDetrLoss(nn.Module):

    def forward(self, logits, pred_boxes, targets):

        prob = logits.softmax(dim=-1)

        pred_boxes_xyxy = box_cxcywh_to_xyxy(pred_boxes)

The model outputs classification logits and predicted boxes. These are converted into probabilities and comparable box formats.

cost_class = -prob[:, tgt_labels]

cost_l1 = torch.cdist(boxes_xyxy, tgt_boxes_xyxy)

cost_giou = -giou

Matching cost is computed using classification similarity, box distance, and overlap quality. Hungarian matching selects the best prediction for each ground truth object.

q_ind, t_ind = linear_sum_assignment(cost)

This ensures each real object is matched to exactly one prediction. Transformer detectors require this matching because predictions are unordered.

loss = classification + L1 + GIoU

Classification loss teaches the model which object category each region belongs to. L1 loss teaches accurate box coordinates. GIoU loss teaches spatial overlap quality.

Model Initialization

OWL-ViT is loaded with pretrained weights.

model = OwlViTForObjectDetection.from_pretrained(
"google/owlvit-base-patch32"
)

Pretrained weights allow the model to start with general visual understanding before adapting to the custom dataset.

Freezing Text Encoder

for n,p in model.named_parameters():

    if "text_model" in n:

        p.requires_grad=False

The text encoder is frozen to preserve language representations. Only the vision and detection components are updated. This stabilizes training and prevents semantic drift.

Training Loop

The training loop feeds batches into the model and updates weights.

outputs = model(**batch)

logits = outputs.logits

pred_boxes = outputs.pred_boxes

The model produces candidate detections for each image.

losses = criterion(logits,pred_boxes,targets)

accelerator.backward(loss)

Loss gradients are computed and propagated backward through the network.

optimizer.step()
optimizer.zero_grad()

Weights are updated after each batch.

Validation Loop

Validation evaluates performance on unseen data.

outputs=model(**batch)

losses=criterion(
outputs.logits,
outputs.pred_boxes,
targets)

Validation loss indicates generalization ability.

Inference Pipeline

During inference, each image or video frame is processed independently.

inputs=processor(
text=[CLASS_NAMES],
images=image,
return_tensors="pt")

The processor converts images and text into tensors.

outputs=model(**inputs)

The model predicts candidate detections.

probs=torch.softmax(logits,dim=-1)

labels=probs.argmax(dim=-1)

scores=probs.max(dim=-1).values

The highest probability class is selected for each candidate region.

Bounding Box Conversion

x1=(cx-bw/2)*W
y1=(cy-bh/2)*H

x2=(cx+bw/2)*W
y2=(cy+bh/2)*H

Normalized coordinates are converted into pixel coordinates.

Confidence Filtering

if score < SCORE_THRESHOLD:
    continue

Low confidence detections are removed.

Non Maximum Suppression

keep=nms(boxes_xyxy,scores,0.45)

NMS removes duplicate detections and keeps only the best bounding box for each object.

Video Processing

Video frames are processed sequentially.

cap=cv2.VideoCapture(INPUT_VIDEO)

Frames are read and passed through the detection pipeline.

detections=detect_frame(frame)

frame=draw(frame,detections)

out.write(frame)

The final annotated frames are saved into an output video.

Complete Training Flow

The complete pipeline consists of:

COCO Dataset
→ Dataset Loader
→ Processor
→ OWL-ViT Model
→ Hungarian Matching
→ Detection Loss
→ Backpropagation
→ Checkpoints
→ Inference
→ Filtering
→ Annotated Output