kalpkanungo/SceneGraphNet

1import json2import os3import cv24import torch5from torch.utils.data import Dataset6 7class RelationshipDataset(Dataset):8    def __init__(self, image_dir, label_path):9        self.image_dir = image_dir10 11        with open(label_path) as f:12            self.data = json.load(f)13 14    def __len__(self):15        return len(self.data)16 17    def __getitem__(self, idx):18        item = self.data[idx]19 20        img_path = os.path.join(self.image_dir, item["image"])21        image = cv2.imread(img_path)22        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)23        image = image / 255.024        image = (image - 0.5) / 0.525 26        image = torch.tensor(image, dtype=torch.float32).permute(2, 0, 1)27 28        label = torch.tensor(item["label"], dtype=torch.long)29 30        return image, label