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()