app.py

import gradio as gr

import os
os.environ['NUMBA_ENABLE_CUDASIM'] = '1'
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' # Suppress TensorFlow warnings unless fatal

# !! For now, because it is a private repo !!
from huggingface_hub import login, hf_hub_download
# login(token=os.environ.get("HF_AUTH_TOKEN"))
# from huggingface_hub import hf_hub_download

from tensorflow.keras.saving import load_model
import numpy as np
from numba import cuda
from scipy.io import readsav
import matplotlib.pyplot as plt
import time

from filters import GaussianFilter
from vec_metrics import get_pred_vis, mae_of_vectors



# ============================================= LOAD MODEL ===========================================
# Load your trained model
model_path = hf_hub_download(repo_id="mervess/FCD-Solar", filename="fcd.keras")
model = load_model(model_path,
                   custom_objects={'GaussianFilter': GaussianFilter},
                   compile=False)
print("#) Model loaded successfully.")

# Load the F matrix
f_matrix = np.load("data/fourier_matrix.npy")
print("#) STIX fourier matrix loaded successfully.")
# ============================================= PREDICT ==============================================
# Define a function for prediction
def predict_image(vis):
    norm_vis = np.empty(vis.shape)
    alphas = np.empty(vis.shape[0], dtype=np.float32)

    threadsperblock = 32
    blockspergrid = (alphas.size + (threadsperblock - 1)) // threadsperblock

    calc_alphas_cpu[blockspergrid, threadsperblock](vis, alphas, norm_vis)

    predicted_img = model.predict(norm_vis, batch_size=128, verbose=False)
    predicted_img = predicted_img * alphas[:, np.newaxis, np.newaxis, np.newaxis]
    output_image = np.squeeze(predicted_img)
    return output_image

@cuda.jit
def calc_alphas_cpu(vis_arrays, alphas, norm_vis):
    pos = cuda.grid(1)
    if pos < alphas.size:
        nvis = np.sqrt(
            np.square(vis_arrays[pos, :24]) + np.square(vis_arrays[pos, 24:]))
        alphas[pos] = np.max(nvis) * 0.5
        for i in range(vis_arrays.shape[1]):
            norm_vis[pos][i] = vis_arrays[pos][i] / alphas[pos]
# ============================================ READ FILES ============================================
def read_file( file_name, file_type='.npy' ):
    if file_type == '.npy':
        vis_array = np.load( file_name )
    elif file_type == '.sav':
        data = readsav( file_name, python_dict=True, verbose=False )
        vis_structure = data['vis']
        vis_components = vis_structure['OBSVIS']
        vis_array = np.hstack( ( np.real(vis_components), np.imag(vis_components) ) )
    
    return vis_array
# ========================================== EXAMPLE INPUTS ==========================================
# Provide sample data
# ===== 1 =====
vis_array = read_file('example_data/20211216_vis.sav', file_type='.sav')
# ===== 2 =====
vis_array_2 = read_file('example_data/20221231_vis.sav', file_type='.sav')
# ===== 3 =====
vis_array_3 = read_file('example_data/20220828_vis.sav', file_type='.sav')

sample_vectors = {
    "20211216T233548-20211216T234036_10-15keV": vis_array,
    "20221231T123044-20221231T123252_4-10keV": vis_array_2,
    "20220828T155202-20220828T160950_25-50keV": vis_array_3,
}
print("#) Sample data loaded successfully.")
# ============================================ VISUALIZE =============================================
def plot_map( map_data, img_name='', color_map='turbo', is_minorticks_on=True, show_colorbar=True, save_plot=False, plot_title="plot.pdf" ):
    fig, ax = plt.subplots(figsize=(3, 3))  

    cax = ax.imshow(map_data, cmap=color_map)
    ax.set_title(img_name)  
    ax.axis("off")  

    if show_colorbar:
        cbar = fig.colorbar(cax, ax=ax, orientation='horizontal', 
                            fraction=0.046, pad=0.07, extend='both')
        cbar.ax.tick_params(labelsize='large')  # Increase colorbar font size

        # Set the colorbar ticks to beginning, middle, and end, optimizing the formatting process
        min_val = map_data.min()
        max_val = map_data.max()
        mid_val = (min_val + max_val) / 2
        cbar.set_ticks([min_val, mid_val, max_val])
        cbar.set_ticklabels([f'{min_val:.4f}', f'{mid_val:.4f}', f'{max_val:.4f}'])

        if is_minorticks_on:
            cbar.minorticks_on()
            
    plt.tight_layout()

    if save_plot:
        f_format = os.path.splitext(plot_title)[-1][1:]
        if show_colorbar:
            plt.savefig(plot_title, format=f_format, bbox_inches='tight')
        else:
            plt.savefig(plot_title, format=f_format, bbox_inches='tight', pad_inches=0)

    plt.show()

