app.py

import os
import io
import cv2
import matplotlib.animation as animation
import matplotlib.pyplot as plt
import numpy as np
import gradio as gr
from scipy.integrate import quad_vec
from math import tau
from PIL import Image


def fourier_transform_drawing(input_image, frames, coefficients, img_size):
    # Convert PIL to OpenCV image(array)
    input_image = np.array(input_image)
    img = cv2.cvtColor(input_image, cv2.COLOR_RGB2BGR)

    # processing
    # resize the image to a smaller size for faster processing
    dim = (img_size, img_size)
    img = cv2.resize(img, dim, interpolation=cv2.INTER_AREA)
    
    imgray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    blurred = cv2.GaussianBlur(imgray, (5, 5), 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([len(c) for c in contours])
    verts = [tuple(coord) for coord in contours[largest_contour_idx].squeeze()]

    xs, ys = zip(*verts)

    # calculate the range of xs and ys
    x_range = np.max(xs) - np.min(xs)
    y_range = np.max(ys) - np.min(ys)
    
    # determine the scale factors
    desired_range = 400
    scale_factor_x = desired_range / x_range
    scale_factor_y = desired_range / y_range
    
    # apply scaling
    # ys needs to be flipped vertically
    xs = (np.asarray(xs) - np.mean(xs)) * scale_factor_x
    ys = (-np.asarray(ys) + np.mean(ys)) * scale_factor_y

    t_list = np.linspace(0, tau, len(xs))

    # compute the Fourier coefficients
    def f(t, t_list, xs, ys):
        return np.interp(t, t_list, xs + 1j*ys)

    def compute_cn(f, n):
        coef = 1/tau*quad_vec(
            lambda t: f(t, t_list, xs, ys)*np.exp(-n*t*1j), 
            0, 
            tau, 
            limit=100, 
            full_output=False)[0]
        return coef

    N = coefficients
    coefs = [(compute_cn(f, 0), 0)] + [(compute_cn(f, j), j) for i in range(1, N+1) for j in (i, -i)]

    # animate the drawings
    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)]
    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)

    draw_x, draw_y = [], []
    
    def animate(i, coefs, time):
        t = time[i]
        coefs = [(c * np.exp(1j*(fr * tau * t)), fr) for c, fr in coefs]
        center = (0, 0)

        for c, _ in coefs:
            r = np.linalg.norm(c)
            theta = np.linspace(0, tau, 80)
            x, y = center[0] + r * np.cos(theta), center[1] + r * np.sin(theta)
            circle_lines[_].set_data([center[0], center[0 ]+ np.real(c)], [center[1], center[1] + np.imag(c)])
            circles[_].set_data(x, y) 
            center = (center[0] + np.real(c), center[1] + np.imag(c))
        
        draw_x.append(center[0])
        draw_y.append(center[1])

        drawing.set_data(draw_x[:i+1], draw_y[:i+1])

    drawing_time = 1
    time = np.linspace(0, drawing_time, num=frames)    
    anim = animation.FuncAnimation(fig, animate, frames=frames, interval=5, fargs=(coefs, time)) 

    # save the animation as an MP4 file
    output_animation = "output.mp4"
    anim.save(output_animation, fps=15)
    plt.close(fig)

    # return the path to the MP4 file
    return output_animation

# 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", precision=0)
    ],
    outputs=gr.Video(),
    title="Fourier Transform Drawing",
    description="Upload an image and generate a Fourier Transform drawing animation. You can find out more about the project here : https://github.com/staghado/fourier-draw",
    examples=[["Fourier2.jpg", 100, 100], ["Luffy.png", 250, 300]]
)

if __name__ == "__main__":
    interface.launch()