More transform functions added

Author: Lxrd-AJ

Dataset/MattingDataset.py

@@ -8,18 +8,18 @@
 from PIL import Image, ImageFilter, ImageChops
 
 class MattingDataset(data.Dataset):
-    def __init__(self, fgDir, bgDir, alphaDir, allTransform):
+    def __init__(self, fgDir, bgDir, alphaDir, allTransform, imageTransforms):
         self.fgDir = fgDir
         self.bgDir = bgDir
         self.alphaDir = alphaDir
 
         self.foregroundImageNames = os.listdir(self.fgDir)
         self.backgroundImageNames = os.listdir(self.bgDir)
         random.shuffle(self.backgroundImageNames) #TODO: Remove
-        self.backgroundImageNames = self.backgroundImageNames[:10] #TODO: Remove
+        self.backgroundImageNames = self.backgroundImageNames[:12] #TODO: Remove
         self.alphaImageNames = os.listdir(self.alphaDir)
         random.shuffle(self.alphaImageNames) #TODO: Remove
-        self.alphaImageNames = self.alphaImageNames[:20] #TODO:Remove
+        self.alphaImageNames = self.alphaImageNames[:22] #TODO:Remove
 
         self.numForeground = len(self.foregroundImageNames)
         self.numBackground = len(self.backgroundImageNames)
@@ -31,6 +31,7 @@ def __init__(self, fgDir, bgDir, alphaDir, allTransform):
         self.imageBackgroundPair = sorted(self.imageBackgroundPair, key=lambda x: x[0])
 
         self.allTransform = allTransform
+        self.imageTransform = imageTransforms
 
         # assert len(self.imageBackgroundPair) == len(self) #TODO: Remove
 
@@ -50,6 +51,9 @@ def __getitem__(self, idx):
         compositeImage = self.composite_image(foregroundImage, backgroundImage, alphaMask)
 
         assert compositeImage.size == trimap.size, f"composite size = {compositeImage.size} and trimap = {trimap.size} and foreground size = {foregroundImage.size}"
+
+        if self.imageTransform:
+            compositeImage = self.imageTransform(compositeImage)
 
         if self.allTransform:
             compositeImage, trimap, alphaMask = self.allTransform((compositeImage, trimap, alphaMask))

TrainLoss200Items.png

@@ -2,8 +2,54 @@
 import random
 import torch
 import torchvision.transforms.functional as TF
-from PIL import Image
+from PIL import Image, ImageFilter
 
+class RandomBlur(object):
+    def __init__(self, probability=0.5):
+        self.p = probability
+
+    def __call__(self, image):
+        if random.random() < self.p:
+            return image.filter(ImageFilter.GaussianBlur(radius=2))
+        return image
+
+class RandomRotation(object):
+    def __init__(self, probability=0.5, angle=45):
+        self.p = probability
+        self.angle = angle
+
+    def __call__(self, items):
+        image, trimap, mask = items
+        angle = random.randint(-self.angle, self.angle)
+        if random.random() < self.p:
+            image = TF.rotate(image, angle)
+            trimap = TF.rotate(trimap, angle)
+            mask = TF.rotate(mask, angle)
+        return image, trimap, mask
+
+class RandomVerticalFlip(object):
+    def __init__(self, probability=0.5):
+        self.p = probability
+
+    def __call__(self, items):
+        image, trimap, mask = items
+        if random.random() < self.p:
+            image = TF.vflip(image)
+            trimap = TF.vflip(trimap)
+            mask = TF.vflip(mask)
+        return image, trimap, mask
+
+class RandomHorizontalFlip(object):
+    def __init__(self, probability=0.5):
+        self.p = probability
+
+    def __call__(self, items):
+        image, trimap, mask = items
+        if random.random() < self.p:
+            image = TF.hflip(image)
+            trimap = TF.hflip(trimap)
+            mask = TF.hflip(mask)
+        return image, trimap, mask
 
 class ToTensor(object):
     def __call__(self, items):

TrainLoss231Items.png

