| import os |
| import tempfile |
| from pathlib import Path |
|
|
| import torch |
| import gradio as gr |
| from PIL import Image |
|
|
| from config import CONFIG |
| from model import load_from_checkpoint, MobileNetUNet |
| from inference_utils import full_inference, make_side_by_side, tiled_inference |
|
|
|
|
| MODEL = None |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
|
|
| def load_model(): |
| global MODEL |
| ckpt_path = CONFIG["checkpoint_path"] |
| if os.path.exists(ckpt_path): |
| print(f"Loading checkpoint from {ckpt_path}...") |
| MODEL = load_from_checkpoint(ckpt_path, device=DEVICE, width_mult=CONFIG["width_mult"]) |
| MODEL.to(DEVICE) |
| MODEL.eval() |
| print("Model loaded successfully.") |
| else: |
| print(f"Checkpoint not found at {ckpt_path}. Using ONNX is not yet supported in this app.") |
| MODEL = None |
|
|
|
|
| def process_image(input_img, use_tiled, tile_overlap): |
| if MODEL is None: |
| load_model() |
| if MODEL is None: |
| raise gr.Error("No model loaded. Please ensure checkpoints/last.ckpt exists.") |
|
|
| img = input_img.convert("RGB") |
|
|
| with torch.no_grad(): |
| if use_tiled: |
| inp_img, outputs = tiled_inference( |
| MODEL, img, CONFIG["image_size"], DEVICE, overlap=tile_overlap, |
| ) |
| else: |
| inp_img, outputs = full_inference(MODEL, img, CONFIG["image_size"], DEVICE) |
|
|
| result = make_side_by_side(inp_img, outputs) |
| return result, outputs |
|
|
|
|
| def run_and_save(input_img, use_tiled, tile_overlap): |
| result, outputs = process_image(input_img, use_tiled, tile_overlap) |
|
|
| tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".png") |
| result.save(tmp.name) |
| return tmp.name |
|
|
|
|
| def run_and_return_maps(input_img, use_tiled, tile_overlap): |
| result, outputs = process_image(input_img, use_tiled, tile_overlap) |
| return result, outputs.get("basecolor"), outputs.get("normal"), outputs.get("rmd"), outputs.get("rgb") |
|
|
|
|
| with gr.Blocks(title="ShadeNet", theme=gr.themes.Soft()) as demo: |
| gr.Markdown( |
| """ |
| # ShadeNet |
| Lightweight inverse rendering model that decomposes an RGB image into **basecolor**, **normal**, **roughness/metallic/depth** maps, |
| then recombines them into a reconstructed RGB. |
| """ |
| ) |
|
|
| with gr.Row(): |
| with gr.Column(): |
| input_img = gr.Image(type="pil", label="Input Image") |
| with gr.Row(): |
| use_tiled = gr.Checkbox(label="Tiled Inference", value=False) |
| tile_overlap = gr.Slider(minimum=0, maximum=128, value=16, step=1, label="Tile Overlap (px)") |
| run_btn = gr.Button("Decompose", variant="primary") |
|
|
| with gr.Column(): |
| output_img = gr.Image(type="pil", label="Side-by-Side Result") |
|
|
| with gr.Row(): |
| gr.Markdown("### Individual Output Maps") |
| with gr.Row(): |
| basecolor_out = gr.Image(type="pil", label="Basecolor") |
| normal_out = gr.Image(type="pil", label="Normal") |
| rmd_out = gr.Image(type="pil", label="RMD (R=roughness, G=metallic, B=depth)") |
| rgb_out = gr.Image(type="pil", label="Reconstructed RGB") |
|
|
| run_btn.click( |
| fn=run_and_return_maps, |
| inputs=[input_img, use_tiled, tile_overlap], |
| outputs=[output_img, basecolor_out, normal_out, rmd_out, rgb_out], |
| ) |
|
|
| gr.Examples( |
| examples=[], |
| inputs=[input_img], |
| ) |
|
|
| gr.Markdown( |
| """ |
| --- |
| ### Model Details |
| - **Architecture**: ShadeNet 28M — MobileNetV2 backbone + Parallel Encoder + UNet decoder with attention |
| - **Input**: RGB image (resized to 512×512 center crop) |
| - **Outputs**: Basecolor (3ch), Normal (3ch), RMD (3ch = roughness, metallic, depth), Reconstructed RGB (3ch) |
| - **ONNX models**: Available in the `onnx/` folder for CPU/edge deployment |
| """ |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| load_model() |
| demo.launch(server_name="0.0.0.0") |
|
|
