| 41 | |
| 42 | # this function is implemented based on Diff-JPEG (https://github.com/necla-ml/Diff-JPEG). |
| 43 | def jpeg_coding_cv2(image_rgb, jpeg_quality): |
| 44 | B, _, _, _ = image_rgb.shape |
| 45 | image_rgb_jpeg = [] |
| 46 | for index in range(B): |
| 47 | encode_parameters = (int(cv2.IMWRITE_JPEG_QUALITY), int(jpeg_quality[index].item())) |
| 48 | _, encoding = cv2.imencode('.jpeg', image_rgb[index].flip(0).permute(1, 2, 0).numpy(), encode_parameters) |
| 49 | image_rgb_jpeg.append(torch.from_numpy(cv2.imdecode(encoding, 1)).permute(2, 0, 1).flip(0)) |
| 50 | image_rgb_jpeg = torch.stack(image_rgb_jpeg, dim=0) |
| 51 | return image_rgb_jpeg |
| 52 | |
| 53 | class JPEGCompressAttack(nn.Module): |
| 54 | def __init__(self, flag=True): |