Run in Google Colab

View source on GitHub

Download notebook

Data Loading with FiftyOneTorchDataset

This recipe introduces FiftyOneTorchDataset, which lets you turn any FiftyOne dataset or view directly into a PyTorch-compatible dataset — no data copying, no format conversion. Specifically, it covers:

• Loading a dataset from the Dataset Zoo and creating a view

• Converting a FiftyOne view to a FiftyOneTorchDataset with a custom get_item function

• Visualizing samples from the dataset

• Wrapping the dataset in a torch.utils.data.DataLoader for use in a training loop

API references: FiftyOneTorchDataset · GetItem

Setup

If you haven’t already, install FiftyOne:

!pip install fiftyone

In this tutorial, we’ll use PyTorch for working with tensors and inspecting sample data. To follow along, you’ll need to install torch and torchvision, if necessary:

!pip install torch torchvision

This recipe requires a helper file, utils.py, which contains reusable functions for building get_item methods and creating dataloaders. The following cell downloads it into your working directory.

import urllib.request

# utils.py is shared across the torch-dataset-examples notebooks
url = "https://cdn.voxel51.com/tutorials_torch_dataset_examples/notebook_the_cache_field_names_argument/utils.py"
urllib.request.urlretrieve(url, "utils.py")

Import Libraries

import fiftyone as fo
import fiftyone.zoo as foz
from fiftyone.utils.torch import FiftyOneTorchDataset

import utils

import torch
from torch.utils.data import DataLoader
import numpy as np
import torchvision.transforms.v2 as transforms
from torchvision import tv_tensors
import matplotlib.pyplot as plt
import matplotlib.patches as plt_patches
from PIL import Image

torch.multiprocessing.set_start_method("forkserver")
torch.multiprocessing.set_forkserver_preload(["torch", "fiftyone"])

Load Dataset

dataset = foz.load_zoo_dataset("quickstart", overwrite=True)

# make sure it's persistent
print(dataset.persistent)

# if it's not, set this to True
if not dataset.persistent:
    dataset.persistent = True

Curate Your View

One of the key advantages of FiftyOneTorchDataset is that you can pass any FiftyOne view — not just a full dataset. Here we take a 100-sample subset, but any filter, sort, or tag expression works equally well.

some_interesting_view = dataset.take(100)

Converting to a FiftyOneTorchDataset

To convert a view to a FiftyOneTorchDataset, call .to_torch(get_item) on any dataset or view. The get_item argument is a callable that receives a fiftyone.core.sample.Sample and returns whatever your model expects.

To best understand what’s happening, start with the identity function — this shows you exactly what FiftyOneTorchDataset hands to your get_item callable:

# to best understand what's happening, let's first pass the identity function
def get_item_identity(x):
    return x

torch_dataset = some_interesting_view.to_torch(get_item_identity)

result = torch_dataset[0]
print(type(result))
print(result["id"])
print(result["filepath"])

The get_item callable can return anything — here’s a minimal example that returns just the sample ID:

def simple_get_item(sample):
    return sample["id"]

torch_dataset = some_interesting_view.to_torch(simple_get_item)

print(torch_dataset[0])

# FiftyOneTorchDataset iterates in the same order as the view
assert [res for res in torch_dataset] == some_interesting_view.values("id")

Writing a Real get_item

Here’s a detection-ready get_item that converts FiftyOne bounding boxes to tv_tensors.BoundingBoxes and applies augmentations:

augmentations = transforms.Compose([
    transforms.CenterCrop(512),
    transforms.ClampBoundingBoxes(),
])


def get_item(sample):
    image = Image.open(sample["filepath"])
    og_wh = np.array([image.width, image.height])
    image = tv_tensors.Image(image)

    detections = sample["ground_truth.detections"]
    if detections is None:
        detections = []
    detections_tensor = (
        torch.tensor([d["bounding_box"] for d in detections])
        if len(detections) > 0
        else torch.zeros((0, 4))
    )

    res = {
        "box": tv_tensors.BoundingBoxes(
            detections_tensor * torch.tensor([*og_wh, *og_wh]),
            format=tv_tensors.BoundingBoxFormat("XYWH"),
            canvas_size=image.shape[-2:],
        ),
        "label": [d["label"] for d in detections],
        "id": sample["id"],
    }
    image, res = augmentations(image, res)
    return image, res

Visualizing Samples

This is also a good place to debug your get_item function before hooking it into a training loop:

torch_dataset = some_interesting_view.to_torch(get_item)

# run this a couple of times to look through samples in the dataset
random_index = np.random.randint(0, len(torch_dataset))
image, res = torch_dataset[random_index]
plt.title(res["id"])
plt.imshow(image.permute(1, 2, 0).numpy())
axes = plt.gca()
for b, l in zip(res["box"], res["label"]):
    rect = plt_patches.Rectangle(
        (b[0], b[1]), b[2], b[3], linewidth=1, edgecolor="r", facecolor="none"
    )
    axes.add_patch(rect)
    axes.annotate(l, rect.get_xy())
plt.show()

Creating a DataLoader

FiftyOneTorchDataset is fully compatible with torch.utils.data.DataLoader, including multi-worker loading. The only required addition is worker_init_fn=FiftyOneTorchDataset.worker_init, which lets each worker process open its own FiftyOne database connection:

def simple_collate_fn(batch):
    return tuple(zip(*batch))

dataloader = torch.utils.data.DataLoader(
    torch_dataset,
    batch_size=5,
    shuffle=True,
    num_workers=2,
    worker_init_fn=FiftyOneTorchDataset.worker_init,
    collate_fn=simple_collate_fn,
)

The cells below use utils.create_dataloader_simple(), which wraps this same pattern:

# Create a fresh dataset object — iterating torch_dataset above opened a DB connection
# that cannot be pickled for multiprocessing workers.
# utils.get_item_quickstart is the same get_item defined above.
torch_dataset = some_interesting_view.to_torch(utils.get_item_quickstart)
dataloader = utils.create_dataloader_simple(torch_dataset)

ids_seen = utils.ids_in_dataloader(dataloader)

from collections import Counter

ids_with_counts = Counter(ids_seen)
assert set(ids_with_counts.keys()) == set(some_interesting_view.values("id"))
assert all(v == 1 for v in ids_with_counts.values())

# visualizing results, run this a couple of times to see different batches
iterator = iter(dataloader)

image, res = next(iterator)

fig, axes = plt.subplots(1, len(image), figsize=(4 * len(image), 3))
for i, img in enumerate(image):
    axes[i].set_title(res[i]["id"])
    axes[i].imshow(img.permute(1, 2, 0).numpy())
    for b, l in zip(res[i]["box"], res[i]["label"]):
        rect = plt_patches.Rectangle(
            (b[0], b[1]), b[2], b[3], linewidth=1, edgecolor="r", facecolor="none"
        )
        axes[i].add_patch(rect)
        axes[i].annotate(l, rect.get_xy())
plt.show()