img_manipulations.py

img_manipulations.py

"""
just selected python functions for EMERGENT project purposes
...to be iported in IPYNB notebook

(C) aleksejalex, SteveLikesToDebug
"""

import numpy as np

# imports for images
from PIL import Image  # PIL
import cv2  # opencv-python aka cv2

from tqdm import tqdm
import os
import sys
import cv2
import random

import zipfile
import urllib.request

from barcode import EAN13
from barcode.writer import ImageWriter

from pyzbar.pyzbar import decode


def zip_directory(directory_name):
    # Get the absolute path of the directory
    abs_dir_path = os.path.abspath(directory_name)

    # Make sure the directory exists
    if not os.path.exists(abs_dir_path):
        raise FileNotFoundError(f"Directory '{directory_name}' not found.")

    # Get the base name of the directory
    base_dir_name = os.path.basename(abs_dir_path)

    # Create a zip file in the same directory
    zip_file_name = f"{base_dir_name}.zip"
    zip_file_path = os.path.join(abs_dir_path, zip_file_name)

    with zipfile.ZipFile(zip_file_path, 'w') as zip_file:
        # Walk through the directory and add all files to the zip file
        for foldername, subfolders, filenames in os.walk(abs_dir_path):
            for filename in filenames:
                file_path = os.path.join(foldername, filename)
                arcname = os.path.relpath(file_path, abs_dir_path)
                zip_file.write(file_path, arcname)

    print(f"Directory '{directory_name}' zipped successfully to '{zip_file_name}'.")


def get_and_unzip(zip_file_url):
    # Get the file name from the URL
    file_name = os.path.basename(zip_file_url)

    # Download the file
    urllib.request.urlretrieve(zip_file_url, file_name)

    # Extract the contents of the zip file
    with zipfile.ZipFile(file_name, 'r') as zip_ref:
        zip_ref.extractall()

    print(f"File '{file_name}' downloaded and extracted successfully.")


def decode_barcode(img=None):
    """
    function that simply reads the barcode from image (should be robust to cases when
    several codes are in single image (todo test this))
    :param img: path to image as 'str' OR 'cv2.Image'
    :param img_show: bool = to show detected code in the picture or not
    :return: string = code detected.
    """
    # Load the image
    if img is None:
        raise Exception("detect_barcode has not received any input image. Abort.")
    if type(img) is str:
        image = cv2.imread(img)
    else:
        image = img.copy()  # already in cv2 image type

    barcode_data = " "
    # Decode barcodes in the image
    barcodes = decode(image)

    # Iterate through the barcodes found
    for barcode in barcodes:
        barcode_data = barcode.data.decode('utf-8')
        barcode_type = barcode.type
        x1, y1, x2, y2 = barcode.rect  # Get the coordinates of the barcode

        # Print the barcode data and type
        print(f"Data: {barcode_data}, Type: {barcode_type}")

        # Draw a rectangle around the barcode on the image
        cv2.rectangle(image, (x1, y1), (x1 + x2, y1 + y2), (0, 255, 0), 2)

    if barcode_data == " ":
        print(f'Error: no code recognized.')
    return barcode_data


def generate_barcode(ean: str, save_png: bool = False, filename: str = None):
    """
    EAN13 = 13 integers, but here i needs to be passed 12 integers, the last one will be added according to the requirements of EAN standard.
    """
    if len(ean) != 12:
        raise Exception(f"generate_barcode: can't proceed, given ean has {len(ean)} chars instead of 12. ")
    code = EAN13(ean, writer=ImageWriter())
    code_image = code.render()
    if save_png:
        code.save(filename)  # '.png' is added automatically

    barcode_in_cv2 = cv2.cvtColor(np.array(code_image), cv2.COLOR_RGB2BGR)  # Convert PIL Image to cv2 format
    return barcode_in_cv2


def generate_random_intstr(seed, num_of_digits):
    random.seed(seed)  # Set the seed for reproducibility
    return ''.join(map(str, [random.randint(0, 9) for _ in range(num_of_digits)]))


def generate_random_angle(seed):
    random.seed(seed)
    return random.randint(0, 360)


