Script for creating occluded glasses and their masks

generate_overlays.py

# Used in: https://www.kaggle.com/datasets/mantasu/glasses-segmentation-synthetic-dataset
# Prerequisities: Portrait Eyeglasses and Shadows Removel Dataset: https://github.com/StoryMY/take-off-eyeglasses

import os
import cv2
import random
import numpy as np
import matplotlib.pyplot as plt

from PIL import Image
from tqdm import tqdm
from glasses_detector import FullGlassesSegmenter # may be updated in future
from scipy.ndimage import binary_dilation, binary_erosion
from skimage.morphology import disk


def generate_random_color():
    if np.random.random() < 0.5:
        # 50% of the time make black
        color = np.array([0, 0, 0])
    elif np.random.random() < 0.2:
        # Sometimes generate bright colors
        color = np.random.randint(70, 140, 3)
    else:
        # Mostly dark color is generated
        color = np.random.randint(0, 70, 3)

    if np.random.random() < 0.2:
        # Sometimes generate high alpha
        alpha = np.random.uniform(0.3, 0.5)
    else:
        # Most of the time, low alpha
        alpha = np.random.uniform(0.0, 0.3)

    return color, alpha

def generate_random_noise(glass_mask):
    # Initialize the noise image with 3 channels
    noise = np.zeros((*glass_mask.shape, 3), dtype=np.uint8)

    if np.random.rand() < 0.85:
        return noise

    # Choose a random type of noise
    noise_type = np.random.choice(["single_spot", "rings", "random_noise"])

    # Generate a random bright color for the noise
    color = np.random.randint(200, 256, 3)

    match noise_type:
        case "single_spot":
            # Add a single spot of noise
            for _ in range(100):
                x, y = np.random.randint(0, glass_mask.shape[1]), np.random.randint(0, glass_mask.shape[0])
                if glass_mask[y,x]:  # Ensure the noise is within the glass
                    cv2.circle(noise, (x, y), 5, color.tolist(), -1)
                    break

        case "rings":
            # Add sun-like rings of noise
            for _ in range(100):
                x, y = np.random.randint(0, glass_mask.shape[1]), np.random.randint(0, glass_mask.shape[0])
                if glass_mask[y,x]:  # Ensure the noise is within the glass
                    for i in range(10):
                        cv2.circle(noise, (x, y), i*5, color.tolist(), 1)
                    break

        case "random_noise":
            # Add random noise
            mask = np.random.random(glass_mask.shape) < 0.05
            mask &= glass_mask.astype(bool)  # Ensure the noise is within the glass
            for i in range(3):  # for each color channel
                noise[..., i][mask] = color[i]
    
    # Remove noise outside glasses
    noise[glass_mask == 0] = 0

    return noise


def overlay_glass(image, glass_frames_mask, frames_and_glass_mask):
    # Dilate and erode the glass_frames_mask to fill holes
    frames_new = binary_dilation(glass_frames_mask, structure=disk(6)).astype(int)
    glass_mask = np.maximum(0, frames_and_glass_mask - frames_new)
    glass_mask = binary_erosion(glass_mask, structure=disk(5), iterations=2).astype(int)
    glass_mask = binary_dilation(glass_mask, structure=disk(5), iterations=4).astype(int)

    # Generate a random color and alpha
    color, alpha = generate_random_color()

    # Create the glass with the given color and alpha
    glass = np.zeros_like(image)
    for i in range(3):  # for each channel
        glass[..., i] = glass_mask * color[i]
    
    # Overlay the glass on the original image using cv2.addWeighted
    overlayed_image = image.copy()
    overlayed_image[glass_mask > 0] = cv2.addWeighted(image, alpha, glass, 1 - alpha, 0)[glass_mask > 0]

    # Generate the noise
    noise = generate_random_noise(glass_mask)

    # Add the noise to the overlayed image
    for i in range(3):
        overlayed_image[..., i][noise[..., i] > 0] = noise[..., i][noise[..., i] > 0]

    return overlayed_image, np.minimum(1, glass_mask + glass_frames_mask)


def load_image_and_masks(image_path, segmenter):
    # Load the image, frames mask and generate full glases mask
    image = np.array(Image.open(image_path))
    glass_frames_mask = (np.array(Image.open(image_path.replace("-all", "-seg"))) > 0).astype(int)
    frames_and_glass_mask = segmenter.predict(image, mask_type="bool").astype(int)

    return image, glass_frames_mask, frames_and_glass_mask


def visualize(overlayed_image, frames_and_glass_mask):
    # Visualize the image with the overlayed glass and the full mask
    plt.figure(figsize=(10,5))
    plt.subplot(1,2,1)
    plt.imshow(frames_and_glass_mask)
    plt.title("Full Mask")
    plt.subplot(1,2,2)
    plt.imshow(overlayed_image)
    plt.title("Image with Overlayed Noised Glass")
    plt.show()


def process_dir(dir_path="data/segmentation/glass-frames/ALIGN_RESULT_v2"):
    # Read the paths to images and initialize the segmenter
    img_paths = [file.path for file in os.scandir(dir_path) if "-all" in file.name]
    segmenter = FullGlassesSegmenter(base_model="large", pretrained=True).eval() # may be updated in future
    segmenter.to("cuda:0")

    for img_path in tqdm(img_paths):
        # Load the image, frames and glasses masks; overlay generated glass on top of original image
        image, glass_frames_mask, frames_and_glass_mask = load_image_and_masks(img_path, segmenter)
        overlayed_image, frames_and_glass_mask = overlay_glass(image, glass_frames_mask, frames_and_glass_mask)

        # visualize(overlayed_image, frames_and_glass_mask)
        # break
        
        # Save the new image and the new full glasses mask
        Image.fromarray(overlayed_image).save(img_path.replace("-all", "-sunglasses"))
        Image.fromarray((frames_and_glass_mask * 255).astype(np.uint8), mode='L').save(img_path.replace("-all", "-sgseg"))


def train_val_test_split(dir_path="data/segmentation/glass-frames/ALIGN_RESULT_v2", val_size=0.14, test_size=0.14):
    # Retrieve the identity list and shuffle
    identities = [int(filename.split('-')[2]) for filename in os.listdir(dir_path) if "-all" in filename]
    identities = list(set(identities))
    random.shuffle(identities)

    # Create val and test identity sets
    num_val = round(len(identities) * val_size)
    num_test = round(len(identities) * test_size)
    val_identities = set(identities[:num_val])
    test_identities = set(identities[num_val:num_val+num_test])

    # Create train/val/test dirs
    os.makedirs(os.path.join(dir_path, "val"), exist_ok=True)
    os.makedirs(os.path.join(dir_path, "test"), exist_ok=True)
    os.makedirs(os.path.join(dir_path, "train"), exist_ok=True)

    for file in os.scandir(dir_path):
        if file.is_dir():
            continue
        
        # Check image identity
        identity = int(file.name.split('-')[2])

        # Move to correct split based on identity
        if identity in val_identities:
            os.rename(file.path, os.path.join(dir_path, "val", file.name))
        elif identity in test_identities:
            os.rename(file.path, os.path.join(dir_path, "test", file.name))
        else:
            os.rename(file.path, os.path.join(dir_path, "train", file.name))


if __name__ == "__main__":
    process_dir()
    train_val_test_split()