Eraser

README.md

@@ -228,6 +228,14 @@ Spatial-level transforms will simultaneously change both an input image as well
 ### Keypoints augmentation
 <img src="https://habrastorage.org/webt/e-/6k/z-/e-6kz-fugp2heak3jzns3bc-r8o.jpeg" width=100%>
 
+### Object Removal augmentation
+Using `startEraser(parentDir, imgType, imgExpType)` one can easily equalise strength of different classes of a dataset. The dataset must contain two folders, 'images' folder and 'labels' folder. Using the 'labels' folder, `startEraser()` finds the minority class and then avoids images that doesn't contain the minority classes for synthesis. Later on the rest of images that has the minority class, `startEraser()` blurs out instances of classes that are in excess from the images. It also returns proper labels of the synthesised images.
+#### Original Image
+![object blurring 1](/albumentations/augmentations/eraser/sampleImages/000277.png)
+#### Sythesised Image
+![object blurring 1](/albumentations/augmentations/eraser/sampleImages/000277.jpg)
+#### Newly Generated Label Comparison
+![old new lable](/albumentations/augmentations/eraser/sampleImages/oldNewLabels.png)
 
 ## Benchmarking results
 To run the benchmark yourself, follow the instructions in [benchmark/README.md](https://github.com/albumentations-team/albumentations/blob/master/benchmark/README.md)

albumentations/augmentations/__init__.py

@@ -12,6 +12,7 @@
 from .dropout.functional import *
 from .dropout.grid_dropout import *
 from .dropout.mask_dropout import *
+from .eraser.metadata import *
 from .functional import *
 from .geometric.functional import *
 from .geometric.resize import *

albumentations/augmentations/eraser/__init__.py

@@ -0,0 +1 @@
+from .metadata import *

albumentations/augmentations/eraser/functional.py

@@ -0,0 +1,187 @@
+import glob
+import os
+import random
+
+import cv2
+import numpy as np
+
+
+class MetaData:
+    """Class to extract data from labels of kitti format"""
+
+    def __init__(self, parentDir, imgType):
+        """Load parent directory and Image type present in the image directory
+
+        Args:
+            parentDir (string): Path to the directory that contains the Image folder and Label folder
+            imgType (string): Image type (Ex: JPEG: jpg, PNG: png)
+        """
+        self.labelPath = os.path.join(parentDir, "labels")
+        self.imgPath = os.path.join(parentDir, "images")
+        self.imgType = imgType
+        self.data = {}
+        self.distribution = {}
+        self.labelStrengthOrder = []
+
+    def loadData(self):
+        """Load data from the txt files and stores it as a dictionary.
+        Also loads the distribution of classification classes present in the dataset.
+        """
+        for filePath in glob.glob(f"{self.labelPath}/*"):
+            filename = os.path.basename(filePath).split(".")[0]
+            if os.path.isfile(os.path.join(self.imgPath, f"{filename}.{self.imgType}")):
+                self.data[filename] = {}
+                self.data[filename]["labels"] = {}
+                self.data[filename]["n_labels"] = {}
+                with open(filePath) as f:
+                    for line in f:
+                        label = line.split()[0]
+                        if label not in self.data[filename]["labels"].keys():
+                            self.data[filename]["labels"][label] = {}
+                            self.data[filename]["n_labels"][label] = 1
+                        else:
+                            self.data[filename]["n_labels"][label] += 1
+
+                        numTag = self.data[filename]["n_labels"][label]
+                        xmin, ymin, xmax, ymax = line.split()[4:8]
+                        length = abs(float(xmin) - float(xmax))
+                        breadth = abs(float(ymin) - float(ymax))
+                        area = length * breadth
+
+                        truncation = line.split()[1]
+                        occlusion = line.split()[2]
+                        alpha = line.split()[3]
+                        threeDdim = line.split()[8:11]
+                        location = line.split()[11:14]
+                        rotationY = line.split()[14]
+
+                        self.data[filename]["labels"][label][f"l{numTag}"] = {
+                            "coord": [(int(float(xmin)), int(float(ymin))), (int(float(xmax)), int(float(ymax)))],
+                            "area": area,
+                            "Truncation": truncation,
+                            "Occlusion": occlusion,
+                            "Alpha": alpha,
+                            "ThreeDdim": threeDdim,
+                            "Location": location,
+                            "RotationY": rotationY,
+                        }
+
+            for label in self.data[filename]["n_labels"]:
+                if label not in self.distribution.keys():
+                    self.distribution[label] = self.data[filename]["n_labels"][label]
+                else:
+                    self.distribution[label] += self.data[filename]["n_labels"][label]
+
+    def identifyMinority(self):
+        """Sorts the distribution of classification classes to return class with least number of detections and its strength.
+
+        Returns:
+            list: the first tells us the minority class and the second its number of occurences.
+        """
+        self.labelStrengthOrder = list(dict(sorted(self.distribution.items(), key=lambda item: item[1])).keys())
+        minorityLabel = self.labelStrengthOrder[0]
+        minorityStrength = self.distribution[minorityLabel]
+
+        return [minorityLabel, minorityStrength]
+
+
+def startEraser(parentDir, imgType, imgExpType):
+    """Start the synthesis process of normalizing the classes and
+    generates a synthesis image directory and a synthesis label directory.
+
+    Args:
+        parentDir (string): Path to parent directory that contains image and label directories.
+        imgType (string): Type of image present in the Image directory
+        imgExpType (string): Type of synthesised image to be generated
+    """
+    imgDir = os.path.join(parentDir, "images")
+    synthImgDir = os.path.join(parentDir, "syntheticImages")
+    synthLabDir = os.path.join(parentDir, "syntheticLabels")
+
+    try:
+        os.mkdir(synthImgDir)
+        os.mkdir(synthLabDir)
+    except FileExistsError:
+        print("Synthetic dir already created!")
+
+    obj1 = MetaData(parentDir, imgType)
+    obj1.loadData()
+    # print(obj1.data)
+    print(f"Identified:{obj1.distribution}")
+    minorityLabel, minorityStrength = obj1.identifyMinority()
+
+    for filePath in glob.glob(f"{imgDir}/*"):
+        filename = os.path.basename(filePath).split(".")[0]
+        img = cv2.imread(filePath)
+        mask = np.full(img.shape[:2], 255, dtype="uint8")
+
+        labelData = obj1.data[filename]["labels"]
+        labelCount = obj1.data[filename]["n_labels"]
+
+        if minorityLabel not in labelData.keys():
+            continue
+
+        for labelname in labelData.keys():
+            for tag in labelData[labelname]:
+                coord = labelData[labelname][tag]["coord"]
+                cv2.rectangle(mask, coord[0], coord[1], 0, -1)
+
+        invMask = 255 - mask
+        inPaint = cv2.inpaint(img, invMask, 3, cv2.INPAINT_NS)
+        text = []
+        for label in labelData.keys():
+            if label != minorityLabel:
+                tags = list(labelData[label].keys())
+                nTags = labelCount[label]
+                nAdd = labelCount[minorityLabel]
+                if nTags > nAdd:
+                    tagAdd = random.sample(tags, nAdd)
+                else:
+                    tagAdd = tags
+                for tag in tagAdd:
+                    set1 = " ".join(
+                        [
+                            label,
+                            labelData[label][tag]["Truncation"],
+                            labelData[label][tag]["Occlusion"],
+                            labelData[label][tag]["Alpha"],
+                        ]
+                    )
+                    set3 = " ".join(labelData[label][tag]["ThreeDdim"])
+                    set4 = " ".join(labelData[label][tag]["Location"])
+                    set5 = labelData[label][tag]["RotationY"]
+
+                    (xmin, ymin), (xmax, ymax) = labelData[label][tag]["coord"]
+                    set2 = " ".join([str(xmin), str(ymin), str(xmax), str(ymax)])
+                    inPaint[ymin:ymax, xmin:xmax, :] = img[ymin:ymax, xmin:xmax, :]
+
+                    line = " ".join([set1, set2, set3, set4, set5])
+                    text.append(line)
+                print(f"Adding these tags of {label}:{tagAdd}")
+            else:
+                minTags = list(labelData[label].keys())
+                for tag in minTags:
+                    set1 = " ".join(
+                        [
+                            label,
+                            labelData[label][tag]["Truncation"],
+                            labelData[label][tag]["Occlusion"],
+                            labelData[label][tag]["Alpha"],
+                        ]
+                    )
+                    set3 = " ".join(labelData[label][tag]["ThreeDdim"])
+                    set4 = " ".join(labelData[label][tag]["Location"])
+                    set5 = labelData[label][tag]["RotationY"]
+                    (xmin, ymin), (xmax, ymax) = labelData[label][tag]["coord"]
+                    set2 = " ".join([str(xmin), str(ymin), str(xmax), str(ymax)])
+                    inPaint[ymin:ymax, xmin:xmax, :] = img[ymin:ymax, xmin:xmax, :]
+
+                    line = " ".join([set1, set2, set3, set4, set5])
+                    text.append(line)
+
+        cv2.imwrite(os.path.join(synthImgDir, f"{filename}.{imgExpType}"), inPaint)
+        print(f"Synthesised synthetic Image for {filename}.")
+        fileContents = "\n".join(text)
+        with open(os.path.join(synthLabDir, f"{filename}.txt"), "w") as f:
+            f.write(fileContents)
+        print(f"New label for synthesised {filename} written.")

albumentations/augmentations/eraser/metadata.py

@@ -0,0 +1,167 @@
+import glob
+import os
+import random
+
+import cv2
+import numpy as np
+
+from ...core.transforms_interface import ImageOnlyTransform
+from . import functional as F
+
+
+def load_label_data(label_file):
+    label_data = {}
+    label_data["labels"] = {}
+    label_data["n_labels"] = {}
+    with open(label_file) as f:
+        for line in f:
+            label = line.split()[0]
+            if label not in label_data["labels"]:
+                label_data["labels"][label] = {}
+                label_data["n_labels"][label] = 1
+            else:
+                label_data["n_labels"][label] += 1
+
+            numTag = label_data["n_labels"][label]
+            xmin, ymin, xmax, ymax = line.split()[4:8]
+            length = abs(float(xmin) - float(xmax))
+            breadth = abs(float(ymin) - float(ymax))
+            area = length * breadth
+
+            truncation = line.split()[1]
+            occlusion = line.split()[2]
+            alpha = line.split()[3]
+            threeDdim = line.split()[8:11]
+            location = line.split()[11:14]
+            rotationY = line.split()[14]
+
+            label_data["labels"][label][f"l{numTag}"] = {
+                "coord": [(int(float(xmin)), int(float(ymin))), (int(float(xmax)), int(float(ymax)))],
+                "area": area,
+                "Truncation": truncation,
+                "Occlusion": occlusion,
+                "Alpha": alpha,
+                "ThreeDdim": threeDdim,
+                "Location": location,
+                "RotationY": rotationY,
+            }
+
+    return label_data
+
+
+def generate_img(img, label_names, label_count, minority_label, new_label_file_name, parent_dir):
+    mask = np.full(img.shape[:2], 255, dtype="uint8")
+    for label_name in label_names:
+        for tag in label_names[label_name]:
+            coord = label_names[label_name][tag]["coord"]
+            cv2.rectangle(mask, coord[0], coord[1], 0, -1)
+
+    inv_mask = 255 - mask
+    inPaint = cv2.inpaint(img, inv_mask, 3, cv2.INPAINT_NS)
+
+    text = []
+
+    for label in label_names:
+        if label != minority_label:
+            tags = list(label_names[label].keys())
+            n_tags = label_count[label]
+            n_add = label_count[minority_label]
+
+            if n_tags > n_add:
+                tag_add = random.sample(tags, n_add)
+            else:
+                tag_add = tags
+
+            for tag in tag_add:
+                set1 = " ".join(
+                    [
+                        label,
+                        label_names[label][tag]["Truncation"],
+                        label_names[label][tag]["Occlusion"],
+                        label_names[label][tag]["Alpha"],
+                    ]
+                )
+                set3 = " ".join(label_names[label][tag]["ThreeDdim"])
+                set4 = " ".join(label_names[label][tag]["Location"])
+                set5 = label_names[label][tag]["RotationY"]
+
+                (xmin, ymin), (xmax, ymax) = label_names[label][tag]["coord"]
+                set2 = " ".join([str(xmin), str(ymin), str(xmax), str(ymax)])
+                inPaint[ymin:ymax, xmin:xmax, :] = img[ymin:ymax, xmin:xmax, :]
+
+                line = " ".join([set1, set2, set3, set4, set5])
+                text.append(line)
+            print(f"Adding these tags of {label}:{tag_add}")
+        else:
+            minTags = list(label_names[label].keys())
+            for tag in minTags:
+                set1 = " ".join(
+                    [
+                        label,
+                        label_names[label][tag]["Truncation"],
+                        label_names[label][tag]["Occlusion"],
+                        label_names[label][tag]["Alpha"],
+                    ]
+                )
+                set3 = " ".join(label_names[label][tag]["ThreeDdim"])
+                set4 = " ".join(label_names[label][tag]["Location"])
+                set5 = label_names[label][tag]["RotationY"]
+                (xmin, ymin), (xmax, ymax) = label_names[label][tag]["coord"]
+                set2 = " ".join([str(xmin), str(ymin), str(xmax), str(ymax)])
+                inPaint[ymin:ymax, xmin:xmax, :] = img[ymin:ymax, xmin:xmax, :]
+
+                line = " ".join([set1, set2, set3, set4, set5])
+                text.append(line)
+    file_contents = "\n".join(text)
+    with open(os.path.join(parent_dir, f"{new_label_file_name}.txt"), "w") as f:
+        f.write(file_contents)
+
+    return inPaint
+
+
+def apply_earser(img, label_file, minority_label, new_label_file_name):
+    try:
+        if os.path.getsize(label_file) > 0:
+            label_data = load_label_data(label_file)
+            label_names = label_data["labels"]
+            if minority_label not in label_names:
+                raise RuntimeWarning(f"Given label file doesn't have {minority_label} label!")
+            else:
+                label_count = label_data["n_labels"]
+                parent_dir = os.path.dirname(label_file)
+                img = generate_img(img, label_names, label_count, minority_label, new_label_file_name, parent_dir)
+        else:
+            raise RuntimeError(f"{label_file} is empty!")
+    except OSError as e:
+        print(f"Error reading label file: {e}")
+
+    return img
+
+
+class MetaData(ImageOnlyTransform):
+    """Class to extract data from labels of kitti format"""
+
+    def __init__(self, always_apply=False, p=0.5):
+        """
+        Erase instances of labels in image.
+
+        Args:
+            label_file (_type_): _description_
+            minority_label (_type_): _description_
+            new_label_file_name (_type_): _description_
+            always_apply (bool, optional): _description_. Defaults to False.
+            p (float, optional): _description_. Defaults to 0.5.
+        """
+        super().__init__(always_apply=always_apply, p=p)
+
+    def __call__(self, **data):
+        self.label_file = data["label_file"]
+        self.minority_label = data["minority_label"]
+        self.new_label_file_name = data["new_label_file_name"]
+
+        data["image"] = self.apply(**data)
+
+        return data
+
+    def apply(self, image, label_file, minority_label, new_label_file_name, **params):
+        return apply_earser(image, label_file, minority_label, new_label_file_name)