Update app.py

app.py

@@ -1,49 +1,40 @@
-import os
-import io
-import cv2
 import numpy as np
 import matplotlib.pyplot as plt
 import matplotlib.animation as animation
 from PIL import Image
+import io
+import cv2
 from math import tau
-from concurrent.futures import ThreadPoolExecutor
 import gradio as gr
+from concurrent.futures import ThreadPoolExecutor
 
 def fourier_transform_drawing(input_image, frames, coefficients, img_size, blur_kernel_size, desired_range, num_points, theta_points):
     # Convert PIL to OpenCV image
     img = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB2BGR)
-
-    # Resize the image for faster processing
     img = cv2.resize(img, (img_size, img_size), interpolation=cv2.INTER_AREA)
-
-    # Image processing
     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)
 
     # 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]
+    largest_contour_idx = np.argmax([cv2.contourArea(c) for c in contours])
+    largest_contour = contours[largest_contour_idx]
     
-    # Combine all contours
-    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
-
-    combined_contour = combine_all_contours(contours)
-    verts = [tuple(coord) for coord in combined_contour.squeeze()]
+    # 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)))
-
-    # Scale the coordinates
     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
 
-    # Compute Fourier 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)
@@ -55,11 +46,9 @@ def fourier_transform_drawing(input_image, frames, coefficients, img_size, blur_
     N = coefficients
     indices = [0] + [j for i in range(1, N + 1) for j in (i, -i)]
 
-    # Parallel computation of coefficients
     with ThreadPoolExecutor(max_workers=8) as executor:
         coefs = list(executor.map(lambda n: (compute_cn(f_precomputed, n, t_values), n), indices))
 
-    # Animation setup
     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)]
@@ -75,26 +64,35 @@ def fourier_transform_drawing(input_image, frames, coefficients, img_size, blur_
     theta = np.linspace(0, tau, theta_points)
     coefs_static = [(np.linalg.norm(c), fr) for c, fr in coefs]
 
-    # Animation function
     def animate(i, coefs, time):
         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) 
+            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])
 
-    # Create and save the animation
+        # Capture and yield the current plot as an image
+        buf = io.BytesIO()
+        plt.savefig(buf, format='png', bbox_inches='tight')
+        buf.seek(0)
+        yield np.array(Image.open(buf))
+
+    # Generate and yield images
+    for frame in range(frames):
+        yield from animate(frame, coefs, np.linspace(0, 1, num=frames))
+
+    # Generate final animation
     anim = animation.FuncAnimation(fig, animate, frames=frames, interval=5, fargs=(coefs, np.linspace(0, 1, num=frames)))
-    output_animation = "output.mp4"
-    anim.save(output_animation, fps=15)
-    plt.close(fig)
-    return output_animation
+    buf = io.BytesIO()
+    anim.save(buf, format='gif', fps=15)
+    buf.seek(0)
+    yield buf
 
 # Gradio interface setup
 interface = gr.Interface(
@@ -109,10 +107,9 @@ interface = gr.Interface(
         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=gr.Video(),
+    outputs=["image", "file"],
     title="Fourier Transform Drawing",
     description="Upload an image and generate a Fourier Transform drawing animation.",
-    examples=[["Fourier2.jpg", 100, 200, 224, 5, 400, 1000, 80], ["Luffy.png", 100, 100, 224, 5, 400, 1000, 80]]
 )
 
 if __name__ == "__main__":