import torch import torch.nn as nn import torch.optim as optim from torchvision import transforms, utils from torchvision.models import vgg16 import numpy as np from PIL import Image import io import os # ------------------------------- # 1. Differentiable JPEG simulator # ------------------------------- class DiffJPEG(nn.Module): def __init__(self, quality=85): super().__init__() # Standard JPEG quality scaling if quality < 50: scale = 5000 / quality else: scale = 200 - 2 * quality # Standard JPEG luminance & chrominance quantization tables base_lum = torch.tensor([ [16, 11, 10, 16, 24, 40, 51, 61], [12, 12, 14, 19, 26, 58, 60, 55], [14, 13, 16, 24, 40, 57, 69, 56], [14, 17, 22, 29, 51, 87, 80, 62], [18, 22, 37, 56, 68, 109, 103, 77], [24, 35, 55, 64, 81, 104, 113, 92], [49, 64, 78, 87, 103, 121, 120, 101], [72, 92, 95, 98, 112, 100, 103, 99] ]).float() self.Q_table_lum = torch.clamp((base_lum * scale + 50) / 100, 1.0, 255.0) base_chrom = torch.tensor([ [17, 18, 24, 47, 99, 99, 99, 99], [18, 21, 26, 66, 99, 99, 99, 99], [24, 26, 56, 99, 99, 99, 99, 99], [47, 66, 99, 99, 99, 99, 99, 99], [99, 99, 99, 99, 99, 99, 99, 99], [99, 99, 99, 99, 99, 99, 99, 99], [99, 99, 99, 99, 99, 99, 99, 99], [99, 99, 99, 99, 99, 99, 99, 99] ]).float() self.Q_table_chrom = torch.clamp((base_chrom * scale + 50) / 100, 1.0, 255.0) def rgb_to_ycbcr(self, x): # x: [B,3,H,W] in [0,1] # JPEG YCbCr: Y: 0-255, CbCr: 0-255 (offset 128) # But we work in 0-1 range for simplicity, just shift r, g, b = x[:,0]*255, x[:,1]*255, x[:,2]*255 y = 0.299*r + 0.587*g + 0.114*b cb = -0.1687*r - 0.3313*g + 0.5*b + 128 cr = 0.5*r - 0.4187*g - 0.0813*b + 128 return torch.stack([y, cb, cr], dim=1) def ycbcr_to_rgb(self, ycbcr): y, cb, cr = ycbcr[:,0], ycbcr[:,1], ycbcr[:,2] cb = cb - 128 cr = cr - 128 r = y + 1.402*cr g = y - 0.34414*cb - 0.71414*cr b = y + 1.772*cb return torch.stack([r, g, b], dim=1) / 255.0 def dct_8x8(self, x): # x: [B, C, H, W], values typically shifted to -128 to 127 B, C, H, W = x.shape x = x - 128 x = x.view(B, C, H//8, 8, W//8, 8).permute(0,1,2,4,3,5) dct_mat = self._dct_matrix(8, device=x.device) # DCT-II: X = M * x * M^T x_dct = torch.matmul(torch.matmul(dct_mat, x), dct_mat.t()) return x_dct # [B, C, H//8, W//8, 8, 8] def idct_8x8(self, x): # x: [B, C, H//8, W//8, 8, 8] B, C, H_b, W_b, _, _ = x.shape dct_mat = self._dct_matrix(8, device=x.device) # IDCT: x = M^T * X * M x_idct = torch.matmul(torch.matmul(dct_mat.t(), x), dct_mat) x_idct = x_idct.permute(0,1,2,4,3,5).reshape(B, C, H_b*8, W_b*8) return x_idct + 128 def _dct_matrix(self, N, device): n = torch.arange(N, device=device).float() k = torch.arange(N, device=device).float().unsqueeze(1) mat = torch.cos(np.pi / N * (n + 0.5) * k) mat[0] = mat[0] * np.sqrt(1/N) mat[1:] = mat[1:] * np.sqrt(2/N) return mat def quantize(self, coeff, table): # Straight‑through estimator t = table.view(1, 1, 1, 1, 8, 8) q = coeff / t hard = torch.round(q) soft = q + (hard - q).detach() bit_proxy = torch.abs(soft).sum() / coeff.numel() return soft * t, bit_proxy def forward(self, x): ycbcr = self.rgb_to_ycbcr(x) y, cb, cr = ycbcr[:,0:1], ycbcr[:,1:2], ycbcr[:,2:3] # Subsample chroma cb_small = nn.functional.avg_pool2d(cb, 2) cr_small = nn.functional.avg_pool2d(cr, 2) # DCT y_dct = self.dct_8x8(y) cb_dct = self.dct_8x8(cb_small) cr_dct = self.dct_8x8(cr_small) # Quantize y_q, bit_y = self.quantize(y_dct, self.Q_table_lum.to(x.device)) cb_q, bit_cb = self.quantize(cb_dct, self.Q_table_chrom.to(x.device)) cr_q, bit_cr = self.quantize(cr_dct, self.Q_table_chrom.to(x.device)) # IDCT y_rec = self.idct_8x8(y_q) cb_rec = self.idct_8x8(cb_q) cr_rec = self.idct_8x8(cr_q) # Upsample cb_rec = nn.functional.interpolate(cb_rec, scale_factor=2, mode='bilinear', align_corners=False) cr_rec = nn.functional.interpolate(cr_rec, scale_factor=2, mode='bilinear', align_corners=False) ycbcr_rec = torch.cat([y_rec, cb_rec, cr_rec], dim=1) rgb_rec = self.ycbcr_to_rgb(ycbcr_rec) return torch.clamp(rgb_rec, 0, 1), bit_y + bit_cb + bit_cr # ---------------------------------- # 2. Perceptual loss (MS‑SSIM + L2) # ---------------------------------- class MS_SSIM_L2(nn.Module): def __init__(self): super().__init__() # Use torchmetrics if available, else simple L2 + gradient penalty def forward(self, pred, target): # Simplified: L2 + edge loss (Sobel) l2 = nn.functional.mse_loss(pred, target) sobel_x = torch.tensor([[-1,0,1],[-2,0,2],[-1,0,1]], device=pred.device).float().view(1,1,3,3) sobel_y = torch.tensor([[-1,-2,-1],[0,0,0],[1,2,1]], device=pred.device).float().view(1,1,3,3) def edges(img): img_gray = 0.299*img[:,0:1] + 0.587*img[:,1:2] + 0.114*img[:,2:3] # Use groups=1 for single channel gray gx = nn.functional.conv2d(img_gray, sobel_x, padding=1) gy = nn.functional.conv2d(img_gray, sobel_y, padding=1) return torch.sqrt(gx**2 + gy**2 + 1e-6) loss_edge = nn.functional.l1_loss(edges(pred), edges(target)) return 10.0 * l2 + 0.1 * loss_edge # ---------------------------------- # 3. Optimization loop # ---------------------------------- def preprocess_for_jpeg(image_path, output_path, num_iter=300, lr=0.01, weight_bit=0.0001): # Load image img_pil = Image.open(image_path).convert('RGB') transform = transforms.Compose([ transforms.Resize((256,256)), # ensure divisible by 8 transforms.ToTensor() ]) img_tensor = transform(img_pil).unsqueeze(0) # [1,3,H,W] # Make trainable copy preprocessed = img_tensor.clone().detach().requires_grad_(True) optimizer = optim.Adam([preprocessed], lr=lr) jpeg_sim = DiffJPEG(quality=90) loss_fn = MS_SSIM_L2() print("Starting optimisation...") for i in range(num_iter): optimizer.zero_grad() rec, bit_estimate = jpeg_sim(preprocessed) mse = nn.functional.mse_loss(rec, img_tensor) loss_id = loss_fn(rec, img_tensor) loss_bit = bit_estimate loss = loss_id + weight_bit * loss_bit loss.backward() if i % 50 == 0: gnorm = preprocessed.grad.norm().item() if preprocessed.grad is not None else 0.0 print(f"iter {i:4d} | loss: {loss.item():.4f} | mse: {mse.item():.4f} | gnorm: {gnorm:.6f}") optimizer.step() # Clamp to [0,1] with torch.no_grad(): preprocessed.clamp_(0, 1) # Save optimized image preprocessed_img = preprocessed.detach().squeeze(0).permute(1,2,0).cpu().numpy() preprocessed_img = (preprocessed_img * 255).astype(np.uint8) Image.fromarray(preprocessed_img).save(output_path) print(f"Optimised image saved to {output_path}") # Compare JPEG sizes (actual PIL JPEG) def jpeg_size(img_path, quality=85): img = Image.open(img_path) buf = io.BytesIO() img.save(buf, format='JPEG', quality=quality) return buf.tell() orig_size = jpeg_size(image_path, quality=85) new_size = jpeg_size(output_path, quality=85) print(f"Original JPEG size: {orig_size/1024:.1f} KB") print(f"Optimised JPEG size: {new_size/1024:.1f} KB") print(f"Reduction: {100*(1 - new_size/orig_size):.1f}%") # ---------------------------------- # 4. Run the experiment # ---------------------------------- if __name__ == "__main__": # Replace with your image path preprocess_for_jpeg("input.jpg", "output_preprocessed.jpg", num_iter=300, weight_bit=0.0001)