Upload app.py

app.py

@@ -0,0 +1,113 @@
+import torch
+import numpy as np
+import gradio as gr
+import cv2
+import h5py
+from test_develop_code.architecture import model_generator
+import PIL.Image
+
+
+device = torch.device("cpu")
+model = model_generator("mst_plus_plus", "mst_plus_plus.pth").to(device)
+model.eval()
+wavelengths = np.linspace(400, 700, 31)
+
+
+def wavelength_to_rgb(wl: float) -> tuple[float]:
+    if 380 <= wl <= 440:
+        R = -(wl - 440) / (440 - 380)
+        G = 0.0
+        B = 1.0
+    elif 440 < wl <= 490:
+        R = 0.0
+        G = (wl - 440) / (490 - 440)
+        B = 1.0
+    elif 490 < wl <= 510:
+        R = 0.0
+        G = 1.0
+        B = -(wl - 510) / (510 - 490)
+    elif 510 < wl <= 580:
+        R = (wl - 510) / (580 - 510)
+        G = 1.0
+        B = 0.0
+    elif 580 < wl <= 645:
+        R = 1.0
+        G = -(wl - 645) / (645 - 580)
+        B = 0.0
+    elif 645 < wl <= 700:
+        R = 1.0
+        G = 0.0
+        B = 0.0
+    else:
+        R = G = B = 0.0
+    return (max(R, 0.0), max(G, 0.0), max(B, 0.0))
+
+
+def predict(img: np.ndarray) -> np.ndarray:
+    # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    img = img.astype(np.float32)
+    img = (img - img.min()) / (img.max() - img.min() + 1e-8)
+    img = np.transpose(img, (2, 0, 1))
+    img_tensor = torch.from_numpy(img).unsqueeze(0).to(device)
+    with torch.no_grad():
+        pred = model(img_tensor)
+    pred = pred.squeeze(0).cpu().numpy()
+    pred = np.clip(pred, 0, 1)
+    return pred
+
+
+def visualize_channel(cube: np.ndarray, index: int) -> PIL.Image.Image:
+    if cube is None:
+        return None
+    band = cube[index]
+    band = (band - band.min()) / (band.max() - band.min() + 1e-8)
+    color = wavelength_to_rgb(wavelengths[index])
+    rgb = np.stack([band * c for c in color], axis=-1)
+    rgb = (rgb * 255).astype(np.uint8)
+    return PIL.Image.fromarray(rgb)
+
+
+def load_mat(mat_file: gr.File) -> np.ndarray:
+    with h5py.File(mat_file.name, "r") as f:
+        cube = np.array(f["cube"])
+    cube = np.transpose(cube, (0, 2, 1))
+    cube = np.clip(cube, 0, 1)
+    return cube
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("## Spectral Reconstruction")
+    with gr.Row():
+        with gr.Column():
+            rgb_input = gr.Image(type="numpy", label="Upload RGB Image")
+            pred_state = gr.State()
+        with gr.Column():
+            pred_output = gr.Image(label="Prediction Visualization")
+            pred_slider = gr.Slider(minimum=0, maximum=30, step=1, label="Channel (Prediction)", value=0)
+    with gr.Row():
+        with gr.Column():
+            mat_input = gr.File(label="Upload .mat file (Ground Truth)")
+            gt_state = gr.State()
+        with gr.Column():
+            gt_output = gr.Image(label="Ground Truth Visualization")
+            gt_slider = gr.Slider(minimum=0, maximum=30, step=1, label="Channel (Ground Truth)", value=0)
+    rgb_input.change(fn=predict, inputs=rgb_input, outputs=pred_state)
+    pred_slider.change(fn=visualize_channel, inputs=[pred_state, pred_slider], outputs=pred_output)
+    mat_input.change(fn=load_mat, inputs=mat_input, outputs=gt_state)
+    gt_slider.change(fn=visualize_channel, inputs=[gt_state, gt_slider], outputs=gt_output)
+    gr.Examples(
+        examples=[
+            ["assets/ARAD_1K_0001.jpg", 0, "assets/ARAD_1K_0001.mat", 0],
+            ["assets/ARAD_1K_0002.jpg", 0, "assets/ARAD_1K_0002.mat", 0],
+            ["assets/ARAD_1K_0003.jpg", 0, "assets/ARAD_1K_0003.mat", 0],
+            ["assets/ARAD_1K_0004.jpg", 0, "assets/ARAD_1K_0004.mat", 0],
+            ["assets/ARAD_1K_0005.jpg", 0, "assets/ARAD_1K_0005.mat", 0],
+        ],
+        inputs=[rgb_input, pred_slider, mat_input, gt_slider],
+        outputs=[pred_output, gt_output],
+        label="Try Examples"
+    )
+
+
+if __name__ == "__main__":
+    demo.launch()