@@ -38,7 +38,6 @@ def forward(self, x):
         """
         x = self.encoder(x)
         x = self.decoder(x)
-
         return x
 
 
@@ -51,9 +50,13 @@ def __init__(self):
         """
         self.encoderBlocks = nn.Sequential(
             convBatchNormReLU(4, 64, 3),
+            convBatchNormReLU(64, 64, 1, pad=0, stride=1),
             convBatchNormReLU(64, 128, 3),
+            convBatchNormReLU(128, 128, 1, pad=0, stride=1),
             convBatchNormReLU(128, 256, 3),
-            convBatchNormReLU(256, 512, 3)
+            convBatchNormReLU(256, 256, 1, pad=0, stride=1),
+            convBatchNormReLU(256, 512, 3),
+            convBatchNormReLU(512, 512, 1, pad=0, stride=1),
         )
 
     def forward(self, x):

TrainLoss264Items.png

@@ -8,7 +8,7 @@
 from PIL import Image
 from Dataset.MattingDataset import MattingDataset
 from model import EncoderDecoderNet, RefinementNet
-from dataset_transforms import RandomTrimapCrop, Resize, ToTensor
+from dataset_transforms import RandomTrimapCrop, Resize, ToTensor, RandomHorizontalFlip, RandomRotation, RandomVerticalFlip, RandomBlur
 from loss import alpha_prediction_loss, compositional_loss
 
 
@@ -45,21 +45,30 @@ def batch_collate_fn(batch):
 _TRAIN_ALPHA_DIR_ = "./Dataset/Training_set/CombinedAlpha"
 _NETWORK_INPUT_ = (320,320)
 _COMPUTE_DEVICE_ = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-_NUM_EPOCHS_ = 30 #200 #TODO: Remove
+_NUM_EPOCHS_ = 60 #200 #TODO: Remove
 _BATCH_SIZE_ = 8 #TODO: Increase this if using a GPU
 _NUM_WORKERS_ = multiprocessing.cpu_count()
 _LOSS_WEIGHT_ = 0.4 #0.5
 _GRADIENT_CLIP_ = 2.5
 
 tripleTransforms = transforms.Compose([
+    RandomRotation(probability=0.5, angle=180),
+    RandomVerticalFlip(probability=0.5),
+    RandomHorizontalFlip(probability=0.5),
     RandomTrimapCrop([(320, 320), (480, 480), (640, 640)], probability=0.7),
     Resize(_NETWORK_INPUT_),
     ToTensor()
 ])
 
+imageTransforms = transforms.Compose([
+    transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.25),
+    transforms.RandomGrayscale(p=0.3),
+    RandomBlur(probability=0.1)
+])
+
 trainingDataset = MattingDataset(
                         _TRAIN_FOREGROUND_DIR_, _TRAIN_BACKGROUND_DIR_, _TRAIN_ALPHA_DIR_, 
-                        allTransform=tripleTransforms
+                        allTransform=tripleTransforms, imageTransforms=imageTransforms
                     )
 trainDataloader = torch.utils.data.DataLoader(
                             trainingDataset, batch_size=_BATCH_SIZE_, shuffle=True, num_workers=_NUM_WORKERS_, collate_fn=batch_collate_fn)
@@ -142,9 +151,11 @@ def batch_collate_fn(batch):
         plt.title(f"Training loss using a dataset of {len(trainingDataset)} images")
         plt.savefig(f"TrainLoss{len(trainingDataset)}Items.png")
 
+    trainingElapsed = time.time() - trainStart
+    print(f"\nTotal training time is {trainingElapsed//60:.0f}m {trainingElapsed%60:.0f}s")
     #Make a sample prediction
     idx = random.choice(range(0, len(trainingDataset)))
-    img_, trimap, gMasks = trainingDataset[0]
+    img_, trimap, gMasks = trainingDataset[idx]
     trimap = trimap.unsqueeze(0)
     gMasks = gMasks.unsqueeze(0)
     img = torch.cat([img_, trimap], 0).unsqueeze(0)