mdbloice / Augmentor

Image augmentation library in Python for machine learning.
https://augmentor.readthedocs.io/en/stable
MIT License
5.08k stars 866 forks source link

p.ground_truth not work for p.add_operation #240

Open ccl-private opened 3 years ago

ccl-private commented 3 years ago

custom functions not work in mask labels. image

ccl-private commented 3 years ago

I used "np.random" in my custom function. how can I convert the same random value to image and mask image.

# Create your new operation by inheriting from the Operation superclass:
class ZoomPadding(Operation):
    # Here you can accept as many custom parameters as required:
    def __init__(self, probability, min_factor=0.5, max_factor=1.0, desired_size=512):
        # Call the superclass's constructor (meaning you must
        # supply a probability value):
        Operation.__init__(self, probability)
        # Set your custom operation's member variables here as required:
        self.desired_size = desired_size
        self.min_factor = min_factor
        self.max_factor = max_factor

    # Your class must implement the perform_operation method:
    def perform_operation(self, images):
        # Start of code to perform custom image operation.
        for i in range(len(images)):
            image = images[i]
            image_numpy = np.array(image).astype('uint8')
            # im = cv2.cvtColor(image_numpy, cv2.COLOR_RGB2BGR)
            im = image_numpy
            old_size = im.shape[:2]  # old_size is in (height, width) format

            back_ground = np.zeros_like(im)
            back_ground = cv2.resize(back_ground, (self.desired_size, self.desired_size))

            ratio = np.random.random() * (self.max_factor - self.min_factor) + self.min_factor
            new_size = [int(x * ratio) for x in old_size]

            if new_size[0] > self.desired_size:
                new_size[0] = self.desired_size

            if new_size[1] > self.desired_size:
                new_size[1] = self.desired_size

            # new_size should be in (width, height) format

            im = cv2.resize(im, (new_size[1], new_size[0]))

            top = np.random.randint(low=0, high=int(self.desired_size - new_size[1]))
            left = np.random.randint(low=0, high=int(self.desired_size - new_size[0]))

            back_ground[top:top + new_size[1], left:left + new_size[0]] = im

            new_im = back_ground

            image = Image.fromarray(new_im)
            images[i] = image
        # End of code to perform custom image operation.

        # Return the image so that it can further processed in the pipeline:
        return images
ccl-private commented 3 years ago

even when I remove the random value codes, the custom function still does not work in the mask image.

ccl-private commented 3 years ago

here is all of my codes.

import numpy as np
import Augmentor
from Augmentor.Pipeline import Operation
import cv2
from PIL import Image
import os
import glob

# Create your new operation by inheriting from the Operation superclass:
class ResizeCenterPadding(Operation):
    # Here you can accept as many custom parameters as required:
    def __init__(self, probability, desired_size):
        # Call the superclass's constructor (meaning you must
        # supply a probability value):
        Operation.__init__(self, probability)
        # Set your custom operation's member variables here as required:
        self.desired_size = desired_size

    # Your class must implement the perform_operation method:
    def perform_operation(self, images):
        # Start of code to perform custom image operation.
        for i in range(len(images)):
            image = images[i]
            image_numpy = np.array(image).astype('uint8')
            # im = cv2.cvtColor(image_numpy, cv2.COLOR_RGB2BGR)
            im = image_numpy
            old_size = im.shape[:2]  # old_size is in (height, width) format

            ratio = float(self.desired_size) / max(old_size)
            new_size = tuple([int(x * ratio) for x in old_size])

            # new_size should be in (width, height) format

            im = cv2.resize(im, (new_size[1], new_size[0]))

            delta_w = self.desired_size - new_size[1]
            delta_h = self.desired_size - new_size[0]
            top, bottom = delta_h // 2, delta_h - (delta_h // 2)
            left, right = delta_w // 2, delta_w - (delta_w // 2)

            color = [0, 0, 0]
            new_im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT,
                                        value=color)
            image = Image.fromarray(new_im)
            images[i] = image
        # End of code to perform custom image operation.

        # Return the image so that it can further processed in the pipeline:
        return images

# Create your new operation by inheriting from the Operation superclass:
class ZoomPadding(Operation):
    # Here you can accept as many custom parameters as required:
    def __init__(self, probability, min_factor=0.5, max_factor=1.0, desired_size=512):
        # Call the superclass's constructor (meaning you must
        # supply a probability value):
        Operation.__init__(self, probability)
        # Set your custom operation's member variables here as required:
        self.desired_size = desired_size
        self.min_factor = min_factor
        self.max_factor = max_factor

    # Your class must implement the perform_operation method:
    def perform_operation(self, images):
        # Start of code to perform custom image operation.
        for i in range(len(images)):
            image = images[i]
            image_numpy = np.array(image).astype('uint8')
            # im = cv2.cvtColor(image_numpy, cv2.COLOR_RGB2BGR)
            im = image_numpy
            old_size = im.shape[:2]  # old_size is in (height, width) format

            back_ground = np.zeros_like(im)
            back_ground = cv2.resize(back_ground, (self.desired_size, self.desired_size))

            ratio = np.random.random() * (self.max_factor - self.min_factor) + self.min_factor
            new_size = [int(x * ratio) for x in old_size]

            if new_size[0] > self.desired_size:
                new_size[0] = self.desired_size

            if new_size[1] > self.desired_size:
                new_size[1] = self.desired_size

            # new_size should be in (width, height) format

            im = cv2.resize(im, (new_size[1], new_size[0]))

            top = np.random.randint(low=0, high=int(self.desired_size - new_size[1]))
            left = np.random.randint(low=0, high=int(self.desired_size - new_size[0]))

            back_ground[top:top + new_size[1], left:left + new_size[0]] = im

            new_im = back_ground

            image = Image.fromarray(new_im)
            images[i] = image
        # End of code to perform custom image operation.

        # Return the image so that it can further processed in the pipeline:
        return images

path = '/data/ccl/linping_ccl/split_house/train/'
path_new = '/data/ccl/linping_ccl/split_house/train_augmented'
mask_path = "/data/ccl/linping_ccl/split_house/trainannot/"
mask_path_new = '/data/ccl/linping_ccl/split_house/trainannot_augmented'
desired_size = 512
pics_num = 1000

if os.path.exists(path_new):
    os.system("rm -r " + path_new)
if not os.path.exists(path_new):
    os.makedirs(path_new)

p = Augmentor.Pipeline(source_directory=path,
                       output_directory=os.path.join(path_new, ''))
p.ground_truth(mask_path)

p.rotate_without_crop(probability=0.7, max_left_rotation=45, max_right_rotation=45, expand=True)
p.flip_left_right(probability=0.5)
p.flip_left_right(probability=0.5)
# p.random_erasing(probability=0.5, rectangle_area=0.5)
p.add_operation(ResizeCenterPadding(probability=1, desired_size=desired_size))
p.add_operation(ZoomPadding(probability=0.5, min_factor=0.5, max_factor=1.0, desired_size=desired_size))
try:
    p.sample(int(pics_num))
except Exception as e:
    print(e)