app.py

import math
import os
import tempfile

import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
import trimesh
from scipy.interpolate import UnivariateSpline

LEN_SPLINE = 1000


def signal_to_spline(data, sr, signal_frame_length, spline_s, mult=1):
    data = data * mult
    data[data < 0] = 0
    data = np.abs(data)

    duration = len(data) / sr
    f_len = int(sr * signal_frame_length)
    frames = [data[i : i + f_len] for i in range(0, len(data), f_len)]

    residue = f_len - len(data) % f_len
    if residue:
        frames[-1] = np.pad(frames[-1], (0, f_len - len(data) % f_len))

    frame_max = np.max(frames, axis=1)
    frame_times = (np.arange(len(frame_max)) * f_len + 0.5 * f_len) / sr

    spline = UnivariateSpline(frame_times, frame_max, s=spline_s, ext=1)
    time_spline = np.linspace(0, duration, LEN_SPLINE)
    smooth_envelope = spline(time_spline)
    smooth_envelope[smooth_envelope < 0] = 0

    return smooth_envelope


def create_mesh_geometry(pos_env, neg_env, n_segs, z_scale, min_girth):
    pos_env[pos_env < min_girth] = min_girth
    neg_env[neg_env < min_girth] = min_girth

    n_rings = len(pos_env)
    theta = np.tile(np.linspace(0, np.pi * 2, n_segs, endpoint=False), (n_rings, 1))
    space_0 = np.linspace(pos_env, neg_env, n_segs // 2, endpoint=False).T
    space = np.concat([space_0, space_0[:, ::-1]], axis=1)

    x = np.cos(theta) * space
    y = np.sin(theta) * space
    z = np.tile(np.linspace(0, z_scale, n_rings, endpoint=False)[:, None], (1, n_segs))

    vertices = np.stack([x, y, z], axis=2).reshape(-1, 3)

    j = np.arange(n_segs)
    j_next = (j + 1) % n_segs
    i = np.arange(n_rings - 1)
    ring_offset = i[:, None] * n_segs

    curr_j = ring_offset + j
    curr_j_next = ring_offset + j_next
    next_j = ring_offset + n_segs + j
    next_j_next = ring_offset + n_segs + j_next

    faces1 = np.stack([curr_j, curr_j_next, next_j], axis=2)
    faces2 = np.stack([next_j, curr_j_next, next_j_next], axis=2)
    faces = np.concatenate([faces1, faces2], axis=0).reshape(-1, 3)

    return vertices, faces, n_segs, n_rings


def add_end_caps(vertices, faces, n_segs, n_rings):
    bottom_center = vertices[:n_segs].mean(axis=0)
    top_center = vertices[-n_segs:].mean(axis=0)

    vertices = np.vstack([vertices, bottom_center, top_center])
    bottom_center_idx = len(vertices) - 2
    top_center_idx = len(vertices) - 1

    bottom_ring_indices = np.arange(n_segs)
    bottom_ring_next = (bottom_ring_indices + 1) % n_segs
    bottom_faces = np.stack(
        [np.full(n_segs, bottom_center_idx), bottom_ring_next, bottom_ring_indices],
        axis=1,
    )

    top_ring_start = (n_rings - 1) * n_segs
    top_ring_indices = top_ring_start + np.arange(n_segs)
    top_ring_next = top_ring_start + (np.arange(n_segs) + 1) % n_segs
    top_faces = np.stack(
        [np.full(n_segs, top_center_idx), top_ring_indices, top_ring_next], axis=1
    )

    faces = np.vstack([faces, bottom_faces, top_faces])
    return vertices, faces


def create_3d_model_from_audio(
    audio_tuple, signal_frame_length, spline_s, n_segs, z_scale, y_scale, min_girth
):
    if audio_tuple is None:
        return None, None

    sr, data = audio_tuple

    if data.ndim > 1:
        data = np.mean(data, axis=1)

    data = data / np.max(data) * y_scale
    duration = len(data) / sr
    time = np.linspace(0, duration, len(data))
    time_spline = np.linspace(0, duration, LEN_SPLINE)

    pos_env = signal_to_spline(data, sr, signal_frame_length, spline_s)
    neg_env = signal_to_spline(data, sr, signal_frame_length, spline_s, mult=-1)

    fig_curve = plt.figure(figsize=(12, 3))
    plt.plot(time, data, alpha=0.3, color="blue")
    plt.plot(time_spline, pos_env, color="red")
    plt.plot(time_spline, -neg_env, color="red")
    plt.grid(True, linestyle="--", alpha=0.6)
    plt.tight_layout()

    vertices, faces, n_segs_actual, n_rings = create_mesh_geometry(
        pos_env, neg_env, n_segs, z_scale, min_girth
    )
    vertices, faces = add_end_caps(vertices, faces, n_segs_actual, n_rings)

    mesh = trimesh.Trimesh(vertices=vertices, faces=faces, process=True)
    mesh.fix_normals()
    mesh.apply_transform(
        trimesh.transformations.rotation_matrix(math.pi / 2, [1, 0, 0])
    )

    torus = trimesh.creation.torus(
        major_radius=min_girth * 1.5,
        minor_radius=min_girth / 2,
        major_sections=n_segs,
        minor_sections=n_segs,
    )
    torus.apply_translation([0, min_girth, 0])
    combined = trimesh.util.concatenate([mesh, torus])

    temp_dir = tempfile.gettempdir()
    output_path = os.path.join(temp_dir, "output_model.stl")
    combined.export(output_path, file_type="stl")

    return fig_curve, output_path


def create_interface():
    with gr.Blocks(theme=gr.themes.Soft()) as demo:
        gr.Markdown(
            """
        # 🔊 Audio to 3D Keychain Generator
        Upload a WAV audio file and adjust the parameters to generate a 3D keychain model.
        """
        )

        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### Input & Parameters")
                audio_input = gr.Audio(type="numpy", label="Upload Audio File")

                gr.Markdown("**Envelope Parameters**")
                signal_frame_length = gr.Slider(
                    0.01,
                    0.5,
                    value=0.1,
                    step=0.01,
                    label="Signal Frame Length (lower -> more details)",
                )
                spline_s = gr.Slider(
                    1e-5, 1e-2, value=1e-4, step=1e-5, label="Spline Smoothing Factor"
                )

                gr.Markdown("**3D Model Parameters**")
                n_segs = gr.Slider(
                    8, 128, value=32, step=1, label="Number of Segments (Resolution)"
                )
                z_scale = gr.Slider(
                    1.0, 20.0, value=5.0, step=0.5, label="Length (Z-Axis Scale)"
                )
                min_girth = gr.Slider(
                    0.0, 1.0, value=0.2, step=0.01, label="Minimum Girth (Radius)"
                )
                y_scale = gr.Slider(
                    0.1, 5.0, value=1.0, step=0.1, label="Radius (Y-Axis Scale)"
                )

                submit_btn = gr.Button("Generate Model", variant="primary")

            with gr.Column(scale=2):
                gr.Markdown("### Outputs")
                plot_curve = gr.Plot(label="Signal Envelope")
                model_3d = gr.Model3D(label="3D Model", camera_position=(0, -15, 5))

        submit_btn.click(
            fn=create_3d_model_from_audio,
            inputs=[
                audio_input,
                signal_frame_length,
                spline_s,
                n_segs,
                z_scale,
                y_scale,
                min_girth,
            ],
            outputs=[plot_curve, model_3d],
        )

        gr.Examples(
            [["voice.wav", 0.1, 1e-4, 64, 5.0, 0.2, 1.0]],
            inputs=[
                audio_input,
                signal_frame_length,
                spline_s,
                n_segs,
                z_scale,
                min_girth,
                y_scale,
            ],
            outputs=[plot_curve, model_3d],
            fn=create_3d_model_from_audio,
            cache_examples=False,  # Use False if you don't have a pre-existing voice.wav file
        )

    return demo


if __name__ == "__main__":
    demo = create_interface()
    demo.launch()