Context Encoders: Cleaned up code. Sample in README

Author: eriklindernoren

README.md

@@ -188,6 +188,13 @@ $ cd implementations/context_encoder/
 $ python3 context_encoder.py
 ```
 
+<p align="center">
+    <img src="assets/context_encoder.png" width="640"\>
+</p>
+<p align="center">
+    Rows: Masked | Inpainted | Original | Masked | Inpainted | Original
+</p>
+
 ### Coupled GAN
 _Coupled Generative Adversarial Networks_
 

assets/context_encoder.png

@@ -31,12 +31,12 @@
 
 parser = argparse.ArgumentParser()
 parser.add_argument('--n_epochs', type=int, default=200, help='number of epochs of training')
-parser.add_argument('--batch_size', type=int, default=4, help='size of the batches')
+parser.add_argument('--batch_size', type=int, default=8, help='size of the batches')
 parser.add_argument('--dataset_name', type=str, default='img_align_celeba', help='name of the dataset')
 parser.add_argument('--lr', type=float, default=0.0002, help='adam: learning rate')
 parser.add_argument('--b1', type=float, default=0.5, help='adam: decay of first order momentum of gradient')
 parser.add_argument('--b2', type=float, default=0.999, help='adam: decay of first order momentum of gradient')
-parser.add_argument('--n_cpu', type=int, default=8, help='number of cpu threads to use during batch generation')
+parser.add_argument('--n_cpu', type=int, default=4, help='number of cpu threads to use during batch generation')
 parser.add_argument('--latent_dim', type=int, default=100, help='dimensionality of the latent space')
 parser.add_argument('--img_size', type=int, default=128, help='size of each image dimension')
 parser.add_argument('--mask_size', type=int, default=64, help='size of random mask')
@@ -83,7 +83,7 @@ def weights_init_normal(m):
                 transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5)) ]
 dataloader = DataLoader(ImageDataset("../../data/%s" % opt.dataset_name, transforms_=transforms_),
                         batch_size=opt.batch_size, shuffle=True, num_workers=opt.n_cpu)
-sample_dataloader = DataLoader(ImageDataset("../../data/%s" % opt.dataset_name, transforms_=transforms_),
+test_dataloader = DataLoader(ImageDataset("../../data/%s" % opt.dataset_name, transforms_=transforms_, mode='val'),
                         batch_size=12, shuffle=True, num_workers=1)
 
 # Optimizers
@@ -92,35 +92,11 @@ def weights_init_normal(m):
 
 Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor
 
-# Adversarial ground truths
-valid = Variable(Tensor(np.ones(patch)), requires_grad=False)
-fake = Variable(Tensor(np.zeros(patch)), requires_grad=False)
-
-def apply_random_mask(imgs):
-    idx = np.random.randint(0, opt.img_size-opt.mask_size, (imgs.shape[0], 2))
-
-    masked_imgs = imgs.clone()
-    masked_parts = None
-    for i, (y1, x1) in enumerate(idx):
-        y2, x2 = y1 + opt.mask_size, x1 + opt.mask_size
-        masked_part = masked_imgs[i:i+1, :, y1:y2, x1:x2].clone()
-        masked_parts = masked_part if masked_parts is None else torch.cat((masked_parts, masked_part),  0)
-        masked_imgs[i, :, y1:y2, x1:x2] = 1
-
-    return masked_imgs, masked_parts
-
-def apply_center_mask(imgs):
-    # Get upper-left pixel coordinate
-    i = (imgs.shape[2] - opt.mask_size) // 2
-
-    masked_imgs = imgs.clone()
-    masked_imgs[:, :, i:i+opt.mask_size, i:i+opt.mask_size] = 1
-
-    return masked_imgs, i
-
 def save_sample(batches_done):
-    samples = Variable(next(iter(sample_dataloader)).type(Tensor))
-    masked_samples, i = apply_center_mask(samples)
+    samples, masked_samples, i = next(iter(test_dataloader))
+    samples = Variable(samples.type(Tensor))
+    masked_samples = Variable(masked_samples.type(Tensor))
+    i = i[0].item() # Upper-left coordinate of mask
     # Generate inpainted image
     gen_mask = generator(masked_samples)
     filled_samples = masked_samples.clone()
@@ -134,9 +110,7 @@ def save_sample(batches_done):
 # ----------
 
 for epoch in range(opt.n_epochs):
-    for i, imgs in enumerate(dataloader):
-
-        masked_imgs, masked_parts = apply_random_mask(imgs)
+    for i, (imgs, masked_imgs, masked_parts) in enumerate(dataloader):
 
         # Adversarial ground truths
         valid = Variable(Tensor(imgs.shape[0], *patch).fill_(1.0), requires_grad=False)

implementations/context_encoder/context_encoder.py

@@ -8,17 +8,45 @@
 import torchvision.transforms as transforms
 
 class ImageDataset(Dataset):
-    def __init__(self, root, transforms_=None, mode='train'):
+    def __init__(self, root, transforms_=None, img_size=128, mask_size=64, mode='train'):
         self.transform = transforms.Compose(transforms_)
+        self.img_size = img_size
+        self.mask_size = mask_size
+        self.mode = mode
+        self.files = sorted(glob.glob('%s/*.jpg' % root))
+        self.files = self.files[:-4000] if mode == 'train' else self.files[-4000:]
 
-        self.files = sorted(glob.glob('%s/*.*' % root))
+    def apply_random_mask(self, img):
+        """Randomly masks image"""
+        y1, x1 = np.random.randint(0, self.img_size-self.mask_size, 2)
+        y2, x2 = y1 + self.mask_size, x1 + self.mask_size
+        masked_part = img[:, y1:y2, x1:x2]
+        masked_img = img.clone()
+        masked_img[:, y1:y2, x1:x2] = 1
+
+        return masked_img, masked_part
+
+    def apply_center_mask(self, img):
+        """Mask center part of image"""
+        # Get upper-left pixel coordinate
+        i = (self.img_size - self.mask_size) // 2
+        masked_img = img.clone()
+        masked_img[:, i:i+self.mask_size, i:i+self.mask_size] = 1
+
+        return masked_img, i
 
     def __getitem__(self, index):
 
         img = Image.open(self.files[index % len(self.files)])
         img = self.transform(img)
+        if self.mode == 'train':
+            # For training data perform random mask
+            masked_img, aux = self.apply_random_mask(img)
+        else:
+            # For test data mask the center of the image
+            masked_img, aux = self.apply_center_mask(img)
 
-        return img
+        return img, masked_img, aux
 
     def __len__(self):
         return len(self.files)

implementations/context_encoder/datasets.py

@@ -11,14 +11,14 @@ def downsample(in_feat, out_feat, normalize=True):
             layers = [nn.Conv2d(in_feat, out_feat, 4, stride=2, padding=1)]
             if normalize:
                 layers.append(nn.BatchNorm2d(out_feat, 0.8))
-            layers.append(nn.LeakyReLU(0.2, inplace=True))
+            layers.append(nn.LeakyReLU(0.2))
             return layers
 
         def upsample(in_feat, out_feat, normalize=True):
             layers = [nn.ConvTranspose2d(in_feat, out_feat, 4, stride=2, padding=1)]
             if normalize:
                 layers.append(nn.BatchNorm2d(out_feat, 0.8))
-            layers.append(nn.ReLU(inplace=True))
+            layers.append(nn.ReLU())
             return layers
 
         self.model = nn.Sequential(