<gh_stars>1-10
import torch
from torch.utils.data import Dataset
import cv2
import os
import numpy as np

class SqBaseDataset(Dataset):
    def __init__(self):
        pass
        
    def __len__(self):
        return 0
    
    def __getitem__(self, idx):
        dct = {"x": [None], "y": None}
        return dct

    def check_shape(self):
        dct = self.__getitem__(0)
        shapes = {
            "x": [x.shape for x in dct["x"]],
            "y": dct["y"].shape
        }
        return shapes

class TabularDataset(SqBaseDataset):
    def __init__(self, features, targets):
        """
        features: list of numpy.array
        targets: numpy.array
        """
        self.features = features
        self.targets = targets
        
    def __len__(self):
        return (self.targets.shape[0])
    
    def __getitem__(self, idx):
        dct = {"x": [torch.tensor(x[idx, :], dtype=torch.float) for x in self.features],
               "y": torch.tensor(self.targets[idx, :], dtype=torch.float)}
        return dct

class ImgClassifierDataset(SqBaseDataset):
    def __init__(self, imgs, targets, img_dir, img_format, transform=None, C_W_H=True):
        """
        imgs: list or numpy.array
        targets: numpy.array
        img_dir: str
        transform: albumentations or None
        """
        self.imgs = imgs
        self.targets = targets
        self.img_dir = img_dir
        self.img_format = img_format
        self.transform = transform
        self.C_W_H = C_W_H
        
    def __len__(self):
        return (self.targets.shape[0])
    
    def __getitem__(self, idx):

        img_path = os.path.join(self.img_dir, self.imgs[idx] + f".{self.img_format}")
        img = cv2.imread(img_path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        
        if self.transform is not None:
            res = self.transform(image=img)
            img = res["image"]
                
        img = img.transpose(2,0,1).astype(np.float32) if self.C_W_H else img.astype(np.float32)

        dct = {"x": [torch.tensor(img, dtype=torch.float)],
               "y": torch.tensor(self.targets[idx, :], dtype=torch.float)}
        return dct
import torch
from torch.utils.data import Dataset
import cv2
import os
import numpy as np

class SqBaseDataset(Dataset):
    def __init__(self):
        pass
        
    def __len__(self):
        return 0
    
    def __getitem__(self, idx):
        dct = {"x": [None], "y": None}
        return dct

    def check_shape(self):
        dct = self.__getitem__(0)
        shapes = {
            "x": [x.shape for x in dct["x"]],
            "y": dct["y"].shape
        }
        return shapes

class TabularDataset(SqBaseDataset):
    def __init__(self, features, targets):
        """
        features: list of numpy.array
        targets: numpy.array
        """
        self.features = features
        self.targets = targets
        
    def __len__(self):
        return (self.targets.shape[0])
    
    def __getitem__(self, idx):
        dct = {"x": [torch.tensor(x[idx, :], dtype=torch.float) for x in self.features],
               "y": torch.tensor(self.targets[idx, :], dtype=torch.float)}
        return dct

class ImgClassifierDataset(SqBaseDataset):
    def __init__(self, imgs, targets, img_dir, img_format, transform=None, C_W_H=True):
        """
        imgs: list or numpy.array
        targets: numpy.array
        img_dir: str
        transform: albumentations or None
        """
        self.imgs = imgs
        self.targets = targets
        self.img_dir = img_dir
        self.img_format = img_format
        self.transform = transform
        self.C_W_H = C_W_H
        
    def __len__(self):
        return (self.targets.shape[0])
    
    def __getitem__(self, idx):

        img_path = os.path.join(self.img_dir, self.imgs[idx] + f".{self.img_format}")
        img = cv2.imread(img_path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        
        if self.transform is not None:
            res = self.transform(image=img)
            img = res["image"]
                
        img = img.transpose(2,0,1).astype(np.float32) if self.C_W_H else img.astype(np.float32)

        dct = {"x": [torch.tensor(img, dtype=torch.float)],
               "y": torch.tensor(self.targets[idx, :], dtype=torch.float)}
        return dct
features: list of numpy.array         targets: numpy.array imgs: list or numpy.array         targets: numpy.array         img_dir: str         transform: albumentations or None