Update app.py

app.py

@@ -2,16 +2,13 @@ import numpy as np
 import matplotlib.pyplot as plt
 import matplotlib.animation as animation
 from PIL import Image
-import io
-import os
 import cv2
 from math import tau
 import gradio as gr
 from concurrent.futures import ThreadPoolExecutor
 import tempfile
 
-def fourier_transform_drawing(input_image, frames, coefficients, img_size, blur_kernel_size, desired_range, num_points, theta_points):
-    # Convert PIL to OpenCV image
+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)
@@ -19,17 +16,8 @@ def fourier_transform_drawing(input_image, frames, coefficients, img_size, blur_
     _, thresh = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)
     contours, _ = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
 
-    # find the contour with the largest area
     largest_contour_idx = np.argmax([cv2.contourArea(c) for c in contours])
     largest_contour = contours[largest_contour_idx]
-    
-    # def combine_all_contours(contours):
-    #     combined_contour = np.array([], dtype=np.int32).reshape(0, 1, 2)
-    #     for contour in contours:
-    #         combined_contour = np.vstack((combined_contour, contour))
-    #     return combined_contour
-
-    # largest_contour = combine_all_contours(contours)
     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)
@@ -37,23 +25,28 @@ def fourier_transform_drawing(input_image, frames, coefficients, img_size, blur_
     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)
 
-    def compute_cn(f_exp, n, t_values):
-        coef = np.trapz(f_exp * np.exp(-n * t_values * 1j), t_values) / tau
-        return coef
-
     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))
+        coefs = list(executor.map(lambda n: compute_cn(f_precomputed, n, t_values), indices))
+    
+    return coefs
 
+def setup_animation_env(img_size, desired_range, coefficients):
     fig, ax = plt.subplots()
-    circles = [ax.plot([], [], 'b-')[0] for _ in range(-N, N + 1)]
-    circle_lines = [ax.plot([], [], 'g-')[0] for _ in range(-N, N + 1)]
+    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)
@@ -61,82 +54,75 @@ def fourier_transform_drawing(input_image, frames, coefficients, img_size, blur_
     ax.set_axis_off()
     ax.set_aspect('equal')
     fig.set_size_inches(15, 15)
-
-    draw_x, draw_y = [], []
-    theta = np.linspace(0, tau, theta_points)
-    coefs_static = [(np.linalg.norm(c), fr) for c, fr in coefs]
-    last_image = None 
-    
-    # Initialize the background
     fig.canvas.draw()
     background = fig.canvas.copy_from_bbox(ax.bbox)
-    
-    def animate(i, coefs, time, fig, ax, background, circles, circle_lines, drawing, draw_x, draw_y, coefs_static, theta):
-        # Restore the background to erase old frames
-        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 * time[i]))
-            x, y = center[0] + r * np.cos(theta), center[1] + r * np.sin(theta)
-            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[:i+1], draw_y[:i+1])
-    
-        # Draw only the updated elements
-        for circle in circles:
-            ax.draw_artist(circle)
-        for line in circle_lines:
-            ax.draw_artist(line)
-        ax.draw_artist(drawing)
-    
-        # Blit only the updated area
-        fig.canvas.blit(ax.bbox)
-    
-        # Capture the current canvas state as a PIL Image
-        canvas = fig.canvas
-        w, h = canvas.get_width_height()
-        buf = np.frombuffer(canvas.buffer_rgba(), dtype=np.uint8)
-        image = Image.fromarray(buf.reshape(h, w, 4), 'RGBA').convert('RGB')
-        last_image = image
-        
-        yield (last_image, None)
-
-    # Generate and yield images for each frame
+
+    return fig, ax, background, circles, circle_lines, drawing
+
+def animate(frame, coefs, time, 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 * time[frame]))
+        x, y = center[0] + r * np.cos(theta), center[1] + r * np.sin(theta)
+        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)
+
+def generate_animation(frames, coefs, img_size, desired_range, theta_points, 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]
     time = np.linspace(0, 1, num=frames)
-    for frame in range(frames):
-        yield from animate(frame, coefs, time, fig, ax, background, circles, circle_lines, drawing, draw_x, draw_y, coefs_static, theta)
+    theta = np.linspace(0, tau, theta_points)
+    draw_x, draw_y = [], []
 
-    # Generate final animation
+    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))
+
+    return anim
+
+def fourier_transform_drawing(input_image, frames, coefficients, img_size, blur_kernel_size, desired_range, num_points, theta_points):
+    xs, ys = process_image(input_image, img_size, blur_kernel_size, desired_range)
+    coefs = calculate_fourier_coefficients(xs, ys, num_points, coefficients)
+    anim = generate_animation(frames, coefs, img_size, desired_range, theta_points, coefficients)
+    
+    # Saving the animation
     with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as temp_file:
-        anim = animation.FuncAnimation(fig, animate, frames=frames, interval=5, fargs=(coefs, np.linspace(0, 1, num=frames)))
         anim.save(temp_file.name, fps=15)
+        return Image.fromarray(np.zeros((img_size, img_size), dtype=np.uint8)), temp_file.name
 
-    yield (last_image, temp_file.name)
-
-# Gradio interface setup
-interface = gr.Interface(
-    fn=fourier_transform_drawing,
-    inputs=[
-        gr.Image(label="Input Image", 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),
-        gr.Slider(minimum=50, maximum=500, value=80, label="Theta Points for Animation")
-    ],
-    outputs=["image", gr.Video()],
-    title="Fourier Transform Drawing",
-    description="Upload an image and generate a Fourier Transform drawing animation.",
-)
+def setup_gradio_interface():
+    interface = gr.Interface(
+        fn=fourier_transform_drawing,
+        inputs=[
+            gr.Image(label="Input Image", 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),
+            gr.Slider(minimum=50, maximum=500, value=80, label="Theta Points for Animation")
+        ],
+        outputs=["image", gr.Video()],
+        title="Fourier Transform Drawing",
+        description="Upload an image and generate a Fourier Transform drawing animation."
+    )
+    return interface
 
 if __name__ == "__main__":
-    # define queue - required for generators
+    interface = setup_gradio_interface()
     interface.queue()
     interface.launch()