app.py

# pyrefly: ignore [missing-import]
import gradio as gr
import tensorflow as tf
# pyrefly: ignore [missing-import]
from PIL import Image
# pyrefly: ignore [missing-import]
import numpy as np
import os

# ---------------------------------------------------------------------------
# Model loading
# ---------------------------------------------------------------------------
MODEL_PATH = os.path.join(os.path.dirname(__file__), "GAN_Sat_image_2_map.h5")
model = tf.keras.models.load_model(MODEL_PATH, compile=False)

# ---------------------------------------------------------------------------
# Prediction function
# ---------------------------------------------------------------------------
def predict(image):
    """Convert a satellite image to a map-style image using the trained
    Pix2Pix generator.

    Args:
        image: PIL.Image uploaded by the user.

    Returns:
        PIL.Image of the generated map.
    """
    if image is None:
        raise gr.Error("Please upload a satellite image first.")

    # Pre-process: resize to the model's expected input & normalise to [0, 1]
    image = image.resize((256, 256))
    img_array = np.array(image).astype("float32") / 255.0
    img_array = np.expand_dims(img_array, axis=0)  # add batch dimension

    # Run inference
    prediction = model.predict(img_array, verbose=0)[0]

    # Post-process: clip & convert back to uint8 image
    prediction = np.clip(prediction, 0.0, 1.0)
    prediction = (prediction * 255).astype(np.uint8)

    return Image.fromarray(prediction)


# ---------------------------------------------------------------------------
# Example images (optional — placed in ./examples/)
# ---------------------------------------------------------------------------
example_dir = os.path.join(os.path.dirname(__file__), "examples")
examples = None
if os.path.isdir(example_dir):
    example_files = sorted(
        [
            os.path.join(example_dir, f)
            for f in os.listdir(example_dir)
            if f.lower().endswith((".png", ".jpg", ".jpeg", ".webp"))
        ]
    )
    if example_files:
        examples = [[f] for f in example_files]


# ---------------------------------------------------------------------------
# Custom CSS for a polished look
# ---------------------------------------------------------------------------
custom_css = """
.gradio-container {
    max-width: 960px !important;
    margin: 0 auto;
}
.gr-button-primary {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    border: none !important;
    font-weight: 600 !important;
}
.gr-button-primary:hover {
    background: linear-gradient(135deg, #764ba2 0%, #667eea 100%) !important;
    transform: translateY(-1px);
    box-shadow: 0 4px 15px rgba(102, 126, 234, 0.4) !important;
}
footer { display: none !important; }
"""

# ---------------------------------------------------------------------------
# Gradio Interface
# ---------------------------------------------------------------------------
with gr.Blocks(css=custom_css, title="Pix2Pix — Satellite to Map", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🛰️ Pix2Pix GAN — Satellite Image to Map
        **Convert satellite / aerial imagery into clean map-style views**
        powered by a Conditional GAN with U-Net generator & PatchGAN discriminator.
        """
    )

    with gr.Row(equal_height=True):
        with gr.Column():
            input_image = gr.Image(
                type="pil",
                label="📡 Upload Satellite Image",
                height=350,
            )
            submit_btn = gr.Button("🗺️ Generate Map", variant="primary", size="lg")

        with gr.Column():
            output_image = gr.Image(
                type="pil",
                label="🗺️ Generated Map",
                height=350,
            )

    if examples:
        gr.Examples(
            examples=examples,
            inputs=input_image,
            outputs=output_image,
            fn=predict,
            cache_examples=False,
            label="📸 Try an Example",
        )

    gr.Markdown(
        """
        ---
        **How it works:** The model uses a *U-Net* generator with skip connections
        trained adversarially against a *PatchGAN* discriminator. Input images are
        resized to **256 × 256** before inference.

        📎 [Kaggle Notebook](https://www.kaggle.com/code/dobariyanaitik/pix2pix-gan-to-generate-maps-from-satellite-images)
         | 
        📄 [Pix2Pix Paper (Isola et al., 2017)](https://arxiv.org/abs/1611.07004)
        """
    )

    submit_btn.click(fn=predict, inputs=input_image, outputs=output_image)

# ---------------------------------------------------------------------------
# Launch
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    demo.launch()