CarDataset2.py

import cv2
import torch
import numpy as np
from torch.utils.data import Dataset



class CarDataset(Dataset):
    def __init__(self, data, IMG_WIDTH = 288, IMG_HEIGHT = 224, HEATMAP_RADIUS = 8, transform=None):
        self.transform = transform
        self.images = []
        self.heatmaps = []
        
        for img_path, label_path in data:
            img = cv2.imread(img_path)
            original_height, original_width, *_= img.shape
            img = cv2.resize(img, (IMG_WIDTH, IMG_HEIGHT))
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            
            heatmap = np.zeros((IMG_HEIGHT, IMG_WIDTH), dtype=np.float32)
            with open(label_path, "r") as f:
                points = [list(map(int, line.split())) for line in f.readlines()]
            
            for x, y in points:
                x = int(x * IMG_WIDTH / original_width)
                y = int(y * IMG_HEIGHT / original_height)
                cv2.circle(heatmap, (x, y), HEATMAP_RADIUS, 1, -1)
            
            if self.transform:
                img = self.transform(img)
            
            self.images.append(img)
            self.heatmaps.append(torch.tensor(heatmap, dtype=torch.float32))
            
        self.images = torch.stack(self.images)
        self.heatmaps = torch.stack(self.heatmaps)
    
    def __len__(self):
        return len(self.images)
    
    def __getitem__(self, idx):
        return self.images[idx], self.heatmaps[idx]