image_to_image.models.pix2pix

  1"""
  2Module to define a Pix2Pix Model. 
  3A UNet CNN Generator combined with a small generative loss.
  4
  5Functions:
  6- unet_down_block
  7- unet_up_block
  8
  9Classes:
 10- MMC
 11- UNetGenerator
 12- Discriminator
 13- Pix2Pix
 14
 15By Tobia Ippolito
 16"""
 17# ---------------------------
 18#        > Imports <
 19# ---------------------------
 20import torch
 21import torch.nn as nn
 22import torch.nn.functional as F
 23from torch.amp import autocast
 24
 25
 26
 27# ---------------------------
 28#       > Generator <
 29# ---------------------------
 30class MMC(nn.Module):  # MinMaxClamping
 31    """
 32    Min-Max Clamping Module.
 33
 34    Clamps input tensor values between a specified minimum and maximum.
 35
 36    Parameter:
 37    - min (float): 
 38        Minimum allowed value (default=0.0).
 39    - max (float): 
 40        Maximum allowed value (default=1.0).
 41
 42    Usage:
 43    - Can be used at the output layer of a generator to ensure predictions remain in a valid range.
 44    """
 45    def __init__(self, min=0.0, max=1.0):
 46        super().__init__()
 47        self.min = min
 48        self.max = max
 49
 50    def forward(self, x):
 51        """
 52        Forward pass.
 53
 54        Parameter:
 55        - x (torch.tensor): 
 56            Input tensor.
 57
 58        Returns:
 59        - torch.tensor: Clamped tensor with values between `min` and `max`.
 60        """
 61        return torch.clamp(x, self.min, self.max)
 62
 63def unet_down_block(in_channels=1, out_channels=1, normalize=True):
 64    """
 65    Creates a U-Net downsampling block.
 66
 67    Parameter:
 68    - in_channels (int): 
 69        Number of input channels.
 70    - out_channels (int): 
 71        Number of output channels.
 72    - normalize (bool): 
 73        Whether to apply instance normalization.
 74
 75    Returns:
 76    - nn.Sequential: Convolutional downsampling block with LeakyReLU activation.
 77    """
 78    layers = [nn.Conv2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1, bias=False)]
 79    if normalize:
 80        # layers += [nn.BatchNorm2d(out_channels)]
 81        layers += [nn.InstanceNorm2d(out_channels, affine=True)]
 82    layers += [nn.LeakyReLU(0.2, inplace=True)]
 83    return nn.Sequential(*layers)
 84
 85def unet_up_block(in_channels=1, out_channels=1, dropout=0.0):
 86    """
 87    Creates a U-Net upsampling block.
 88
 89    Parameter:
 90    - in_channels (int): 
 91        Number of input channels.
 92    - out_channels (int): 
 93        Number of output channels.
 94    - dropout (float): 
 95        Dropout probability (default=0).
 96
 97    Returns:
 98    - nn.Sequential: Transposed convolutional block with ReLU activation and optional dropout.
 99    """
