Spaces:
Sleeping
Sleeping
import numpy as np | |
import matplotlib.pyplot as plt | |
import matplotlib.animation as animation | |
from PIL import Image | |
import cv2 | |
from math import tau | |
import gradio as gr | |
from concurrent.futures import ThreadPoolExecutor | |
import tempfile | |
def process_image(input_image, img_size, blur_kernel_size, desired_range): | |
img = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB2BGR) | |
img = cv2.resize(img, (img_size, img_size), interpolation=cv2.INTER_AREA) | |
imgray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) | |
blurred = cv2.GaussianBlur(imgray, (blur_kernel_size, blur_kernel_size), 0) | |
_, thresh = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU) | |
contours, _ = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) | |
largest_contour_idx = np.argmax([cv2.contourArea(c) for c in contours]) | |
largest_contour = contours[largest_contour_idx] | |
verts = [tuple(coord) for coord in largest_contour.squeeze()] | |
xs, ys = np.asarray(list(zip(*verts))) | |
x_range, y_range = np.max(xs) - np.min(xs), np.max(ys) - np.min(ys) | |
scale_x, scale_y = desired_range / x_range, desired_range / y_range | |
xs = (xs - np.mean(xs)) * scale_x | |
ys = -(ys - np.mean(ys)) * scale_y | |
return xs, ys | |
def compute_cn(f_exp, n, t_values): | |
coef = np.trapz(f_exp * np.exp(-n * t_values * 1j), t_values) / tau | |
return coef | |
def calculate_fourier_coefficients(xs, ys, num_points, coefficients): | |
t_list = np.linspace(0, tau, len(xs)) | |
t_values = np.linspace(0, tau, num_points) | |
f_precomputed = np.interp(t_values, t_list, xs + 1j * ys) | |
N = coefficients | |
indices = [0] + [j for i in range(1, N + 1) for j in (i, -i)] | |
with ThreadPoolExecutor(max_workers=8) as executor: | |
coefs = list(executor.map(lambda n: (compute_cn(f_precomputed, n, t_values), n), indices)) | |
return coefs | |
def setup_animation_env(img_size, desired_range, coefficients): | |
fig, ax = plt.subplots() | |
circles = [ax.plot([], [], 'b-')[0] for _ in range(-coefficients, coefficients + 1)] | |
circle_lines = [ax.plot([], [], 'g-')[0] for _ in range(-coefficients, coefficients + 1)] | |
drawing, = ax.plot([], [], 'r-', linewidth=2) | |
ax.set_xlim(-desired_range, desired_range) | |
ax.set_ylim(-desired_range, desired_range) | |
ax.set_axis_off() | |
ax.set_aspect('equal') | |
fig.set_size_inches(15, 15) | |
fig.canvas.draw() | |
background = fig.canvas.copy_from_bbox(ax.bbox) | |
return fig, ax, background, circles, circle_lines, drawing | |
def animate(frame, coefs, frame_times, fig, ax, background, circles, circle_lines, drawing, draw_x, draw_y, coefs_static, theta): | |
fig.canvas.restore_region(background) | |
center = (0, 0) | |
for idx, (r, fr) in enumerate(coefs_static): | |
c_dynamic = coefs[idx][0] * np.exp(1j * (fr * tau * frame_times[frame])) | |
x, y = center[0] + r * np.cos(theta[frame]), center[1] + r * np.sin(theta[frame]) | |
circle_lines[idx].set_data([center[0], center[0] + np.real(c_dynamic)], [center[1], center[1] + np.imag(c_dynamic)]) | |
circles[idx].set_data([x], [y]) | |
center = (center[0] + np.real(c_dynamic), center[1] + np.imag(c_dynamic)) | |
draw_x.append(center[0]) | |
draw_y.append(center[1]) | |
drawing.set_data(draw_x, draw_y) | |
for circle in circles: | |
ax.draw_artist(circle) | |
for line in circle_lines: | |
ax.draw_artist(line) | |
ax.draw_artist(drawing) | |
fig.canvas.blit(ax.bbox) | |
# Convert canvas to PIL Image using buffer_rgba | |
#fig.canvas.draw() | |
buffer = fig.canvas.buffer_rgba() | |
pil_image = Image.frombuffer("RGBA", fig.canvas.get_width_height(), buffer, "raw", "RGBA", 0, 1) | |
return (pil_image, None) | |
def generate_animation(frames, coefs, img_size, desired_range, coefficients): | |
fig, ax, background, circles, circle_lines, drawing = setup_animation_env(img_size, desired_range, coefficients) | |
coefs_static = [(np.linalg.norm(c), fr) for c, fr in coefs] | |
frame_times = np.linspace(0, 1, num=frames) | |
thetas = np.linspace(0, tau, num=frames) | |
draw_x, draw_y = [], [] | |
anim = animation.FuncAnimation(fig, animate, frames=frames, interval=5, fargs=(coefs, frame_times, fig, ax, background, circles, circle_lines, drawing, draw_x, draw_y, coefs_static, thetas)) | |
return fig, ax, background, circles, circle_lines, drawing, draw_x, draw_y, anim, frame_times, thetas, coefs_static | |
def fourier_transform_drawing(input_image, frames, coefficients, img_size, blur_kernel_size, desired_range, num_points): | |
xs, ys = process_image(input_image, img_size, blur_kernel_size, desired_range) | |
coefs = calculate_fourier_coefficients(xs, ys, num_points, coefficients) | |
# Setup animation environment | |
fig, ax, background, circles, circle_lines, drawing, draw_x, draw_y, anim, frame_times, thetas, coefs_static = generate_animation(frames, coefs, img_size, desired_range, coefficients) | |
# Create a temporary file for the video | |
with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as temp_file: | |
video_path = temp_file.name | |
# Generate and save each frame as a PIL image, and ultimately the video | |
for frame in range(frames): | |
pil_image, _ = animate(frame, coefs, time, fig, ax, background, circles, circle_lines, drawing, draw_x, draw_y, coefs_static, theta) | |
yield pil_image, video_path | |
# Save the animation as a video | |
anim = animation.FuncAnimation(fig, animate, frames=frames, interval=5, fargs=(coefs, time, fig, ax, background, circles, circle_lines, drawing, draw_x, draw_y, coefs_static, theta)) | |
anim.save(video_path, fps=15) | |
def setup_gradio_interface(): | |
interface = gr.Interface( | |
fn=fourier_transform_drawing, | |
inputs=[ | |
gr.Image(label="Drawing Progress", sources=['upload'], type="pil"), | |
gr.Slider(minimum=5, maximum=500, value=100, label="Number of Frames"), | |
gr.Slider(minimum=1, maximum=500, value=50, label="Number of Coefficients"), | |
gr.Number(value=224, label="Image Size (px)", precision=0), | |
gr.Slider(minimum=3, maximum=11, step=2, value=5, label="Blur Kernel Size (odd number)"), | |
gr.Number(value=400, label="Desired Range for Scaling", precision=0), | |
gr.Number(value=1000, label="Number of Points for Integration", precision=0), | |
], | |
outputs=["image", gr.Video()], | |
title="Fourier Transform Drawing", | |
description="Upload an image and generate a Fourier Transform drawing animation." | |
) | |
return interface | |
if __name__ == "__main__": | |
interface = setup_gradio_interface() | |
interface.queue() | |
interface.launch() |