|
from matplotlib import cm |
|
import matplotlib.pyplot as plt |
|
|
|
|
|
import numpy as np |
|
|
|
|
|
import onnxruntime as ort |
|
|
|
from optimum.onnxruntime import ORTModel |
|
|
|
import pandas as pd |
|
|
|
from scipy import special |
|
|
|
|
|
|
|
|
|
import gradio as gr |
|
|
|
|
|
|
|
model_path = 'chlab/planet_detection_models/' |
|
|
|
|
|
labels = 20 |
|
ticks = 14 |
|
legends = 14 |
|
text = 14 |
|
titles = 22 |
|
lw = 3 |
|
ps = 200 |
|
cmap = 'magma' |
|
|
|
def normalize_array(x: list): |
|
|
|
'''Makes array between 0 and 1''' |
|
|
|
x = np.array(x) |
|
|
|
return (x - np.min(x)) / np.max(x - np.min(x)) |
|
|
|
def load_model(model: str, activation: bool=True): |
|
|
|
if activation: |
|
model += '_w_activation' |
|
|
|
options = ort.SessionOptions() |
|
options.intra_op_num_threads = 1 |
|
options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL |
|
provider = "CPUExecutionProvider" |
|
ort_session = ort.InferenceSession(model_path + '%s.onnx' % (model), options, providers=[provider]) |
|
|
|
|
|
return ort_session |
|
|
|
def get_activations(intermediate_model, image: list, |
|
layer=None, vmax=2.5, sub_mean=True): |
|
|
|
'''Gets activations for a given input image''' |
|
|
|
|
|
input_name = intermediate_model.get_inputs()[0].name |
|
outputs = intermediate_model.run(None, {input_name: image}) |
|
|
|
|
|
output_1 = outputs[1] |
|
output_2 = outputs[2] |
|
|
|
output = outputs[0] |
|
output = special.softmax(output) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
in_image = np.sum(image[0, :, :, :], axis=0) |
|
in_image = normalize_array(in_image) |
|
|
|
|
|
|
|
if layer is None: |
|
activation_1 = np.sum(output_1[0, :, :, :], axis=0) |
|
activation_2 = np.sum(output_2[0, :, :, :], axis=0) |
|
else: |
|
activation_1 = output_1[0, layer, :, :] |
|
activation_2 = output_2[0, layer, :, :] |
|
|
|
if sub_mean: |
|
activation_1 -= np.mean(activation_1) |
|
activation_1 = np.abs(activation_1) |
|
|
|
activation_2 -= np.mean(activation_2) |
|
activation_2 = np.abs(activation_2) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
return in_image, activation_1, activation_2 |
|
|
|
|
|
def predict_and_analyze(model_name, num_channels, dim, image): |
|
|
|
'''Loads a model with activations, passes through image and shows activations |
|
|
|
The image must be a pandas dataframe that can be made from a (C, W, H) numpy array |
|
using |
|
|
|
m,n,r = X.shape |
|
X = np.column_stack((np.repeat(np.arange(C), W), |
|
X.reshape(C * W, -1))) |
|
df = pd.DataFrame(X) |
|
|
|
then get the image back with |
|
|
|
X = df.values |
|
X = X[:, 1:] |
|
X = X.reshape((C, W, W)) |
|
|
|
image = 2d numpy array in shape (C, W*W) |
|
i.e. take a C,W,W array and reshape into (C, W*W) |
|
|
|
''' |
|
|
|
num_channels = int(num_channels) |
|
W = int(dim) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
print("Loading data") |
|
image = np.load(image.name, allow_pickle=True) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
if len(image.shape) != 4: |
|
image = image[np.newaxis, :, :, :] |
|
|
|
assert image.shape == (1, num_channels, W, W), "Data is the wrong shape" |
|
|
|
input_image = np.sum(image[0, :, :, :], axis=0) |
|
|
|
model_name += '_%i' % (num_channels) |
|
|
|
|
|
print("Loading model") |
|
model = load_model(model_name, activation=True) |
|
|
|
print("Looking at activations") |
|
output, activation_1, activation_2 = get_activations(model, image, sub_mean=True) |
|
|
|
output = 'Planet prediction with %f percent confidence' % (100*output) |
|
|
|
return output, input_image, activation_1, activation_2 |
|
|
|
|
|
demo = gr.Interface( |
|
fn=predict_and_analyze, |
|
inputs=[gr.Dropdown(["regnet", "efficientnet"], |
|
value="efficientnet", |
|
label="Model Selection", |
|
show_label=True), |
|
gr.Dropdown(["45", "61", "75"], |
|
value="61", |
|
label="Number of Velocity Channels", |
|
show_label=True), |
|
gr.Dropdown(["600"], |
|
value="600", |
|
label="Image Dimensions", |
|
show_label=True), |
|
gr.File(label="Input Data", show_label=True)], |
|
outputs=[gr.Textbox(lines=1, label="Prediction", show_label=True), |
|
gr.Image(label="Input Image", show_label=True), |
|
gr.Image(label="Activation 1", show_label=True), |
|
gr.Image(label="Actication 2", show_label=True)], |
|
title="Kinematic Planet Detector" |
|
) |
|
demo.launch() |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|