def rotate_image(image, angle):
    # Get image dimensions
    height, width = image.shape[:2]

    # Calculate the rotation matrix
    rotation_matrix = cv2.getRotationMatrix2D((width / 2, height / 2), angle, 1)

    # Calculate the bounding box for the rotated image
    cos_theta = np.abs(rotation_matrix[0, 0])
    sin_theta = np.abs(rotation_matrix[0, 1])

    new_width = int((width * cos_theta) + (height * sin_theta))
    new_height = int((width * sin_theta) + (height * cos_theta))

    rotation_matrix[0, 2] += (new_width / 2) - (width / 2)
    rotation_matrix[1, 2] += (new_height / 2) - (height / 2)

    # Apply rotation to the image without cropping
    rotated_image = cv2.warpAffine(image, rotation_matrix, (new_width, new_height), flags=cv2.INTER_LINEAR)

    return rotated_image


def upscale_image(original_image, scale_factor=4):
    # Upscale the image
    upscaled_image = cv2.resize(original_image, None, fx=scale_factor, fy=scale_factor, interpolation=cv2.INTER_LINEAR)

    return upscaled_image


def resize_image(input_img, new_width=None, new_height=None):
    # Get the original dimensions of the image
    original_height, original_width = input_img.shape[:2]

    # Calculate the aspect ratio
    aspect_ratio = original_width / float(original_height)

    # If new_width is specified, calculate new_height to maintain aspect ratio
    if new_width is not None:
        new_height = int(new_width / aspect_ratio)

    # If new_height is specified, calculate new_width to maintain aspect ratio
    elif new_height is not None:
        new_width = int(new_height * aspect_ratio)

    # Resize the image using the calculated dimensions
    resized_img = cv2.resize(input_img, (new_width, new_height))

    # Save the resized image
    return resized_img


