|
|
|
|
|
import gradio as gr |
|
from utils import initialize_gmm, generate_grid, generate_contours, generate_intermediate_points, plot_samples_and_contours |
|
import matplotlib.pyplot as plt |
|
import torch |
|
import numpy as np |
|
|
|
def validate_inputs(mu_list, Sigma_list, pi_list): |
|
try: |
|
mu = eval(mu_list) |
|
Sigma = eval(Sigma_list) |
|
pi = eval(pi_list) |
|
|
|
if not (isinstance(mu, list) and all(isinstance(i, list) for i in mu)): |
|
return False, "Mu list is invalid." |
|
if not (isinstance(Sigma, list) and all(isinstance(i, list) for i in Sigma)): |
|
return False, "Sigma list is invalid." |
|
if not isinstance(pi, list): |
|
return False, "Pi list is invalid." |
|
|
|
if not torch.isclose(torch.tensor(pi).sum(), torch.tensor(1.0)): |
|
return False, "Mixture weights must sum to 1." |
|
|
|
return True, "" |
|
except Exception as e: |
|
return False, str(e) |
|
|
|
def visualize_gmm(mu_list, Sigma_list, pi_list, dx, dtheta, T, N): |
|
is_valid, error_message = validate_inputs(mu_list, Sigma_list, pi_list) |
|
if not is_valid: |
|
fig, ax = plt.subplots() |
|
ax.text(0.5, 0.5, f'Invalid input: {error_message}', horizontalalignment='center', verticalalignment='center') |
|
ax.set_xlim(-5, 5) |
|
ax.set_ylim(-5, 5) |
|
ax.set_aspect('equal', adjustable='box') |
|
plt.close(fig) |
|
return fig, fig |
|
|
|
try: |
|
gmm = initialize_gmm(eval(mu_list), eval(Sigma_list), eval(pi_list)) |
|
grid_points = generate_grid(dx) |
|
std_normal_contours = generate_contours(dtheta) |
|
gmm_samples = gmm.sample(500) |
|
normal_samples = torch.distributions.MultivariateNormal(torch.zeros(2), torch.eye(2)).sample((500,)) |
|
(intermediate_points_gmm_to_normal, contour_intermediate_points_gmm_to_normal, grid_intermediate_points_gmm_to_normal, |
|
intermediate_points_normal_to_gmm, contour_intermediate_points_normal_to_gmm, grid_intermediate_points_normal_to_gmm) = \ |
|
generate_intermediate_points(gmm, grid_points, std_normal_contours, gmm_samples, normal_samples, T, N) |
|
|
|
final_frame_gmm_to_normal = intermediate_points_gmm_to_normal.cpu().detach().numpy() |
|
final_frame_normal_to_gmm = intermediate_points_normal_to_gmm.cpu().detach().numpy() |
|
|
|
fig1, ax1 = plot_samples_and_contours(final_frame_gmm_to_normal, contour_intermediate_points_gmm_to_normal.cpu().detach().numpy(), grid_intermediate_points_gmm_to_normal.cpu().detach().numpy(), "GMM to Normal Final Frame") |
|
fig2, ax2 = plot_samples_and_contours(final_frame_normal_to_gmm, contour_intermediate_points_normal_to_gmm.cpu().detach().numpy(), grid_intermediate_points_normal_to_gmm.cpu().detach().numpy(), "Normal to GMM Final Frame") |
|
|
|
return fig1, fig2 |
|
except Exception as e: |
|
fig, ax = plt.subplots() |
|
ax.text(0.5, 0.5, f'Error during visualization: {str(e)}', horizontalalignment='center', verticalalignment='center') |
|
ax.set_xlim(-5, 5) |
|
ax.set_ylim(-5, 5) |
|
ax.set_aspect('equal', adjustable='box') |
|
plt.close(fig) |
|
return fig, fig |
|
|
|
demo = gr.Interface( |
|
fn=visualize_gmm, |
|
inputs=[ |
|
gr.Textbox(label="Mu List", value="[[2, 1], [-1, -2], [3, -2]]", placeholder="Enter means as a list of lists, e.g., [[0,0], [1,1]]"), |
|
gr.Textbox(label="Sigma List", value="[[[0.2, 0.1], [0.1, 0.3]], [[1.0, -0.1], [-0.1, 0.1]], [[0.05, 0.0], [0.0, 0.05]]]", placeholder="Enter covariances as a list of lists, e.g., [[[0.2, 0.1], [0.1, 0.3]], [[1.0, -0.1], [-0.1, 0.1]]]"), |
|
gr.Textbox(label="Pi List", value="[0.05, 0.8, 0.15]", placeholder="Enter weights as a list, e.g., [0.5, 0.5]"), |
|
gr.Slider(minimum=0.01, maximum=1.0, label="dx", value=0.1), |
|
gr.Slider(minimum=2*np.pi/3600, maximum=2*np.pi/36, label="dtheta", value=2*np.pi/360), |
|
gr.Slider(minimum=1, maximum=100, label="T", value=10), |
|
gr.Slider(minimum=1, maximum=500, label="N", value=100) |
|
], |
|
outputs=[ |
|
gr.Plot(label="GMM to Normal Flow Final Frame"), |
|
gr.Plot(label="Normal to GMM Flow Final Frame") |
|
], |
|
live=True |
|
) |
|
|
|
demo.launch() |
|
|
