File size: 2,625 Bytes
5ade2ef
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
---
license: mit
tags:
- image-to-image
- gan
- pix2pix
- pytorch
- conditional-gan
pipeline_tag: image-to-image
---

# Pix2Pix — Edges to Handbags

A conditional GAN implementing the pix2pix framework (Isola et al., 2017) for paired image-to-image translation. This model translates edge maps of handbags into realistic photographic renderings.

## Model Description

Pix2Pix learns a mapping from input condition images to output target images using an adversarial training objective. The generator is supervised by both an adversarial loss (fooling the discriminator) and an L1 reconstruction loss (staying close to the ground truth).

### Architecture

**Generator — U-Net**
- Encoder: stacked `Conv2d` + `BatchNorm2d` + `ReLU` + `MaxPool2d` blocks, doubling channels at each stage
- Decoder: `Upsample` (bilinear) + `Conv2d` blocks with skip connections from the corresponding encoder stage
- Output: Tanh activation to produce pixel values in [-1, 1]

**Discriminator — PatchGAN**
- Takes the concatenation of condition and real/fake image as input (6 channels)
- Classifies whether overlapping image patches are real or generated
- Trained with `BCEWithLogitsLoss`

### Loss Functions

```
L_total = L_adversarial + lambda_recon * L_L1
        = BCEWithLogitsLoss + 200 * L1Loss
```

## Training Details

| Parameter | Value |
|---|---|
| Dataset | edges2handbags |
| Image resolution | 256 x 256 |
| Epochs | 50 (checkpoint saved at epoch 59) |
| Batch size | 4 |
| Learning rate | 0.0002 |
| Optimizer | Adam (both G and D) |
| Weight initialization | Normal distribution (mean=0, std=0.02) |
| lambda_recon (L1 weight) | 200 |

## Repository Contents

| File | Description |
|---|---|
| train.py | Training loop |
| UNet.py | Generator (U-Net) and discriminator architecture |
| utils.py | Helper functions |
| dataset.sh | Downloads the edges2handbags dataset |
| Pix2Pix_Epoch59.pth | Saved generator checkpoint |

## How to Use

```python
import torch
from UNet import UNet

device = "cuda" if torch.cuda.is_available() else "cpu"
gen = UNet(input_dim=3, real_dim=3).to(device)
checkpoint = torch.load("Pix2Pix_Epoch59.pth", map_location=device)
gen.load_state_dict(checkpoint)
gen.eval()

# condition: edge map tensor of shape (B, 3, 256, 256), normalized to [-1, 1]
with torch.no_grad():
    fake = gen(condition)
```

## References

- Isola et al. (2017). [Image-to-Image Translation with Conditional Adversarial Networks](https://arxiv.org/abs/1611.07004)
- Ronneberger et al. (2015). [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://arxiv.org/abs/1505.04597)

## License

MIT