Spaces:
Sleeping
Sleeping
| # 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() | |