develop #4
|
@ -18,7 +18,7 @@ from torch.utils.checkpoint import checkpoint
|
||||||
|
|
||||||
class aiuNNDataset(torch.utils.data.Dataset):
|
class aiuNNDataset(torch.utils.data.Dataset):
|
||||||
def __init__(self, parquet_path):
|
def __init__(self, parquet_path):
|
||||||
self.df = pd.read_parquet(parquet_path, columns=['image_512', 'image_1024']).head(2000)
|
self.df = pd.read_parquet(parquet_path, columns=['image_512', 'image_1024']).head(10000)
|
||||||
self.augmentation = Compose([
|
self.augmentation = Compose([
|
||||||
RandomBrightnessContrast(p=0.5),
|
RandomBrightnessContrast(p=0.5),
|
||||||
HorizontalFlip(p=0.5),
|
HorizontalFlip(p=0.5),
|
||||||
|
|
Loading…
Reference in New Issue