Sum absolute difference and mean squared error metrics added

Author: Lxrd-AJ

Dataset/MattingDataset.py

@@ -16,10 +16,10 @@ def __init__(self, fgDir, bgDir, alphaDir, allTransform, imageTransforms):
         self.foregroundImageNames = os.listdir(self.fgDir)
         self.backgroundImageNames = os.listdir(self.bgDir)
         random.shuffle(self.backgroundImageNames) #TODO: Remove
-        self.backgroundImageNames = self.backgroundImageNames[:3] #TODO: Remove 13
+        self.backgroundImageNames = self.backgroundImageNames[:4] #TODO: Remove 13
         self.alphaImageNames = os.listdir(self.alphaDir)
         random.shuffle(self.alphaImageNames) #TODO: Remove
-        self.alphaImageNames = self.alphaImageNames[:3] #TODO:Remove 23
+        self.alphaImageNames = self.alphaImageNames[:4] #TODO:Remove 23
 
         self.numForeground = len(self.foregroundImageNames)
         self.numBackground = len(self.backgroundImageNames)

TrainLoss16Items.png

@@ -4,6 +4,20 @@
 import torchvision.transforms.functional as TF
 from PIL import Image, ImageFilter
 
+class RandomAffine(object):
+    def __init__(self, probability=0.5):
+        self.p = probability
+
+    def __call__(self, items):
+        image, trimap, mask = items
+        if random.random() < self.p:
+            angle = random.randint(-180, 180)
+            image = TF.affine(image, angle, translate=[0,0], scale=1.0, shear=0, resample=Image.BICUBIC)
+            # use nearest so the values of the trimap and alpha mask are not changed
+            trimap = TF.affine(trimap, angle, translate=[0,0], scale=1.0, shear=0, resample=Image.NEAREST)
+            mask = TF.affine(mask, angle, translate=[0,0], scale=1.0, shear=0, resample=Image.NEAREST)
+        return image, trimap, mask
+
 class RandomBlur(object):
     def __init__(self, probability=0.5):
         self.p = probability

TrainLoss25Items.png

@@ -61,4 +61,22 @@ def show(xf):
     sumTrueForeground = trueForeground.sum(dim=[2,3]) + eps
     totalLoss = rootDiff.sum(dim=[2,3]) / sumTrueForeground
     avgLoss = totalLoss.mean().mean() # average over the RGB channels and also across the batch
-    return avgLoss
+    return avgLoss
+
+
+def sum_absolute_difference(trueAlpha, predAlpha):
+    """
+    calculates the sum of absolute differences between images and predictions in batches
+    As the calculation is done over a batch, the mean is used to reduce the results
+    """
+    difference = predAlpha - trueAlpha
+    avgDiff = difference.sum(dim=[1,2]).mean()
+    return avgDiff
+
+def mean_squared_error(trueAlpha, predAlpha, compositeImage):
+    trimaps = compositeImage[:,3,:] * 255
+    blackMask = torch.zeros_like(trueAlpha)
+    unknownRegions = torch.where(trimaps == 127, trueAlpha, blackMask)
+    mse = torch.pow(predAlpha - trueAlpha, 2).sum() / unknownRegions.sum()
+
+    return mse

dataset_transforms.py

@@ -8,8 +8,8 @@
 from PIL import Image
 from Dataset.MattingDataset import MattingDataset
 from model import EncoderDecoderNet, RefinementNet
-from dataset_transforms import RandomTrimapCrop, Resize, ToTensor, RandomHorizontalFlip, RandomRotation, RandomVerticalFlip, RandomBlur
-from loss import alpha_prediction_loss, compositional_loss
+from dataset_transforms import RandomTrimapCrop, Resize, ToTensor, RandomHorizontalFlip, RandomRotation, RandomVerticalFlip, RandomBlur, RandomAffine
+from loss import alpha_prediction_loss, compositional_loss, sum_absolute_difference, mean_squared_error
 
 
 
@@ -45,7 +45,7 @@ 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_ = 8 #200 #TODO: Remove 90
+_NUM_EPOCHS_ = 30 #200 #TODO: Remove 90
 _BATCH_SIZE_ = 8 #TODO: Increase this if using a GPU
 _NUM_WORKERS_ = multiprocessing.cpu_count()
 _LOSS_WEIGHT_ = 0.6
@@ -68,7 +68,7 @@ def batch_collate_fn(batch):
 
 trainingDataset = MattingDataset(
                         _TRAIN_FOREGROUND_DIR_, _TRAIN_BACKGROUND_DIR_, _TRAIN_ALPHA_DIR_, 
-                        allTransform=tripleTransforms, imageTransforms=imageTransforms
+                        allTransform=tripleTransforms, imageTransforms=None
                     )
 trainDataloader = torch.utils.data.DataLoader(
                             trainingDataset, batch_size=_BATCH_SIZE_, shuffle=True, num_workers=_NUM_WORKERS_, collate_fn=batch_collate_fn)
@@ -115,10 +115,12 @@ def batch_collate_fn(batch):
                 modelAlphaLoss = alpha_prediction_loss(predictedMasks, groundTruthMasks)                
                 refinedAlphaLoss = alpha_prediction_loss(refinedMasks, groundTruthMasks)
                 lossAlpha = modelAlphaLoss + refinedAlphaLoss
-                # lossComposition = compositional_loss(predictedMasks, groundTruthMasks, compositeImages)
-                lossComposition = compositional_loss(refinedMasks, groundTruthMasks, compositeImages)
+                lossComposition = compositional_loss(predictedMasks, groundTruthMasks, compositeImages)
                 totalLoss = _LOSS_WEIGHT_ * lossAlpha + (1 - _LOSS_WEIGHT_) * lossComposition
                 epochLoss += totalLoss.item()
+                with torch.no_grad():
+                    sad = sum_absolute_difference(groundTruthMasks, refinedMasks)
+                    mse = mean_squared_error(groundTruthMasks, refinedMasks, compositeImages)
 
                 if idx % 100 == 0:
                     print(f"\tIteration {idx+1}/{len(trainingDataset)}")
@@ -128,6 +130,11 @@ def batch_collate_fn(batch):
                     print(f"\t Alpha loss = {lossAlpha}")
                     print(f"\t Composition loss = {lossComposition}")
                     print(f"\t Total Loss = {totalLoss}")
+                    print(f"\t {'***' * 5}")
+                    print(f"\t Metrics:")
+                    print(f"\t {'***' * 5}")
+                    print(f"\t Sum absolute difference: {sad}")
+                    print(f"\t Mean Squared Error: {mse}")
                     print()
 
                 optimiser.zero_grad()