70 lines
2.4 KiB
Python
70 lines
2.4 KiB
Python
import numpy as np
|
|
import cv2
|
|
import argparse
|
|
import time
|
|
from numpy import random
|
|
|
|
def load_image(image_path):
|
|
image = cv2.imread(image_path, cv2.IMREAD_COLOR)
|
|
if image is None:
|
|
raise ValueError(f"Image at path {image_path} could not be loaded.")
|
|
return image.astype(np.float32)
|
|
|
|
def draw_circle(image, reference_image, center, radius):
|
|
mask = np.zeros(image.shape[:2], dtype=np.uint8)
|
|
cv2.circle(mask, center, int(radius), 255, -1)
|
|
|
|
mean_color = cv2.mean(reference_image, mask=mask)[:3]
|
|
mean_color = tuple(map(int, mean_color))
|
|
|
|
tmp_image = image.copy()
|
|
cv2.circle(tmp_image, center, int(radius), mean_color, -1)
|
|
return tmp_image
|
|
|
|
def calculate_vector_distance(image1, image2):
|
|
return np.linalg.norm(image1 - image2)
|
|
|
|
def main(in_path, out_path, iterations, seed):
|
|
random.seed(seed)
|
|
image1 = load_image(in_path)
|
|
height, width, _ = image1.shape
|
|
image2 = np.zeros((height,width,3), np.uint8)
|
|
|
|
last_time = time.time()
|
|
|
|
for i in range(iterations):
|
|
print(f"Image {in_path} pass {i}", end="\r")
|
|
|
|
center = (random.randint(width), random.randint(height))
|
|
radius = max(width, height) / 2
|
|
step = radius
|
|
|
|
for _ in range(int(np.log2(max(width, height))) + 1):
|
|
candidate_up = draw_circle(image2, image1, center, radius + step)
|
|
candidate_up_distance = calculate_vector_distance(image1, candidate_up)
|
|
|
|
candidate_down = draw_circle(image2, image1, center, max(0, radius - step))
|
|
candidate_down_distance = calculate_vector_distance(image1, candidate_down)
|
|
|
|
if candidate_down_distance <= candidate_up_distance:
|
|
radius = max(0, radius - step)
|
|
else:
|
|
radius += step
|
|
|
|
step /= 2
|
|
|
|
image2 = draw_circle(image2, image1, center, radius)
|
|
|
|
print(f"{out_path} took {-last_time + time.time()} seconds")
|
|
cv2.imwrite(f"{out_path}.png", np.uint8(image2))
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Draw and optimize a circle on an image using gradient ascent.")
|
|
parser.add_argument("input", help="Path to the reference image.")
|
|
parser.add_argument("output", help="Path to save the output image.")
|
|
parser.add_argument("circles", help="The amount of circles to draw.")
|
|
parser.add_argument("seed", help="The seed to use.")
|
|
|
|
args = parser.parse_args()
|
|
main(args.input, args.output, int(args.circles), int(args.seed))
|