fo_img_cls.py

from typing import NamedTuple
from typing import List
from typing import Callable, Optional

import fiftyone as fo
import fiftyone.zoo as foz

from pathlib import Path
import json
import numpy as np

import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import models
from torchvision.io import read_image


class NativeSample(NamedTuple):
    img: Path
    label: str


class ImagenetTorch(Dataset):
    def __init__(
        self, samples: List[NativeSample], transforms: Optional[Callable] = None
    ):
        self.samples = samples
        self.transforms = transforms

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, index):
        label = self.samples[index].label
        img_path = self.samples[index].img
        image = read_image(str(img_path))
        # image = image.float()

        if self.transforms:
            image = self.transforms(image)

        return image, label


# loading the zoo dataset
dataset = foz.load_zoo_dataset("imagenet-sample")

all_samples: List[NativeSample] = []
for sample in dataset.iter_samples():
    all_samples.append(
        NativeSample(
            img=sample.filepath,
            label=sample.ground_truth.label,
        )
    )

torch_ds = ImagenetTorch(all_samples)
img_net_dataloader = DataLoader(torch_ds, batch_size=1)

# initialize the ResNet 50 model with Imagenet weights
model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
model.eval() # set it into eval mode

# download the "idx_to_labels.json" file from this link [https://gist.github.com/mohan-aditya05/783a56ba62e0605824489dea552531d5]
with open("idx_to_labels.json", "rb") as f:
    idx_to_label = json.load(f)

preprocess = models.ResNet50_Weights.IMAGENET1K_V1.transforms() # transforms for resnet50

preds = []
confs = []
it = iter(img_net_dataloader)
for i in range(len(img_net_dataloader)):
    ipt = next(it)
    img, label = ipt
    img = preprocess(img)

    out = model(img)
    pred = torch.argmax(out, 1)
    pred = int(pred[0])

    out = out.detach().numpy()
    odds = np.exp(out)
    conf = np.max(odds, axis=1) / np.sum(odds, axis=1)

    preds.append(pred)
    confs.append(conf)

for i, sample in enumerate(dataset.iter_samples()):
    # assigning predictions back to samples with the predicted labels and confidence scores
    sample["predictions"] = fo.Classification(
        label=idx_to_label[str(preds[i])].split(",")[0],
        confidence=confs[i]*100,
    )
    sample.save()

if fo.dataset_exists(name="imagenet-preds"):
    fo.delete_dataset(name="imagenet-preds")

# creating a separate dataset with the name "imagenet-preds"  
new_ds = fo.Dataset(name="imagenet-preds")
new_ds.merge_samples(dataset)
new_ds.persistent = True # since we want the dataset to exist after the process is ended we set persistent as True
new_ds.save()