Anomaly Simulation

AS.py

"""
# Anomaly Simulation for MVTecAD 
# Download describable textures dataset
wget https://www.robots.ox.ac.uk/~vgg/data/dtd/download/dtd-r1.0.1.tar.gz
# Download MVTec anomaly detection dataset
wget https://www.mydrive.ch/shares/38536/3830184030e49fe74747669442f0f282/download/420938113-1629952094/mvtec_anomaly_detection.tar.xz
"""

import os
import random
from glob import glob

import cv2
import imgaug.augmenters as iaa
import matplotlib.pyplot as plt
import numpy as np
from einops import rearrange

DTD_ROOT = "./data/dtd/images"
IMAGE_ROOT = "./data/mvtecad/images"
IMAGE_SIZE = (512, 512)
STRUCTRUE_GRID_SIZE = 8

texture_list = glob(os.path.join(DTD_ROOT, "*/*.jpg"))
texture_len = len(texture_list)
obj_list = [
    "capsule",
    "bottle",
    "carpet",
    "leather",
    "pill",
    "transistor",
    "tile",
    "cable",
    "zipper",
    "toothbrush",
    "metal_nut",
    "hazelnut",
    "screw",
    "grid",
    "wood",
]

# https://github.com/VitjanZ/DRAEM/blob/main/data_loader.py
augmenters = [
    iaa.GammaContrast((0.5, 2.0), per_channel=True),
    iaa.MultiplyAndAddToBrightness(mul=(0.8, 1.2), add=(-30, 30)),
    iaa.pillike.EnhanceSharpness(),
    iaa.AddToHueAndSaturation((-50, 50), per_channel=True),
    iaa.Solarize(0.5, threshold=(32, 128)),
    iaa.Posterize(),
    iaa.Invert(),
    iaa.pillike.Autocontrast(),
    iaa.pillike.Equalize(),
    iaa.Affine(rotate=(-45, 45)),
]


def interpolant(t):
    return t * t * t * (t * (t * 6 - 15) + 10)


