DeForge-AI/dataset.py at main · tbtiberiu/DeForge-AI · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
import torch
import torch.nn as nn
from torch.utils.data import Dataset
from torchvision.transforms import InterpolationMode, v2
from torchvision.transforms.v2 import functional as tvf


class AIGIBenchDataset(Dataset):
    def __init__(self, hf_data, transform=None):
        self.hf_data = hf_data
        self.transform = transform

    def __len__(self):
        return len(self.hf_data)

    def __getitem__(self, idx):
        item = self.hf_data[idx]
        image = item['image'].convert('RGB')
        label = item['label']  # 0 for Real, 1 for Fake

        if self.transform:
            image = self.transform(image)

        return image, torch.tensor(label, dtype=torch.float32)


class DownUpResize(nn.Module):
    def __init__(self, scale_range=(0.55, 0.9)):
        super().__init__()
        self.scale_range = scale_range
        self.interpolations = [
            InterpolationMode.BILINEAR,
            InterpolationMode.BICUBIC,
            InterpolationMode.NEAREST_EXACT,
        ]

    def forward(self, image):
        height, width = tvf.get_size(image)
        scale = torch.empty(1).uniform_(*self.scale_range).item()
        down_height = max(32, int(round(height * scale)))
        down_width = max(32, int(round(width * scale)))

        down_interpolation = self.interpolations[
            torch.randint(len(self.interpolations), size=()).item()
        ]
        up_interpolation = self.interpolations[
            torch.randint(len(self.interpolations), size=()).item()
        ]

        image = tvf.resize(
            image,
            [down_height, down_width],
            interpolation=down_interpolation,
            antialias=False,
        )
        return tvf.resize(
            image,
            [height, width],
            interpolation=up_interpolation,
            antialias=False,
        )


class GaussianNoise(nn.Module):
    def __init__(self, sigma_range=(0.002, 0.02)):
        super().__init__()
        self.sigma_range = sigma_range

    def forward(self, image):
        sigma = torch.empty(1).uniform_(*self.sigma_range).item()
        noise = torch.randn_like(image) * sigma
        return (image + noise).clamp(0.0, 1.0)


def get_train_transforms(size=256):
    resize_size = max(int(round(size * 1.15)), size)
    return v2.Compose(
        [
            v2.ToImage(),
            v2.Resize(resize_size, antialias=False),
            v2.RandomCrop(size),
            v2.RandomHorizontalFlip(),
            v2.RandomApply([DownUpResize()], p=0.1),
            v2.RandomApply(
                [
                    v2.ColorJitter(
                        brightness=0.2, contrast=0.2, saturation=0.1, hue=0.02
                    )
                ],
                p=0.25,
            ),
            v2.RandomApply([v2.GaussianBlur(kernel_size=5, sigma=(0.1, 1.2))], p=0.25),
            v2.RandomAdjustSharpness(sharpness_factor=1.5, p=0.15),
            v2.RandomApply([v2.JPEG(quality=(40, 95))], p=0.5),
            v2.ToDtype(torch.float32, scale=True),
            v2.RandomApply([GaussianNoise()], p=0.25),
            v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ]
    )


def get_val_transforms(size=256):
    resize_size = max(int(round(size * 1.15)), size)
    return v2.Compose(
        [
            v2.ToImage(),
            v2.Resize(resize_size, antialias=False),
            v2.CenterCrop(size),
            v2.ToDtype(torch.float32, scale=True),
            v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ]
    )