File size: 4,195 Bytes
7bc83e9 7b1ccbd 7bc83e9 7b1ccbd 7bc83e9 7b1ccbd 7bc83e9 7b1ccbd 7bc83e9 7b1ccbd 7bc83e9 7b1ccbd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
# app.py
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()
|