|
import gradio as gr |
|
import numpy as np |
|
import matplotlib as mpl |
|
mpl.use('agg') |
|
import matplotlib.pyplot as plt |
|
import torch |
|
import torch.nn as nn |
|
import torch.optim as optim |
|
import torch.nn.functional as F |
|
from torch.utils.data import TensorDataset, DataLoader |
|
from sklearn.decomposition import PCA |
|
from sklearn.cluster import KMeans |
|
from sklearn.manifold import TSNE |
|
from umap import UMAP |
|
import plotly.express as px |
|
import pandas as pd |
|
|
|
class recon_encoder(nn.Module): |
|
|
|
def __init__(self, latent_size, nconv=16, pool=4, drop=0.05): |
|
super(recon_encoder, self).__init__() |
|
|
|
|
|
self.encoder = nn.Sequential( |
|
nn.Conv2d(in_channels=1, out_channels=nconv, kernel_size=3, stride=1, padding=(1,1)), |
|
nn.Dropout(drop), |
|
nn.ReLU(), |
|
nn.Conv2d(nconv, nconv, 3, stride=1, padding=(1,1)), |
|
nn.Dropout(drop), |
|
nn.ReLU(), |
|
nn.MaxPool2d((pool,pool)), |
|
|
|
nn.Conv2d(nconv, nconv*2, 3, stride=1, padding=(1,1)), |
|
nn.Dropout(drop), |
|
nn.ReLU(), |
|
nn.Conv2d(nconv*2, nconv*2, 3, stride=1, padding=(1,1)), |
|
nn.Dropout(drop), |
|
nn.ReLU(), |
|
nn.MaxPool2d((pool,pool)), |
|
|
|
nn.Conv2d(nconv*2, nconv*4, 3, stride=1, padding=(1,1)), |
|
nn.Dropout(drop), |
|
nn.ReLU(), |
|
nn.Conv2d(nconv*4, nconv*4, 3, stride=1, padding=(1,1)), |
|
nn.Dropout(drop), |
|
nn.ReLU(), |
|
nn.MaxPool2d((pool,pool)), |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
) |
|
|
|
|
|
self.bottleneck = nn.Sequential( |
|
|
|
nn.Flatten(), |
|
nn.Linear(1024, latent_size), |
|
|
|
nn.ReLU(), |
|
|
|
|
|
|
|
|
|
) |
|
|
|
|
|
self.decoder1 = nn.Sequential( |
|
|
|
nn.Conv2d(nconv*4, nconv*4, 3, stride=1, padding=(1,1)), |
|
nn.Dropout(drop), |
|
nn.ReLU(), |
|
nn.Conv2d(nconv*4, nconv*4, 3, stride=1, padding=(1,1)), |
|
nn.Dropout(drop), |
|
nn.ReLU(), |
|
nn.Upsample(scale_factor=pool, mode='bilinear'), |
|
|
|
nn.Conv2d(nconv*4, nconv*4, 3, stride=1, padding=(1,1)), |
|
nn.Dropout(drop), |
|
nn.ReLU(), |
|
nn.Conv2d(nconv*4, nconv*4, 3, stride=1, padding=(1,1)), |
|
nn.Dropout(drop), |
|
nn.ReLU(), |
|
nn.Upsample(scale_factor=pool, mode='bilinear'), |
|
|
|
nn.Conv2d(nconv*4, nconv*2, 3, stride=1, padding=(1,1)), |
|
nn.Dropout(drop), |
|
nn.ReLU(), |
|
nn.Conv2d(nconv*2, nconv*2, 3, stride=1, padding=(1,1)), |
|
nn.Dropout(drop), |
|
nn.ReLU(), |
|
nn.Upsample(scale_factor=pool, mode='bilinear'), |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
nn.Conv2d(nconv*2, 1, 3, stride=1, padding=(1,1)), |
|
nn.Sigmoid() |
|
) |
|
|
|
|
|
def forward(self,x): |
|
with torch.cuda.amp.autocast(): |
|
x1 = self.encoder(x) |
|
x1 = self.bottleneck(x1) |
|
|
|
return x1 |
|
|
|
|
|
|
|
def calc_fc_shape(self): |
|
x0 = torch.zeros([256,256]).unsqueeze(0) |
|
x0 = self.encoder(x0) |
|
|
|
self.conv_bock_output_shape = x0.shape |
|
|
|
self.flattened_size = x0.flatten().shape[0] |
|
|
|
return self.flattened_size |
|
|
|
class recon_model(nn.Module): |
|
|
|
def __init__(self, latent_size, nconv=16, pool=4, drop=0.05): |
|
super(recon_model, self).__init__() |
|
|
|
|
|
self.encoder = nn.Sequential( |
|
nn.Conv2d(in_channels=1, out_channels=nconv, kernel_size=3, stride=1, padding=(1,1)), |
|
nn.Dropout(drop), |
|
nn.ReLU(), |
|
nn.Conv2d(nconv, nconv, 3, stride=1, padding=(1,1)), |
|
nn.Dropout(drop), |
|
nn.ReLU(), |
|
nn.MaxPool2d((pool,pool)), |
|
|
|
nn.Conv2d(nconv, nconv*2, 3, stride=1, padding=(1,1)), |
|
nn.Dropout(drop), |
|
nn.ReLU(), |
|
nn.Conv2d(nconv*2, nconv*2, 3, stride=1, padding=(1,1)), |
|
nn.Dropout(drop), |
|
nn.ReLU(), |
|
nn.MaxPool2d((pool,pool)), |
|
|
|
nn.Conv2d(nconv*2, nconv*4, 3, stride=1, padding=(1,1)), |
|
nn.Dropout(drop), |
|
nn.ReLU(), |
|
nn.Conv2d(nconv*4, nconv*4, 3, stride=1, padding=(1,1)), |
|
nn.Dropout(drop), |
|
nn.ReLU(), |
|
nn.MaxPool2d((pool,pool)), |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
) |
|
|
|
|
|
self.bottleneck = nn.Sequential( |
|
|
|
nn.Flatten(), |
|
nn.Linear(1024, latent_size), |
|
|
|
nn.ReLU(), |
|
nn.Linear(latent_size, 1024), |
|
|
|
nn.ReLU(), |
|
nn.Unflatten(1,(64,4,4)) |
|
) |
|
|
|
|
|
self.decoder1 = nn.Sequential( |
|
|
|
nn.Conv2d(nconv*4, nconv*4, 3, stride=1, padding=(1,1)), |
|
nn.Dropout(drop), |
|
nn.ReLU(), |
|
nn.Conv2d(nconv*4, nconv*4, 3, stride=1, padding=(1,1)), |
|
nn.Dropout(drop), |
|
nn.ReLU(), |
|
nn.Upsample(scale_factor=pool, mode='bilinear'), |
|
|
|
nn.Conv2d(nconv*4, nconv*4, 3, stride=1, padding=(1,1)), |
|
nn.Dropout(drop), |
|
nn.ReLU(), |
|
nn.Conv2d(nconv*4, nconv*4, 3, stride=1, padding=(1,1)), |
|
nn.Dropout(drop), |
|
nn.ReLU(), |
|
nn.Upsample(scale_factor=pool, mode='bilinear'), |
|
|
|
nn.Conv2d(nconv*4, nconv*2, 3, stride=1, padding=(1,1)), |
|
nn.Dropout(drop), |
|
nn.ReLU(), |
|
nn.Conv2d(nconv*2, nconv*2, 3, stride=1, padding=(1,1)), |
|
nn.Dropout(drop), |
|
nn.ReLU(), |
|
nn.Upsample(scale_factor=pool, mode='bilinear'), |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
nn.Conv2d(nconv*2, 1, 3, stride=1, padding=(1,1)), |
|
nn.Sigmoid() |
|
) |
|
|
|
|
|
def forward(self,x): |
|
with torch.cuda.amp.autocast(): |
|
x1 = self.encoder(x) |
|
x1 = self.bottleneck(x1) |
|
|
|
return self.decoder1(x1) |
|
|
|
|
|
|
|
def calc_fc_shape(self): |
|
x0 = torch.zeros([256,256]).unsqueeze(0) |
|
x0 = self.encoder(x0) |
|
|
|
self.conv_bock_output_shape = x0.shape |
|
|
|
self.flattened_size = x0.flatten().shape[0] |
|
|
|
return self.flattened_size |
|
|
|
full_model = torch.load('betst_model_100x_0064.pth',map_location=torch.device('cpu')) |
|
encoder_model = recon_encoder(latent_size=64) |
|
encoder_state_dict = encoder_model.state_dict() |
|
|
|
checkpoint = torch.load('betst_model_100x_0064_statedict.pth',map_location=torch.device('cpu')) |
|
pretrained_dict = {k: v for k, v in checkpoint.items() if k in encoder_state_dict} |
|
|
|
encoder_model.load_state_dict(pretrained_dict) |
|
|
|
|
|
|
|
|
|
|
|
def load_data(file): |
|
all_data = np.load(file.name).astype(np.float32) |
|
all_data = all_data.reshape(-1,1,256,256) |
|
dataloader = DataLoader(all_data,batch_size=32,shuffle=False) |
|
return all_data, dataloader, 'upload complete: {}'.format(all_data.shape) |
|
|
|
|
|
|
|
def show_image(selection, all_data): |
|
fig1, ax1 = plt.subplots() |
|
ax1.imshow(all_data[selection][0],plt.cm.inferno,origin='lower') |
|
ax1.axis('off') |
|
fig1.tight_layout() |
|
|
|
fig2, ax2 = plt.subplots() |
|
prediction = full_model(torch.tensor(all_data[selection].reshape(-1,1,256,256))).detach().cpu().numpy() |
|
ax2.imshow(prediction[0,0],plt.cm.inferno,origin='lower') |
|
ax2.axis('off') |
|
fig2.tight_layout() |
|
|
|
return fig1, fig2 |
|
|
|
def encode_data(dataloader): |
|
preds_full = [] |
|
preds_enc = [] |
|
|
|
|
|
for i, images in enumerate(dataloader): |
|
if i > 5: |
|
break |
|
pred_full = full_model(images) |
|
pred_enc = encoder_model(images) |
|
for j in range(images.shape[0]): |
|
preds_full.append(pred_full[j].detach().cpu().numpy()) |
|
preds_enc.append(pred_enc[j].detach().cpu().numpy()) |
|
|
|
|
|
processed_images = np.array(preds_full).squeeze() |
|
encoded_images = np.array(preds_enc) |
|
message = 'finished' |
|
|
|
return message, processed_images, encoded_images |
|
|
|
def print_state(state): |
|
return state.shape |
|
|
|
def latent_vis(encoded_data,decomp_method,clustering_method,cluster_number,all_data): |
|
if decomp_method == 'PCA': |
|
pca = PCA(n_components=2) |
|
decomp = pca.fit_transform(encoded_data) |
|
elif decomp_method == 'tSNE': |
|
tsne = TSNE(n_components=2) |
|
decomp = tsne.fit_transform(encoded_data) |
|
elif decomp_method == 'UMAP': |
|
reducer = UMAP() |
|
decomp = reducer.fit_transform(encoded_data) |
|
|
|
if clustering_method == 'KMeans': |
|
kmeans = KMeans(n_clusters=int(cluster_number)) |
|
cluster_labels = kmeans.fit_predict(encoded_data) |
|
|
|
|
|
df = pd.DataFrame(decomp,columns=['x','y']) |
|
df['cluster'] = cluster_labels |
|
df['value'] = np.ones_like(cluster_labels) * np.arange(len(decomp)) |
|
|
|
fig = px.scatter(df,x='x',y='y',color='cluster',color_continuous_scale='viridis',hover_name='value',hover_data={'x': False, |
|
'y': False, |
|
'cluster': False, |
|
'value': False}) |
|
|
|
fig.update_layout(clickmode='event+select') |
|
fig.update_traces(marker=dict(size=12), |
|
selector=dict(mode='markers')) |
|
|
|
fig1 = plt.figure(figsize=(20,5)) |
|
n_rows = 1 |
|
n_cols = int(cluster_number) |
|
colors = plt.cm.viridis(np.linspace(0,1,len(np.unique(cluster_labels)))) |
|
|
|
for i in np.unique(cluster_labels): |
|
ind = np.where(cluster_labels[:] == i)[0] |
|
|
|
|
|
r = np.random.choice(ind) |
|
ax1 = fig1.add_subplot(n_rows,n_cols,i+1) |
|
ax1.imshow(all_data[r][0],plt.cm.inferno,origin='lower') |
|
ax1.set_title('Class {}: {}'.format(i,len(ind)),color=colors[i],fontsize=20,weight='bold') |
|
|
|
|
|
|
|
|
|
fig1.tight_layout() |
|
return decomp, cluster_labels, fig, fig1 |
|
|
|
def interactive_vis(decomp,clusters,images): |
|
df = pd.DataFrame(decomp,columns=['x','y']) |
|
df['cluster'] = clusters |
|
df['value'] = np.ones_like(clusters) * np.arange(len(decomp)) |
|
df['im'] = images |
|
|
|
fig = px.scatter(df,x='x',y='y',color='cluster',custom_data='im',color_continuous_scale='viridis',hover_name='value',hover_data={'x': False, |
|
'y': False, |
|
'cluster': False, |
|
'value': False}) |
|
|
|
fig.update_layout(clickmode='event+select') |
|
fig.update_traces(marker=dict(size=20), |
|
selector=dict(mode='markers')) |
|
|
|
return fig |
|
|
|
def neighbor_vis(decomp,neighbor_index,n_neighbors,all_data): |
|
neighbor_index = int(neighbor_index) |
|
|
|
d = np.sqrt((decomp[:,0] - decomp[neighbor_index,0]) ** 2 + (decomp[:,1] - decomp[neighbor_index,1]) ** 2) |
|
ar = np.argsort(d) |
|
|
|
n_rows = int(np.ceil(n_neighbors/5)) |
|
n_cols = 5 |
|
fig = plt.figure(figsize=(20,5*n_rows)) |
|
|
|
n = 1 |
|
ax = fig.add_subplot(n_rows,n_cols,n) |
|
ax.imshow(all_data[neighbor_index][0],plt.cm.inferno,origin='lower') |
|
ax.set_title('{}'.format(neighbor_index),fontsize=20,weight='bold') |
|
ax.axis('off') |
|
n += 1 |
|
|
|
neighbors = ar[1:1+n_neighbors-1] |
|
|
|
for i in neighbors: |
|
ax = fig.add_subplot(n_rows,n_cols,n) |
|
ax.imshow(all_data[i][0],plt.cm.inferno,origin='lower') |
|
ax.set_title('{}'.format(i),fontsize=20) |
|
ax.axis('off') |
|
n += 1 |
|
|
|
return fig |
|
|
|
|
|
intro_text1 = '# AI-NERD: Artificial Intelligence for Non-Equilibrium Relaxation Dynamics' |
|
intro_text2 = 'AI-NERD is a platform for applying unsupervised image classification to X-ray Photon Corrleation Spectroscopy (XPCS) data. Here, we demonstrate how raw experimental data can be automatically processed and clustered, and how latent space analysis can be used to understand the physics of relaxing systems without any background information or assumptions.<br><br>Please see out [preprint](https://arxiv.org/abs/2212.03984) for more information.<br><br>' |
|
l = 900 |
|
with gr.Blocks() as demo: |
|
gr.Markdown(intro_text1) |
|
gr.Markdown(intro_text2) |
|
|
|
gr.Markdown('### Evaluation of Training Results') |
|
gr.Markdown('Use the dropdown menu below to select a sample image. The frame on the left will show the raw C2 data, and the frame on the right will show the neural network output. After sampling individual images, click _Process All Images_ to run the entire dataset through the Autoencoder') |
|
|
|
with gr.Row(): |
|
file_path = gr.File() |
|
with gr.Column(): |
|
upload_status = gr.Textbox(label='file upload status') |
|
file_upload = gr.Button(value='load data') |
|
|
|
all_data = gr.State() |
|
dataloader = gr.State() |
|
file_upload.click(load_data,file_path,[all_data,dataloader,upload_status]) |
|
|
|
selection = gr.Dropdown(list(np.arange(2000)),value=200,label='select sample image') |
|
with gr.Row(): |
|
output_image_1 = gr.Plot(label='input C2 data') |
|
output_image_2 = gr.Plot(label='Autoencoder Reproduction') |
|
|
|
selection.change(show_image,[selection, all_data],[output_image_1,output_image_2]) |
|
|
|
with gr.Row(): |
|
process_all = gr.Button(value='Process All Images') |
|
status = gr.Textbox(label='batch processing status') |
|
|
|
|
|
proc_im = gr.State() |
|
enc_im = gr.State() |
|
process_all.click(encode_data,inputs=[dataloader],outputs=[status,proc_im,enc_im],show_progress=True,status_tracker=None) |
|
|
|
|
|
|
|
|
|
gr.Markdown('<br><br>') |
|
gr.Markdown('### Latent Space Visualization') |
|
gr.Markdown('Select the decomposition and clustering method for latent space visualization') |
|
with gr.Row(): |
|
with gr.Column(): |
|
decomp_method = gr.Dropdown(choices=['PCA','tSNE','UMAP'],label='select decomposition method',value='UMAP') |
|
with gr.Row(): |
|
clustering_method = gr.Dropdown(choices=['KMeans','Agglomerative','DBSCAN'],label='select clusterting algorithm',value='KMeans') |
|
cluster_number = gr.Number(label='input number of clusters',value=5) |
|
|
|
process_vis = gr.Button(value='Visualize Latent Space') |
|
latent_scatter = gr.Plot() |
|
latent_sample = gr.Plot() |
|
|
|
save_decomp_coords = gr.State() |
|
save_cluster_labels = gr.State() |
|
process_vis.click(latent_vis,[enc_im,decomp_method,clustering_method,cluster_number,all_data],[save_decomp_coords,save_cluster_labels,latent_scatter,latent_sample]) |
|
|
|
gr.Markdown('<br><br><br>') |
|
gr.Markdown('### Visualize Nearest Neighbors') |
|
gr.Markdown('Hover over data points in the scatter plot above, to identify the index of points of interest. Enter the desired index in the box below, and click _Visualize Neighbors_.') |
|
|
|
with gr.Row(): |
|
with gr.Column(): |
|
neighbor_index = gr.Number(label='input point index',value=110) |
|
n_neighbors = gr.Slider(label='select number of neighbors to view',minimum=5,maximum=10,value=5,step=1) |
|
|
|
neighbor_button = gr.Button(value='Visualize Neighbors') |
|
|
|
neighbor_plot = gr.Plot() |
|
neighbor_button.click(neighbor_vis,[save_decomp_coords,neighbor_index,n_neighbors,all_data],neighbor_plot) |
|
|
|
|
|
|
|
|
|
demo.launch() |
|
|
|
|
