I’m working with this directory:
import torch
from torch.utils.data 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:
self.data = json.load(f)
self.transforms = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
item = self.data[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 = imgs.to(device)
kps = kps.to(device)
optimizer.zero_grad()
preds = model(imgs)
loss = loss_func(preds, kps)
loss.backward()
optimizer.step()
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