| 125 | self.model = bmshj2018_factorized(quality=3, pretrained=True).eval().cuda() |
| 126 | |
| 127 | class Attacker(nn.Module): |
| 128 | def __init__(self, attack_types, flag=True): |
| 129 | super().__init__() |
| 130 | for attack_type in filter(lambda x : x != 'JPEGCompress', attack_types): |
| 131 | setattr(self, attack_type, globals()[f'{attack_type}Attack']()) |
| 132 | |
| 133 | if 'JPEGCompress' in attack_types: |
| 134 | self.JPEGCompress = JPEGCompressAttack(flag) |
| 135 | |
| 136 | def forward(self, image, attack_type): |
| 137 | return getattr(self, attack_type)(image) |