def overlay_images(img1, img2, transparency=0.5):
    # Ensure both images have the same data type
    img1 = img1.astype(np.float32)
    img2 = img2.astype(np.float32)

    # Get the dimensions of the images
    h1, w1, _ = img1.shape
    h2, w2, _ = img2.shape

    # Check if img1 is larger than img2, and resize if necessary
    if h1 > h2 or w1 > w2:
        img1 = cv2.resize(img1, (w1 // 2, h1 // 2))

    # Get the new dimensions of img1
    h1, w1, _ = img1.shape

    # Generate random position for img1 within img2
    x_position = np.random.randint(0, w2 - w1)
    y_position = np.random.randint(0, h2 - h1)

    # Create a mask for img1
    img1_mask = (img1 != 0)

    # Copy img2 to a new variable to avoid modifying the original image
    result_img = img2.copy()

    # Apply img1 to the selected position on img2 with transparency
    result_img[y_position:y_position + h1, x_position:x_position + w1, :] *= 1 - transparency
    result_img[y_position:y_position + h1, x_position:x_position + w1, :] += img1 * transparency

    return result_img.astype(np.uint8)


def add_realistic_noise(image, gaussian_mean=0, gaussian_sigma=25, poisson_scale=0.02):
    """
    Add realistic noise (Gaussian and Poisson) to the input image.

    Parameters:
        image (numpy.ndarray): Input image (BGR or grayscale).
        gaussian_mean (float): Mean of the Gaussian distribution for Gaussian noise.
        gaussian_sigma (float): Standard deviation of the Gaussian distribution for Gaussian noise.
        poisson_scale (float): Scale parameter for Poisson noise.

    Returns:
        numpy.ndarray: Noisy image.
    """
    # Generate Gaussian noise
    gaussian_noise = np.random.normal(gaussian_mean, gaussian_sigma, image.shape)

    # Add Gaussian noise to the image
    noisy_image = np.clip(image + gaussian_noise, 0, 255).astype(np.uint8)

    # Generate Poisson noise
    poisson_noise = np.random.poisson(poisson_scale, image.shape).astype(np.uint8)  # Convert to uint8

    # Add Poisson noise to the image
    noisy_image += poisson_noise
    noisy_image = np.clip(noisy_image, 0, 255).astype(np.uint8)

    return noisy_image


def generate_dataset(trainval_ratio: float = 0.8, dir_name: str = None, path_to_backgrounds: str = None,
                     num_bc: int = 3, num_rot: int = 5, num_bcgr: int = 1, num_placings: int = 4, num_noises: int = 2,
                     noise_strength: int = 40):
    """
    tba
    """
    # prepare dirs
    os.makedirs(f"{dir_name}/train/input")
    os.makedirs(f"{dir_name}/train/target")
    os.makedirs(f"{dir_name}/val/input")
    os.makedirs(f"{dir_name}/val/target")
    # how many imgs we expect to generate and when to stop saving them as training set and start saving as validation set.
    total_num_of_runs_to_expect = num_bc * num_rot * num_bcgr * num_placings * num_noises
    split_train_and_val = round(trainval_ratio * total_num_of_runs_to_expect)
    # in for cycle - generate a barcode, save it in target, then damage it, save to input
    naming_idx = 1
    for j in range(num_bc):  # different EANs
        fake_ean = generate_random_intstr(seed=42 + j, num_of_digits=12)
        good_bc_in_cv = generate_barcode(fake_ean)
        # good_bc_in_cv = resize_image(good_bc_in_cv0, new_width=500)
        for jj in range(num_rot):  # rotations
            rot_bc0 = rotate_image(good_bc_in_cv, 0)  # generate_random_angle(seed=jj))
            print(
                f"currently rendering img_{naming_idx} out of {num_bc * num_rot * num_bcgr * num_placings * num_noises}. Please wait...")
            for jjj in range(num_bcgr):  # background pictures
                background = cv2.imread("trojanka.png")
                for jjjj in range(num_placings):  # placing on background - different positions
                    # placed_bc = place_image_inside_bigger(rot_bc, background_image=background, random_position=True, specified_coordinates=None,
                    #           background_color=(0, 0, 0), transparency=0.5)
                    rot_bc = resize_image(rot_bc0, new_width=500)
                    placed_bc = overlay_images(rot_bc, box, transparency=0.5)
                    for jjjjj in range(num_noises):
                        noisy_bc = add_realistic_noise(placed_bc, gaussian_mean=0, gaussian_sigma=noise_strength,
                                                       poisson_scale=0.03)
                        # to black and white:
                        gray_noisy_bc = cv2.cvtColor(noisy_bc, cv2.COLOR_BGR2GRAY)
                        gray_good_bc_in_cv = cv2.cvtColor(good_bc_in_cv, cv2.COLOR_BGR2GRAY)
                        # saving:
                        if naming_idx < split_train_and_val:
                            path_to_save_good_bc = f"{dir_name}/train/target/img_{naming_idx}.png"
                            path_to_save_damaged_bc = f"{dir_name}/train/input/img_{naming_idx}.png"
                        else:
                            path_to_save_good_bc = f"{dir_name}/val/target/img_{naming_idx}.png"
                            path_to_save_damaged_bc = f"{dir_name}/val/input/img_{naming_idx}.png"
                        resized_input_img = resize_and_crop_image(gray_noisy_bc, desired_height=200, desired_width=200)
                        resized_target_img = resize_and_crop_image(gray_good_bc_in_cv, desired_height=200,
                                                                   desired_width=200)
                        cv2.imwrite(path_to_save_good_bc, resized_target_img)
                        cv2.imwrite(path_to_save_damaged_bc, resized_input_img)
                        naming_idx = naming_idx + 1
    print(f'Done.')

    
def resize_and_crop_image(input_image, desired_height, desired_width):
    # Read the input image
    #img = cv2.imread(input_image)
    img = input_image.copy()

    # Get the original image dimensions
    original_height, original_width = img.shape[:2]

    # Calculate the aspect ratios
    aspect_ratio_original = original_width / original_height
    aspect_ratio_desired = desired_width / desired_height

    # Calculate the new dimensions for resizing
    if aspect_ratio_original > aspect_ratio_desired:
        new_width = int(desired_width)
        new_height = int(desired_width / aspect_ratio_original)
    else:
        new_width = int(desired_height * aspect_ratio_original)
        new_height = int(desired_height)

    # Resize the image while maintaining the aspect ratio
    img_resized = cv2.resize(img, (new_width, new_height))

    # Calculate the cropping parameters
    crop_start_x = max(0, int((new_width - desired_width) / 2))
    crop_start_y = max(0, int((new_height - desired_height) / 2))

    # Crop the image
    img_cropped = img_resized[crop_start_y:crop_start_y + desired_height, crop_start_x:crop_start_x + desired_width]

    return img_cropped