def to_html( msg, color='red' ):
    return f"<html><p style='color:{color};'>{msg}</p></html>"
# =============================================== DEMO ===============================================
# Gradio interface
def demo( sample_choice, input_vector, uploaded_file ):
    """
    Function to demonstrate the FCD model.
    
    Parameters:
    -----------
    sample_choice: str
        Sample choice from the dropdown list.
    input_vector: str
        Input vector from the textbox.
    uploaded_file: file
        Uploaded file.
    
    Returns:
    --------
    numpy array (image)
        Reconstructed image.
    str
        (-) Message to display. / (+) Prediction time.
    str
        Chi2 value.
    """
    vis = None
    # ==== SAMPLE CHOICE ====
    if sample_choice != "" and \
        sample_choice is not None and \
        sample_choice in sample_vectors:
        vis = sample_vectors[sample_choice]

    # ==== INPUT VECTOR ====
    elif input_vector != "" and input_vector is not None and \
        isinstance(input_vector, str):
        try:
            vis = np.array( [ float(i) for i in input_vector.split(",") ], dtype=np.float32 )
            if vis.shape != (48,):
                return None, None, None, to_html("Please provide an input vector of 48 real numbers.")
        except:
            return None, None, None, to_html("Please provide an input consisting of real numbers.")
    
    # ==== UPLOADED FILE ====
    elif uploaded_file is not None:
        try:
            # Get the file name
            file_name = uploaded_file.name
            # Split the file name to get the extension
            _, file_extension = os.path.splitext(file_name)
            vis = read_file( file_name, file_type=file_extension )
        except Exception as e:
            return None, None, None, to_html(f"Error processing the uploaded file: {e}.")

    # ==== PREDICT ====
    if vis is None:
        return None, None, None, to_html("Please provide a sample, input vector, or upload a file.")
    else:
        start_time = time.time()
        image = predict_image(vis.reshape(1, -1))
        prediction_time = time.time() - start_time
        
        # Calculate MAE between predicted visibility and input visibility
        pred_vis = get_pred_vis( image, f_matrix )
        mae = mae_of_vectors( vis, pred_vis )
        plot_map(image, save_plot=True, plot_title="fcd_image.png")
        return "fcd_image.png", \
                    f"{prediction_time:.2f} seconds", \
                    f"{mae:.3f}", \
                    to_html("Image has been reconstructed.", color='green')

# =============================================== MAIN ===============================================
# Build Gradio Interface
interface = gr.Interface(
    fn=demo,
    inputs=[
        gr.Dropdown(
            choices=[""] + list(sample_vectors.keys()),
            label="Select a sample vector (i.e., visibility, Fourier components):",
            value="",  # Placeholder
        ),
        gr.Textbox(
            label="Enter your vector (comma-separated 48 real numbers):",
            placeholder="e.g., 0.1, 0.2, 0.3",
        ),
        gr.File(
            label="Upload your file (either a .sav or .npy file):",
            type="filepath",  # Accept file uploads
        ),
    ],
    outputs=[
        gr.Image(type="filepath", label="Reconstructed Image"),
        gr.Textbox(label="Reconstruction Time"),  # Line for prediction time
        gr.Textbox(label="MAE Metric"),           # Line for MAE metric
        gr.HTML(label="Message"),                 # Line for messages
    ],
    title="FCD Demo",
    description="Select a sample, enter vector (visibility) values, or upload a file to see the reconstructed image.",
)

# Launch the app
interface.launch()