|
|
""" |
|
|
Gradio interface for HATSAT application. |
|
|
""" |
|
|
|
|
|
import gradio as gr |
|
|
from PIL import Image |
|
|
|
|
|
from config import REQUIRED_IMAGE_SIZE, UPSCALE_FACTOR |
|
|
from utils.image_utils import validate_image_size, upscale_image |
|
|
from interface.css_styles import generate_css, get_sample_images |
|
|
|
|
|
|
|
|
def upscale_and_display(image, model, device): |
|
|
"""Process image upload and return upscaled result.""" |
|
|
if image is None: |
|
|
return None, "Please upload an image or select a sample image." |
|
|
|
|
|
|
|
|
is_valid, message = validate_image_size(image) |
|
|
if not is_valid: |
|
|
return None, f"❌ Error: {message}" |
|
|
|
|
|
try: |
|
|
|
|
|
upscaled = upscale_image(image, model, device) |
|
|
return upscaled, "✅ Image successfully enhanced!" |
|
|
except Exception as e: |
|
|
return None, f"❌ Error processing image: {str(e)}" |
|
|
|
|
|
|
|
|
def select_sample_image(image_path): |
|
|
"""Load and return a sample image.""" |
|
|
if image_path: |
|
|
return Image.open(image_path) |
|
|
return None |
|
|
|
|
|
|
|
|
def create_interface(model, device): |
|
|
"""Create and configure the Gradio interface.""" |
|
|
css = generate_css() |
|
|
|
|
|
with gr.Blocks(css=css, title="HATSAT - Super-Resolution for Satellite Images") as iface: |
|
|
gr.Markdown("# HATSAT - Super-Resolution for Satellite Images") |
|
|
gr.Markdown(f"Upload a satellite image or select a sample to enhance its resolution by {UPSCALE_FACTOR}x.") |
|
|
gr.Markdown(f"⚠️ **Important**: Images must be exactly **{REQUIRED_IMAGE_SIZE[0]}x{REQUIRED_IMAGE_SIZE[1]} pixels** for the model to work properly.") |
|
|
|
|
|
|
|
|
with gr.Accordion("Acknowledgments", open=False): |
|
|
gr.Markdown(""" |
|
|
### Base Model: HAT (Hybrid Attention Transformer) |
|
|
This model is a fine tuned version of **HAT**: |
|
|
- **GitHub Repository**: [https://github.com/XPixelGroup/HAT](https://github.com/XPixelGroup/HAT) |
|
|
- **Paper**: [Activating More Pixels in Image Super-Resolution Transformer](https://arxiv.org/abs/2205.04437) |
|
|
- **Authors**: Xiangyu Chen, Xintao Wang, Jiantao Zhou, Yu Qiao, Chao Dong |
|
|
|
|
|
### Training Dataset: SEN2NAIPv2 |
|
|
The model was fine-tuned using the **SEN2NAIPv2** dataset: |
|
|
- **HuggingFace Dataset**: [https://huggingface.co/datasets/tacofoundation/SEN2NAIPv2](https://huggingface.co/datasets/tacofoundation/SEN2NAIPv2) |
|
|
- **Description**: High-resolution satellite imagery dataset for super-resolution tasks |
|
|
""") |
|
|
|
|
|
|
|
|
sample_images = get_sample_images() |
|
|
sample_buttons = [] |
|
|
if sample_images: |
|
|
gr.Markdown("**Sample Images (click to select):**") |
|
|
with gr.Row(): |
|
|
for i, img_path in enumerate(sample_images): |
|
|
btn = gr.Button( |
|
|
"", |
|
|
elem_id=f"sample_btn_{i}", |
|
|
elem_classes="sample-image-btn" |
|
|
) |
|
|
sample_buttons.append((btn, img_path)) |
|
|
|
|
|
with gr.Row(): |
|
|
input_image = gr.Image( |
|
|
type="pil", |
|
|
label=f"Input Image (must be {REQUIRED_IMAGE_SIZE[0]}x{REQUIRED_IMAGE_SIZE[1]} pixels)", |
|
|
elem_classes="image-container", |
|
|
sources=["upload"], |
|
|
height=500, |
|
|
width=500 |
|
|
) |
|
|
|
|
|
output_image = gr.Image( |
|
|
type="pil", |
|
|
label=f"Enhanced Output ({UPSCALE_FACTOR}x)", |
|
|
elem_classes="image-container", |
|
|
interactive=False, |
|
|
height=500, |
|
|
width=500, |
|
|
show_download_button=True |
|
|
) |
|
|
|
|
|
submit_btn = gr.Button("Enhance Image", variant="primary") |
|
|
|
|
|
|
|
|
status_message = gr.Textbox( |
|
|
label="Status", |
|
|
interactive=False, |
|
|
show_label=True |
|
|
) |
|
|
|
|
|
|
|
|
if sample_images: |
|
|
for btn, img_path in sample_buttons: |
|
|
btn.click(fn=lambda path=img_path: select_sample_image(path), outputs=input_image) |
|
|
|
|
|
submit_btn.click( |
|
|
fn=lambda img: upscale_and_display(img, model, device), |
|
|
inputs=input_image, |
|
|
outputs=[output_image, status_message] |
|
|
) |
|
|
|
|
|
return iface |