def generate_perlin_noise_2d(
    shape, res, tileable=(False, False), interpolant=interpolant
):
    """Generate a 2D numpy array of perlin noise.
    Args:
        shape: The shape of the generated array (tuple of two ints).
            This must be a multple of res.
        res: The number of periods of noise to generate along each
            axis (tuple of two ints). Note shape must be a multiple of
            res.
        tileable: If the noise should be tileable along each axis
            (tuple of two bools). Defaults to (False, False).
        interpolant: The interpolation function, defaults to
            t*t*t*(t*(t*6 - 15) + 10).
    Returns:
        A numpy array of shape shape with the generated noise.
    Raises:
        ValueError: If shape is not a multiple of res.
    """
    delta = (res[0] / shape[0], res[1] / shape[1])
    d = (shape[0] // res[0], shape[1] // res[1])
    grid = np.mgrid[0 : res[0] : delta[0], 0 : res[1] : delta[1]].transpose(1, 2, 0) % 1
    # Gradients
    angles = 2 * np.pi * np.random.rand(res[0] + 1, res[1] + 1)
    gradients = np.dstack((np.cos(angles), np.sin(angles)))
    if tileable[0]:
        gradients[-1, :] = gradients[0, :]
    if tileable[1]:
        gradients[:, -1] = gradients[:, 0]
    gradients = gradients.repeat(d[0], 0).repeat(d[1], 1)
    g00 = gradients[: -d[0], : -d[1]]
    g10 = gradients[d[0] :, : -d[1]]
    g01 = gradients[: -d[0], d[1] :]
    g11 = gradients[d[0] :, d[1] :]
    # Ramps
    n00 = np.sum(np.dstack((grid[:, :, 0], grid[:, :, 1])) * g00, 2)
    n10 = np.sum(np.dstack((grid[:, :, 0] - 1, grid[:, :, 1])) * g10, 2)
    n01 = np.sum(np.dstack((grid[:, :, 0], grid[:, :, 1] - 1)) * g01, 2)
    n11 = np.sum(np.dstack((grid[:, :, 0] - 1, grid[:, :, 1] - 1)) * g11, 2)
    # Interpolation
    t = interpolant(grid)
    n0 = n00 * (1 - t[:, :, 0]) + t[:, :, 0] * n10
    n1 = n01 * (1 - t[:, :, 0]) + t[:, :, 0] * n11
    return np.sqrt(2) * ((1 - t[:, :, 1]) * n0 + t[:, :, 1] * n1)


def do_one_img(image_list, savename):
    img = cv2.imread(image_list[random.randint(0, len(image_list) - 1)])
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, dsize=IMAGE_SIZE)

    threshold = 50
    img_gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    white_pixel = len(img_gray[img_gray > threshold])
    all_pixel = img_gray.size
    if white_pixel / all_pixel > 0.5:
        _, mask_target_background = cv2.threshold(
            img_gray, threshold, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU
        )
    else:
        _, mask_target_background = cv2.threshold(
            img_gray, threshold, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU
        )

    mask_target_background = mask_target_background.astype(np.bool).astype(np.int)
    mask_target_foreground = -(mask_target_background - 1)

    min_perlin_scale = 0
    perlin_scale = 6
    perlin_scalex = 2 ** (random.randint(min_perlin_scale, perlin_scale))
    perlin_scaley = 2 ** (random.randint(min_perlin_scale, perlin_scale))
    perlin_noise = generate_perlin_noise_2d(IMAGE_SIZE, (perlin_scalex, perlin_scaley))

    rot = iaa.Affine(rotate=(-90, 90))
    perlin_noise = rot(image=perlin_noise)

    threshold = 0.5
    mask_noise = np.where(
        perlin_noise > threshold,
        np.ones_like(perlin_noise),
        np.zeros_like(perlin_noise),
    )

    mask = mask_noise * mask_target_foreground
    mask = np.expand_dims(mask, axis=2)

    texture_img = cv2.imread(texture_list[random.randint(0, texture_len) - 1])
    texture_img = cv2.cvtColor(texture_img, cv2.COLOR_BGR2RGB)
    texture_img = cv2.resize(texture_img, dsize=IMAGE_SIZE).astype(np.float32)
    aug_ind = np.random.choice(np.arange(len(augmenters)), 3, replace=False)
    aug = iaa.Sequential(
        [augmenters[aug_ind[0]], augmenters[aug_ind[1]], augmenters[aug_ind[2]]]
    )
    structure_source_img = aug(image=img)

    grid_w = IMAGE_SIZE[1] // STRUCTRUE_GRID_SIZE
    grid_h = IMAGE_SIZE[0] // STRUCTRUE_GRID_SIZE
    structure_source_img = rearrange(
        tensor=structure_source_img,
        pattern="(h gh) (w gw) c -> (h w) gw gh c",
        gw=grid_w,
        gh=grid_h,
    )
    disordered_idx = np.arange(structure_source_img.shape[0])
    np.random.shuffle(disordered_idx)

    structure_source_img = rearrange(
        tensor=structure_source_img[disordered_idx],
        pattern="(h w) gw gh c -> (h gh) (w gw) c",
        h=STRUCTRUE_GRID_SIZE,
        w=STRUCTRUE_GRID_SIZE,
    ).astype(np.float32)

    factor = np.random.uniform(0.15, 1, size=1)[0]

    texture_source_img = factor * (mask * texture_img) + (1 - factor) * (mask * img)
    structure_source_img = factor * (mask * structure_source_img) + (1 - factor) * (
        mask * img
    )
    texture_anomaly = ((-mask + 1) * img) + texture_source_img
    structure_anomaly = ((-mask + 1) * img) + structure_source_img
    out_list = [
        img,
        (mask_target_foreground * 255).astype(np.uint8)[:, :, None].repeat(3, 2),
        (perlin_noise * 255).astype(np.uint8)[:, :, None].repeat(3, 2),
        np.repeat((mask * 255).astype(np.uint8), 3, axis=-1),
        texture_img.astype(np.uint8),
        texture_source_img.astype(np.uint8),
        texture_anomaly.astype(np.uint8),
        structure_source_img.astype(np.uint8),
        structure_anomaly.astype(np.uint8),
    ]

    img_out = np.hstack(out_list)
    print(img_out.shape)
    img_out = cv2.cvtColor(img_out, cv2.COLOR_RGB2BGR)
    cv2.imwrite(savename, img_out)
    print(f"{savename} saved.")

    # out_list = [
    #     img,
    #     mask_target_foreground,
    #     perlin_noise,
    #     mask_noise,
    #     mask,
    #     texture_source_img.astype(np.uint8),
    #     structure_source_img.astype(np.uint8),
    #     texture_anomaly.astype(np.uint8),
    #     structure_anomaly.astype(np.uint8),
    # ]
    # fig = plt.figure()
    # plt.subplots_adjust(wspace=0.05, hspace=0)
    # for i in range(len(out_list)):
    #     plt.subplot(1, 10, i + 1)
    #     plt.imshow(out_list[i])
    #     plt.axis("off")
    # # plt.show()
    # plt.savefig(savename, bbox_inches="tight", pad_inches=0.0, dpi=200)
    # plt.close(fig)


def gen_anomaly(count=25):
    for obj in obj_list:
        root_path = f"./outputs/{obj}"
        os.makedirs(root_path, exist_ok=True)
        image_list = glob(os.path.join(IMAGE_ROOT, f"{obj}/*/good/*.png"))
        for idx in range(count):
            do_one_img(image_list, savename=os.path.join(root_path, f"{idx}.png"))


if __name__ == "__main__":
    gen_anomaly()