96 lines
3.3 KiB
Python
96 lines
3.3 KiB
Python
from PIL import Image, ImageOps
|
|
import random
|
|
import os
|
|
import numpy as np
|
|
|
|
RANDOM_SEED = 42
|
|
random.seed(RANDOM_SEED)
|
|
|
|
def resize_image(image, target_size):
|
|
original_width, original_height = image.size
|
|
target_width, target_height = target_size
|
|
|
|
# Calculate the aspect ratios
|
|
width_ratio = target_width / original_width
|
|
height_ratio = target_height / original_height
|
|
|
|
# Use the smaller ratio to ensure the image fits within the target dimensions
|
|
scaling_factor = min(width_ratio, height_ratio)
|
|
|
|
# Calculate new dimensions
|
|
new_width = int(original_width * scaling_factor)
|
|
new_height = int(original_height * scaling_factor)
|
|
|
|
# Resize the image
|
|
return image.resize((new_width, new_height), Image.LANCZOS)
|
|
|
|
def rotate_image(image, angle):
|
|
return image.rotate(angle, expand=True, resample=Image.Resampling.BICUBIC)
|
|
|
|
def tilt_image(image, max_tilt_angle=10):
|
|
width, height = image.size
|
|
|
|
# Simulate tilting by applying an affine transformation
|
|
tilt_angle = random.uniform(-max_tilt_angle, max_tilt_angle)
|
|
x_shift = width * np.sin(np.radians(tilt_angle))
|
|
|
|
# Affine transformation matrix
|
|
if tilt_angle > 0:
|
|
transform_matrix = (1, x_shift / height, -x_shift / 2, 0, 1, 0)
|
|
else:
|
|
transform_matrix = (1, x_shift / height, x_shift / 2, 0, 1, 0)
|
|
|
|
return image.transform(image.size, Image.AFFINE, transform_matrix, resample=Image.Resampling.BICUBIC)
|
|
|
|
def generate_images(object_image, output_folder, background_image_path,
|
|
background_width=960, background_height=720, num_images=50,
|
|
max_rotation_angle=45, max_tilt_angle=10):
|
|
|
|
background_image = Image.open(background_image_path).convert("RGB")
|
|
background_image = background_image.resize((background_width, background_height))
|
|
|
|
for i in range(num_images):
|
|
# Create a copy of the background image
|
|
background = background_image.copy()
|
|
|
|
# Randomly rotate the object image
|
|
rotation_angle = random.uniform(-max_rotation_angle, max_rotation_angle)
|
|
rotated_image = rotate_image(object_image, rotation_angle)
|
|
|
|
# Apply random tilt to the object image
|
|
tilted_image = tilt_image(rotated_image, max_tilt_angle)
|
|
|
|
# Random position for the object
|
|
max_x = background_width - tilted_image.width
|
|
max_y = background_height - tilted_image.height
|
|
x = random.randint(0, max_x)
|
|
y = random.randint(0, max_y)
|
|
|
|
# Paste the object image onto the background
|
|
background.paste(tilted_image, (x, y), tilted_image)
|
|
|
|
# Save the new image
|
|
background.save(f"{output_folder}/variation_{i + 1}.png")
|
|
|
|
def main():
|
|
background_width = 960
|
|
background_height = 720
|
|
new_object_size = (300, 300)
|
|
|
|
os.makedirs("gen_data", exist_ok=True)
|
|
output_folder = "gen_data/"
|
|
|
|
# Load the object image (with transparency)
|
|
object_image = Image.open("img/drone.png").convert("RGBA")
|
|
object_image = resize_image(object_image, new_object_size)
|
|
|
|
# Load the background image from a second source
|
|
background_image_path = "img/background.jpg"
|
|
|
|
# Generate images with object placed on top of the background
|
|
generate_images(object_image, output_folder, background_image_path,
|
|
background_width, background_height, num_images=50)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|