100    layers = [
101        nn.ConvTranspose2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1, bias=False),
102        # nn.BatchNorm2d(out_channels),
103        nn.InstanceNorm2d(out_channels, affine=True),
104        nn.ReLU(inplace=True)
105    ]
106    
107    if dropout:
108        layers += [nn.Dropout(dropout)]
109    return nn.Sequential(*layers)
110
111class UNetGenerator(nn.Module):
112    """
113    U-Net Generator for image-to-image translation.
114
115    Architecture:
116    - 8 downsampling blocks (encoder)
117    - 8 upsampling blocks (decoder) with skip connections
118    - Sigmoid activation at output for [0,1] pixel normalization
119
120    Parameter:
121    - input_channels (int): 
122        Number of input image channels.
123    - output_channels (int): 
124        Number of output image channels.
125    - hidden_channels (int): 
126        Base hidden channels for first layer.
127    """
128    def __init__(self, input_channels=1, output_channels=1, hidden_channels=64):
129        super().__init__()
130        # Encoder
131        self.down1 = unet_down_block(input_channels, hidden_channels, normalize=False) # 128
132        self.down2 = unet_down_block(hidden_channels, hidden_channels*2)                    # 64
133        self.down3 = unet_down_block(hidden_channels*2, hidden_channels*4)                  # 32
134        self.down4 = unet_down_block(hidden_channels*4, hidden_channels*8)                  # 16
135        self.down5 = unet_down_block(hidden_channels*8, hidden_channels*8)                  # 8
136        self.down6 = unet_down_block(hidden_channels*8, hidden_channels*8)                  # 4
137        self.down7 = unet_down_block(hidden_channels*8, hidden_channels*8)                  # 2
138        self.down8 = unet_down_block(hidden_channels*8, hidden_channels*8, normalize=False) # 1
139
140        # Decoder
141        self.up1 = unet_up_block(hidden_channels*8, hidden_channels*8, dropout=0.5)
142        self.up2 = unet_up_block(hidden_channels*16, hidden_channels*8, dropout=0.5)
143        self.up3 = unet_up_block(hidden_channels*16, hidden_channels*8, dropout=0.5)
144        self.up4 = unet_up_block(hidden_channels*16, hidden_channels*8)
145        self.up5 = unet_up_block(hidden_channels*16, hidden_channels*4)
146        self.up6 = unet_up_block(hidden_channels*8, hidden_channels*2)
147        self.up7 = unet_up_block(hidden_channels*4, hidden_channels)
148        self.up8 = nn.Sequential(
149            nn.ConvTranspose2d(hidden_channels*2, output_channels, 4, 2, 1),
150            # nn.Tanh()
151            ## MMC(min=0.0, max=1.0)
152            nn.Sigmoid()
153        )
154
155    def forward(self, x):
156        """
157        Forward pass of the U-Net generator.
158
159        Parameter:
160        - x (torch.tensor): 
161            Input image tensor (batch_size, input_channels, H, W).
162
163        Returns:
164        - torch.tensor: Generated output tensor (batch_size, output_channels, H, W).
165        """
166        d1 = self.down1(x)
167        d2 = self.down2(d1)
168        d3 = self.down3(d2)
169        d4 = self.down4(d3)
170        d5 = self.down5(d4)
171        d6 = self.down6(d5)
172        d7 = self.down7(d6)
173        d8 = self.down8(d7)
174
175        u1 = self.up1(d8)
176        u2 = self.up2(torch.cat([u1, d7], 1))
177        u3 = self.up3(torch.cat([u2, d6], 1))
178        u4 = self.up4(torch.cat([u3, d5], 1))
179        u5 = self.up5(torch.cat([u4, d4], 1))
180        u6 = self.up6(torch.cat([u5, d3], 1))
181        u7 = self.up7(torch.cat([u6, d2], 1))
182        u8 = self.up8(torch.cat([u7, d1], 1))
183
184        return u8
185
186
187
188# ---------------------------
189#      > Discriminator <
190# ---------------------------
191# PatchGAN
192class Discriminator(nn.Module):
193    """
194    PatchGAN Discriminator for Pix2Pix GAN.
195
196    Parameter:
197    - input_channels (int): 
198        Number of input channels (typically input + target channels concatenated).
199    - hidden_channels (int): 
200        Base hidden channels for first layer.
201
202    Architecture:
203    - 5 convolutional blocks with LeakyReLU and batch normalization.
204    - Outputs a 2D patch map of predictions.
205    """
206    def __init__(self, input_channels=6, hidden_channels=64):
207        """
208        Initializes a PatchGAN discriminator.
209
210        The discriminator evaluates input-target image pairs to determine
211        if they are real or generated (fake). It progressively downsamples
212        the spatial dimensions while increasing the number of feature channels.
213
214        Parameters:
215        - input_channels (int): 
216            Number of input channels, typically input + target concatenated (default=6).
217        - hidden_channels (int): 
218            Number of channels in the first convolutional layer; doubled in subsequent layers (default=64).
219        """
220        super().__init__()
221        self.model = nn.Sequential(
222            nn.Conv2d(input_channels, hidden_channels, kernel_size=4, stride=2, padding=1),
223            nn.LeakyReLU(0.2, inplace=True),
224
225            nn.Conv2d(hidden_channels, hidden_channels*2, kernel_size=4, stride=2, padding=1, bias=False),
226            nn.BatchNorm2d(hidden_channels*2),
227            nn.LeakyReLU(0.2, inplace=True),
228
229            nn.Conv2d(hidden_channels*2, hidden_channels*4, kernel_size=4, stride=2, padding=1, bias=False),
230            nn.BatchNorm2d(hidden_channels*4),
231            nn.LeakyReLU(0.2, inplace=True),
232
233            nn.Conv2d(hidden_channels*4, hidden_channels*8, kernel_size=4, stride=1, padding=1, bias=False),
234            nn.BatchNorm2d(hidden_channels*8),
235            nn.LeakyReLU(0.2, inplace=True),
236
237            nn.Conv2d(hidden_channels*8, 1, kernel_size=4, stride=1, padding=1)
238        )
239
240    def forward(self, x, y):
241        """
242        Forward pass of the discriminator.
243
244        Parameter:
245        - x (torch.tensor): 
246            Input image tensor.
247        - y (torch.tensor): 
248            Target or generated image tensor.
249
250        Returns:
251        - torch.tensor: PatchGAN output tensor predicting real/fake for each patch.
252        """
253        # concatenate input and target channels
254        return self.model(torch.cat([x, y], dim=1))
255
256
257
258# ---------------------------
259#         > Pix2Pix <
260# ---------------------------
261class Pix2Pix(nn.Module):
262    """
263    Pix2Pix GAN for image-to-image translation.
264
265    Components:
266    - Generator: U-Net generator producing synthetic images.
267    - Discriminator: PatchGAN discriminator evaluating real vs fake images.
268    - Adversarial loss: Binary cross-entropy.
269    - Optional second loss for pixel-wise supervision.
270
271    Parameter:
272    - input_channels (int): 
273        Number of input channels.
274    - output_channels (int): 
275        Number of output channels.
276    - hidden_channels (int): 
277        Base hidden channels for both generator and discriminator.
278    - second_loss (nn.Module): 
279        Optional secondary loss (default: L1Loss).
280    - lambda_second (float): 
281        Weight for secondary loss in generator optimization.
282    """
283    def __init__(self, input_channels=1, output_channels=1, hidden_channels=64, 
284                 second_loss=nn.L1Loss(), lambda_second=100):
285        """
286        Initializes the Pix2Pix GAN model.
287
288        Components:
289        - Generator: U-Net architecture for producing synthetic images.
290        - Discriminator: PatchGAN for evaluating real vs. fake images.
291        - Adversarial loss: Binary cross-entropy to train the generator to fool the discriminator.
292        - Optional secondary loss: Pixel-wise supervision (default: L1Loss).
293
294        Parameter:
295        - input_channels (int): 
296            Number of channels in the input images (default=1).
297        - output_channels (int): 
298            Number of channels in the output images (default=1).
299        - hidden_channels (int): 
300            Base number of hidden channels in the generator and discriminator (default=64).
301        - second_loss (nn.Module): 
302            Optional secondary loss for the generator (default: nn.L1Loss()).
303        - lambda_second (float): 
304            Weight applied to the secondary loss in generator optimization (default=100).
305        """
306        super().__init__()
307        self.input_channels = input_channels
308        self.output_channels = output_channels
309
310        self.generator = UNetGenerator(input_channels=input_channels, 
311                                       output_channels=output_channels, 
312                                       hidden_channels=hidden_channels)
313        self.discriminator = Discriminator(input_channels=input_channels+output_channels, hidden_channels=hidden_channels)
314
315        self.adversarial_loss = nn.BCEWithLogitsLoss()
316        self.second_loss = second_loss
317        self.lambda_second = lambda_second
318
319        self.last_generator_loss = float("inf")
320        self.last_generator_adversarial_loss = float("inf")
321        self.last_generator_second_loss = float("inf")
322        self.last_discriminator_loss = float("inf")
323
324    def get_input_channels(self):
325        """
326        Returns the number of input channels used by the model.
327
328        Returns:
329        - int: 
330            Number of input channels expected by the model.
331        """
332        return self.input_channels
333    
334    def get_output_channels(self):
335        """
336        Returns the number of output channels produced by the model.
337
338        Returns:
339        - int: 
340            Number of output channels the model generates
341        """
342        return self.output_channels
343
344    def get_dict(self):
345        """
346        Returns a dictionary with the most recent loss values.
347
348        Returns:
349        - dict: Loss components (base, complex).
350
351        Notes:
352        - Useful for logging or monitoring training progress.
353        """
354        return {
355                f"loss_generator": self.last_generator_loss, 
356                f"loss_generator_adversarial": self.last_generator_adversarial_loss, 
357                f"loss_generator_second": self.last_generator_second_loss,
358                f"loss_discriminator": self.last_discriminator_loss
359               }
360
361    def forward(self, x):
362        """
363        Forward pass through the generator.
364
365        Parameter:
366        - x (torch.tensor): 
367            Input tensor.
368
369        Returns:
370        - torch.tensor: Generated output image.
371        """
372        return self.generator(x)
373
374    def generator_step(self, x, y, optimizer, amp_scaler, device, gradient_clipping_threshold):
375        """
376        Performs a single optimization step for the generator.
377
378        This includes:
379        - Forward pass through the generator and discriminator.
380        - Computing adversarial loss (generator tries to fool the discriminator).
381        - Computing optional secondary loss (e.g., L1 or MSE).
382        - Backpropagation and optimizer step, optionally with AMP and gradient clipping.
383
384        Parameters:
385        - x (torch.tensor): 
386            Input tensor for the generator (e.g., source image).
387        - y (torch.tensor): 
388            Target tensor for supervised secondary loss.
389        - optimizer (torch.optim.Optimizer): 
390            Optimizer for the generator parameters.
391        - amp_scaler (torch.cuda.amp.GradScaler or None): 
392            Automatic mixed precision scaler.
393        - device (torch.device): 
394            Device for AMP autocast.
395        - gradient_clipping_threshold (float or None): 
396            Max norm for gradient clipping; if None, no clipping.
397
398        Returns:
399        - tuple(torch.tensor, torch.tensor, torch.tensor):
400            - Total generator loss (adversarial + secondary).
401            - Adversarial loss component.
402            - Secondary loss component (weighted by `lambda_second`).
403
404        Notes:
405        - If AMP is enabled, gradients are scaled and unscaled appropriately.
406        - `last_generator_loss`, `last_generator_adversarial_loss`, and `last_generator_second_loss` are updated.
407        """
408        if amp_scaler:
409            with autocast(device_type=device.type):
410                # make predictions
411                fake_y = self.generator(x)
412
413                discriminator_fake = self.discriminator(x, fake_y)
414
415                # calc loss -> discriminator thinks it is real?
416                loss_adversarial = self.adversarial_loss(discriminator_fake, torch.ones_like(discriminator_fake))
417                loss_second = self.second_loss(fake_y, y) * self.lambda_second
418                loss_total = loss_adversarial + loss_second
419
420            # backward pass -> calc gradients and change the weights towards the opposite of gradients via optimizer
421    
422            optimizer.zero_grad(set_to_none=True)
423            amp_scaler.scale(loss_total).backward()
424            if gradient_clipping_threshold:
425                # Unscale first!
426                amp_scaler.unscale_(optimizer)
427                torch.nn.utils.clip_grad_norm_(self.generator.parameters(), max_norm=gradient_clipping_threshold)
428            amp_scaler.step(optimizer)
429            amp_scaler.update()
430        else:
431            # make predictions
432            fake_y = self.generator(x)
433
434            discriminator_fake = self.discriminator(x, fake_y)
435
436            # calc loss -> discriminator thinks it is real?
437            loss_adversarial = self.adversarial_loss(discriminator_fake, torch.ones_like(discriminator_fake))
438            loss_second = self.second_loss(fake_y, y) * self.lambda_second
439            loss_total = loss_adversarial + loss_second
440            optimizer.zero_grad()
441            loss_total.backward()
442            if gradient_clipping_threshold:
443                torch.nn.utils.clip_grad_norm_(self.generator.parameters(), max_norm=gradient_clipping_threshold)
444            optimizer.step()
445
446        self.last_generator_loss = loss_total.item()
447        self.last_generator_adversarial_loss = loss_adversarial.item()
448        self.last_generator_second_loss = loss_second.item()
449
450        return loss_total, loss_adversarial, loss_second
451
452    def discriminator_step(self, x, y, optimizer, amp_scaler, device, gradient_clipping_threshold):
453        """
454        Performs a single optimization step for the discriminator.
455
456        This includes:
457        - Forward pass through the discriminator for both real and fake samples.
458        - Computing adversarial loss (binary cross-entropy) for real vs fake patches.
459        - Backpropagation and optimizer step, optionally with AMP and gradient clipping.
460
461        Parameters:
462        - x (torch.tensor): 
463            Input tensor (e.g., source image).
464        - y (torch.tensor): 
465            Target tensor (real image) for the discriminator.
466        - optimizer (torch.optim.Optimizer): 
467            Optimizer for the discriminator parameters.
468        - amp_scaler (torch.cuda.amp.GradScaler or None): 
469            Automatic mixed precision scaler.
470        - device (torch.device): 
471            Device for AMP autocast.
472        - gradient_clipping_threshold (float or None): 
473            Max norm for gradient clipping; if None, no clipping.
474
475        Returns:
476        - tuple(torch.tensor, torch.tensor, torch.tensor):
477            - Total discriminator loss (mean of real and fake losses).
478            - Loss for real samples.
479            - Loss for fake samples.
480
481        Notes:
482        - Fake images are detached from the generator to prevent updating its weights.
483        - `last_discriminator_loss` is updated.
484        - Supports AMP and optional gradient clipping for stability.
485        """
486        if amp_scaler:
487            with autocast(device_type=device.type): 
488                # make predictions
489                fake_y = self.generator(x).detach()  # don't update generator!!
490
491                discriminator_real = self.discriminator(x, y)
492                discriminator_fake = self.discriminator(x, fake_y)
493
494                # calc loss -> 1: predictions = real, 0: predictions = fake
495                loss_real = self.adversarial_loss(discriminator_real, torch.ones_like(discriminator_real))  # torch.full_like(discriminator_real, 0.9)
496                loss_fake = self.adversarial_loss(discriminator_fake, torch.zeros_like(discriminator_fake))
497                loss_total = (loss_real + loss_fake) * 0.5
498
499            # backward pass -> calc gradients and change the weights towards the opposite of gradients via optimizer
500            optimizer.zero_grad(set_to_none=True)
501            amp_scaler.scale(loss_total).backward()
502            if gradient_clipping_threshold:
503                # Unscale first!
504                amp_scaler.unscale_(optimizer)
505                torch.nn.utils.clip_grad_norm_(self.discriminator.parameters(), max_norm=gradient_clipping_threshold)
506            amp_scaler.step(optimizer)
507            amp_scaler.update()
508        else:
509            # make predictions
510            fake_y = self.generator(x).detach()
511
512            discriminator_real = self.discriminator(x, y)
513            discriminator_fake = self.discriminator(x, fake_y)
514
515            # calc loss -> 1: predictions = real, 0: predictions = fake
516            loss_real = self.adversarial_loss(discriminator_real, torch.ones_like(discriminator_real))  # torch.full_like(discriminator_real, 0.9)
517            loss_fake = self.adversarial_loss(discriminator_fake, torch.zeros_like(discriminator_fake))
518            loss_total = (loss_real + loss_fake) * 0.5
519            optimizer.zero_grad()
520            loss_total.backward()
521            if gradient_clipping_threshold:
522                torch.nn.utils.clip_grad_norm_(self.discriminator.parameters(), max_norm=gradient_clipping_threshold)
523            optimizer.step()
524
525        self.last_discriminator_loss = loss_total.item()
526
527        return loss_total, loss_real, loss_fake