How would you refactor this pytorch dataset and dataloader to Fastai library

I’m working with this directory:

import torch
from import Dataset, DataLoader
from torchvision import models, transforms
import json
import cv2
import numpy as np

This are my imports

class KeypointsDataset(Dataset):
    def __init__(self, img_dir, data_file):
        self.img_dir = img_dir
        with open(data_file, "r") as f:
   = json.load(f)
        self.transforms = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    def __len__(self):
        return len(
    def __getitem__(self, idx):
        item =[idx]
        img = cv2.imread(f"{self.img_dir}/{item['id']}.png")
        h,w = img.shape[:2]

        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = self.transforms(img)
        kps = np.array(item['kps']).flatten()
        kps = kps.astype(np.float32)

        kps[::2] *= 224.0 / w # Adjust x coordinates
        kps[1::2] *= 224.0 / h # Adjust y coordinates

        return img, kps

This is my dataset where I need to predict the keypoints given an img. The json where I have the keypoints is as follows:

My dataloader is as follow

train_dataset = KeypointsDataset("data/images", "data/data_train.json")
val_dataset = KeypointsDataset("data/images", "data/data_val.json")

train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
train_loader = DataLoader(val_dataset, batch_size=16, shuffle=True)

I’m using a pretrain resnet50 model

model = models.resnet50(pretrained=True)

model.fc = torch.nn.Linear(model.fc.in_features, 14*2) # 14 kps * 2

It’s a regression problem because we want to predict the keypoints

loss_func = torch.nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

This is my training loop

epochs = 20
for epoch in range(epochs):
  for i, (imgs, kps) in enumerate(train_loader):
    imgs =
    kps =

    preds = model(imgs)
    loss = loss_func(preds, kps)

    if i % 10 == 0:
      print(f"Epoch: {epoch}, iter {i}, loss: {loss.item()}")

Anyone could help me in trying to refactor this into Fastai. It would be really good to understand a little bit more about how to use fastai on a img regression problem. Thanks