Skip to content

Instantly share code, notes, and snippets.

@aleksejalex
Last active December 17, 2023 01:25
Show Gist options
  • Save aleksejalex/04fd96bb046ec9b2643912296b1ef005 to your computer and use it in GitHub Desktop.
Save aleksejalex/04fd96bb046ec9b2643912296b1ef005 to your computer and use it in GitHub Desktop.
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment