updated model

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,67 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# pyenv
+.python-version
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# vscode
+.vscode/
+
+# tmp folder
+.tmp

README.md

@@ -1,12 +1,13 @@
 ---
 title: Infrastructure Damage Assessment
-emoji: 🌍
-colorFrom: red
-colorTo: indigo
+emoji: 🏠
+color: green
 sdk: gradio
-sdk_version: 5.24.0
+sdk_version: 5.21.0
 app_file: app.py
 pinned: false
+license: mit
+short_description: Road and buildings damage assessment
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,445 @@
+import gradio as gr
+import numpy as np
+import os
+from PIL import Image
+from math import ceil, floor
+from numpy import ndarray
+from typing import Callable, List
+import scipy.signal
+import onnxruntime as ort
+from tqdm import tqdm
+
+# needed to run locally
+os.environ["GRADIO_TEMP_DIR"] = ".tmp"
+
+WINDOW_CACHE = dict()
+
+
+def _spline_window(window_size: int, power: int = 2) -> np.ndarray:
+    """Generates a 1-dimensional spline of order 'power' (typically 2), in the designated
+    window.
+    Args:
+        window_size (int): size of the interested window
+        power (int, optional): Order of the spline. Defaults to 2.
+    Returns:
+        np.ndarray: 1D spline
+    """
+    intersection = int(window_size / 4)
+    wind_outer = (
+        abs(2 * (scipy.signal.windows.triang(window_size))) ** power) / 2
+    wind_outer[intersection:-intersection] = 0
+    wind_inner = (
+        1 - (abs(2 * (scipy.signal.windows.triang(window_size) - 1)) ** power) / 2
+    )
+    wind_inner[:intersection] = 0
+    wind_inner[-intersection:] = 0
+    wind = wind_inner + wind_outer
+    wind = wind / np.average(wind)
+    return wind
+
+
+def _spline_2d(window_size: int, power: int = 2) -> ndarray:
+    """Makes a 1D window spline function, then combines it to return a 2D window function.
+    The 2D window is useful to smoothly interpolate between patches.
+    Args:
+        window_size (int): size of the window (patch)
+        power (int, optional): Which order for the spline. Defaults to 2.
+    Returns:
+        np.ndarray: numpy array containing a 2D spline function
+    """
+    # Memorization to avoid remaking it for every call
+    # since the same window is needed multiple times
+    wind = _spline_window(window_size, power)
+    # make it 2d
+    wind2 = wind[:, None] * wind[None, :]
+    wind2 = wind2 / np.max(wind2)
+    return wind2
+
+
+def _spline_4d(
+    window_size: int,
+    power: int = 2,
+    batch_size: int = 1,
+    channels: int = 1
+) -> ndarray:
+    """Makes a 4D window spline function
+    Same as the 2D version, but repeated across all channels and batch"""
+    global WINDOW_CACHE
+    key = f"{window_size}_{power}"
+    if key in WINDOW_CACHE:
+        wind4 = WINDOW_CACHE[key]
+    else:
+        wind2 = _spline_2d(window_size, power)
+        wind4 = wind2[None, None, :, :] * np.ones((batch_size, channels, 1, 1))
+        WINDOW_CACHE[key] = wind2
+    return wind4
+
+
+def pad_image(image: np.array, tile_size: int, subdivisions: int) -> np.array:
+    """Add borders to the given image for a "valid" border pattern according to "window_size" and "subdivisions".
+    Image is expected as a numpy array with shape (width, height, channels).
+    Args:
+        image (torch.Tensor): input image, 3D channels-last tensor
+        tile_size (int): size of a single patch, useful to compute padding
+        subdivisions (int): amount of overlap, useful for padding
+    Returns:
+        torch.Tensor: same image, padded specularly by a certain amount in every direction
+    """
+    step = tile_size // subdivisions
+    _, in_h, in_w = image.shape
+    pad_h = step - (in_h % step)
+    pad_w = step - (in_w % step)
+    pad_h_l = pad_h // 2
+    pad_h_r = (pad_h // 2) + (pad_h % 2)
+    pad_w_l = pad_w // 2
+    pad_w_r = (pad_w // 2) + (pad_w % 2)
+    pad = int(round(tile_size * (1 - 1.0 / subdivisions)))
+    image = np.pad(
+        image,
+        ((0, 0), (pad + pad_h_l, pad + pad_h_r), (pad + pad_w_l, pad + pad_w_r)),
+        mode="reflect",
+    )
+    return image, [pad + pad_h_l, pad + pad_h_r, pad + pad_w_l, pad + pad_w_r]
+
+
+def unpad_image(padded_image: ndarray, pads) -> ndarray:
+    """Reverts changes made by 'pad_image'. The same padding is removed, so tile_size and subdivisions
+    must be coherent.
+
+    Args:
+        padded_image (torch.Tensor): image with padding still applied
+        tile_size (int): size of a single patch
+        subdivisions (int): subdivisions to compute overlap
+
+    Returns:
+        torch.Tensor: image without padding, 2D channels-last tensor
+    """
+    pad_left, pad_right, pad_top, pad_bottom = pads
+    # crop the image left, right, top and bottom
+    # get number of dimensions of padded_image
+    n_dims = len(padded_image.shape)
+    # if padded_image is 2d
+    if n_dims == 2:
+        result = padded_image[pad_left:-pad_right, pad_top:-pad_bottom]
+    # if padded_image is 3d
+    elif n_dims == 3:
+        result = padded_image[:, pad_left:-pad_right, pad_top:-pad_bottom]
+    else:
+        raise ValueError(
+            f"padded_image has {n_dims} dimensions, expected 2 or 3.")
+    return result
+
+
+def windowed_generator(
+    padded_image: ndarray, window_size: int, subdivisions: int, batch_size: int = None
+):
+    """Generator that yield tiles grouped by batch size.
+    Args:
+        padded_image (np.ndarray): input image to be processed (already padded), supposed channels-first
+        window_size (int): size of a single patch
+        subdivisions (int): subdivision count on each patch to compute the step
+        batch_size (int, optional): amount of patches in each batch. Defaults to None.
+
+    Yields:
+        Tuple[List[tuple], np.ndarray]: list of coordinates and respective patches as single batch array
+    """
+    step = window_size // subdivisions
+    channel, width, height = padded_image.shape
+    batch_size = batch_size or 1
+    batch = []
+    coords = []
+    for x in range(0, width - window_size + 1, step):
+        for y in range(0, height - window_size + 1, step):
+            coords.append((x, y))
+            # extract the tile, place channels first for batch
+            tile = padded_image[:, x: x + window_size, y: y + window_size]
+            batch.append(tile)
+            # yield the batch once full and restore lists right after
+            if len(batch) == batch_size:
+                yield coords, np.stack(batch)
+                coords.clear()
+                batch.clear()
+    # handle last (possibly unfinished) batch
+    if len(batch) > 0:
+        yield coords, np.stack(batch)
+
+
+def reconstruct(
+    canvas: ndarray, tile_size: int, coords: List[tuple], predictions: ndarray
+) -> ndarray:
+    """Helper function that iterates the result batch onto the given canvas to reconstruct
+    the final result batch after batch.
+    Args:
+        canvas (torch.Tensor): container for the final image.
+        tile_size (int): size of a single patch.
+        coords (List[tuple]): list of pixel coordinates corresponding to the batch items
+        predictions (torch.Tensor): array containing patch predictions, shape (batch, tile_size, tile_size, num_classes)
+
+    Returns:
+        torch.Tensor: the updated canvas, shape (padded_w, padded_h, num_classes)
+    """
+    for (x, y), patch in zip(coords, predictions):
+        # get canvas number of dimensions
+        n_dims = len(canvas.shape)
+        # if canvas is 2d
+        if n_dims == 2:
+            canvas[x: x + tile_size, y: y + tile_size] += patch
+        # if canvas is 3d
+        elif n_dims == 3:
+            canvas[:, x: x + tile_size, y: y + tile_size] += patch
+        else:
+            raise ValueError(
+                f"Canvas has {n_dims} dimensions, expected 2 or 3.")
+    return canvas
+
+
+def predict_smooth_windowing(
+    image: ndarray,
+    tile_size: int,
+    subdivisions: int,
+    prediction_fn: Callable,
+    batch_size: int = 1,
+    out_dim: int = 1,
+) -> np.ndarray:
+    """Allows to predict a large image in one go, dividing it in squared, fixed-size tiles and smoothly
+    interpolating over them to produce a single, coherent output with the same dimensions.
+    Args:
+        image (np.ndarray): input image, expected a 3D vector
+        tile_size (int): size of each squared tile
+        subdivisions (int): number of subdivisions over the single tile for overlaps
+        prediction_fn (Callable): callback that takes the input batch and returns an output tensor
+        batch_size (int, optional): size of each batch. Defaults to None.
+        channels_first (int, optional): whether the input image is channels-first or not
+        mirrored (bool, optional): whether to use dihedral predictions (every simmetry). Defaults to False.
+
+    Returns:
+        np.ndarray: numpy array with dimensions (w, h), containing smooth predictions
+    """
+    img, pads = pad_image(image=image, tile_size=tile_size,
+                          subdivisions=subdivisions)
+    spline = _spline_4d(window_size=tile_size, power=2)
+    # canvas = np.zeros(img.shape[1], img.shape[2])
+    canvas = np.zeros((out_dim, img.shape[1], img.shape[2]))
+    loop = tqdm(windowed_generator(
+        padded_image=img,
+        window_size=tile_size,
+        subdivisions=subdivisions,
+        batch_size=batch_size,
+    ))
+    for coords, batch in loop:
+        pred_batch = prediction_fn(batch)  # .permute(0, 2, 3, 1)
+        # must be 3d for reconstruction to work
+        pred_batch = pred_batch * spline
+        canvas = reconstruct(
+            canvas, tile_size=tile_size, coords=coords, predictions=pred_batch
+        )
+    prediction = unpad_image(canvas, pads=pads)
+    return prediction
+
+
+def center_pad(x, padding, div_factor=32, mode="reflect"):
+    # center pad with different padding for each city
+    # pads the image with the same padding on all sides
+    # the output size must be at least the size + 2*padding
+    # and divisible by div_factor
+    # first, compute the size of the padded image
+    size_x = x.shape[3]
+    size_y = x.shape[2]
+    # get the min padding
+    min_padding_x = size_x + 2 * padding
+    min_padding_y = size_y + 2 * padding
+    # get the new size
+    new_size_x = int(ceil(min_padding_x / div_factor) * div_factor)
+    new_size_y = int(ceil(min_padding_y / div_factor) * div_factor)
+    # get the padding
+    pad_x = new_size_x - size_x
+    pad_y = new_size_y - size_y
+    pad_left = int(floor(pad_x / 2))
+    pad_right = int(ceil(pad_x / 2))
+    pad_top = int(floor(pad_y / 2))
+    pad_bottom = int(ceil(pad_y / 2))
+    if pad_x > size_x or pad_y > size_y:
+        padded = np.pad(
+            x,
+            (
+                (0, 0),
+                (0, 0),
+                (int(floor(size_x / 2)), int(ceil(size_x / 2))),
+                (int(floor(size_y / 2)), int(ceil(size_y / 2))),
+            ),
+            mode=mode,
+        )
+        # and then pad to size
+        padded = np.pad(
+            x,
+            (
+                (0, 0),
+                (0, 0),
+                (int(floor(new_size_x / 2)), int(ceil(new_size_x / 2))),
+                (int(floor(new_size_y / 2)), int(ceil(new_size_y / 2))),
+            ),
+            mode=mode,
+        )
+    else:
+        padded = np.pad(
+            x,
+            (
+                (0, 0),
+                (0, 0),
+                (pad_top, pad_bottom),
+                (pad_left, pad_right),
+            ),
+            mode=mode,
+        )
+    paddings = (pad_top, pad_bottom, pad_left, pad_right)
+    return padded, paddings
+
+
+class ChangeDetectionModel:
+    def __init__(self):
+        path = "assets/models/change_detection.onnx"
+        self.model = ort.InferenceSession(path)
+        self.size = 256
+        self.subdivisions = 2
+        self.batch_size = 2
+        self.out_dim = 1
+
+    def forward(self, x):
+        assert x.ndim == 3, "Expected 3D tensor"
+        # remove batch dimension
+        x = x/255
+        # cast to fp32
+        x = x.astype(np.float32)
+        pred = predict_smooth_windowing(
+            image=x,
+            tile_size=self.size,
+            subdivisions=self.subdivisions,
+            prediction_fn=self.callback,
+            batch_size=self.batch_size,
+            out_dim=self.out_dim
+        )
+        # apply sigmoid
+        pred = 1 / (1 + np.exp(-pred))
+        # set pred to 0 if less than 0.25, 1 if more than .25 and less then .5
+        # to 2 if more than .5 and less than .75, and to 3 if more than .75
+        pred = pred * 3
+        pred = np.round(pred)
+        return pred[0]
+
+    def callback(self, x: ndarray) -> ndarray:
+        # run onnx inference
+        out = self.model.run(None, {"input": x})[0]
+        return out
+
+
+class LocalizationModel:
+    def __init__(self):
+        path = "assets/models/localization.onnx"
+        self.model = ort.InferenceSession(path)
+        self.size = 384
+        self.subdivisions = 2
+        self.batch_size = 2
+        self.out_dim = 3
+
+    def forward(self, x):
+        assert x.ndim == 3, "Expected 3D tensor"
+        # remove batch dimension
+        x = x/255
+        # cast to fp32
+        x = x.astype(np.float32)
+        pred = predict_smooth_windowing(
+            image=x,
+            tile_size=self.size,
+            subdivisions=self.subdivisions,
+            prediction_fn=self.callback,
+            batch_size=self.batch_size,
+            out_dim=self.out_dim
+        )
+        # compute the argmax
+        pred = np.argmax(pred, axis=0)
+        return pred
+
+    def callback(self, x: ndarray) -> ndarray:
+        # run onnx inference
+        out = self.model.run(None, {"input": x})[0]
+        return out
+
+
+def infer(image1, image2):
+    localization_model = LocalizationModel()
+    change_detection_model = ChangeDetectionModel()
+    # resize image1 to image2
+    image1 = image1.resize(image2.size)
+    # half resolution
+    image1 = image1.resize((image1.width // 2, image1.height // 2))
+    image2 = image2.resize((image2.width // 2, image2.height // 2))
+    # convert images to numpy arrays
+    image1 = np.array(image1)
+    image2 = np.array(image2)
+    # from whc to cwh
+    image1_array = np.transpose(image1, (2, 0, 1))
+    image2_array = np.transpose(image2, (2, 0, 1))
+    output_image1 = localization_model.forward(image1_array)
+    # concatenate the images
+    cat_image_array = np.concatenate([image1_array, image2_array], axis=0)
+    output_image2 = change_detection_model.forward(cat_image_array)
+    output_image1_color = np.zeros(
+        (output_image1.shape[0], output_image1.shape[1], 3))
+    # set output_image1_color to gray where output_image1 is 1
+    output_image1_color[output_image1 == 0] = [0, 0, 0]  # Class 0: bg
+    output_image1_color[output_image1 == 1] = [150, 150, 150]  # Class 1: road
+    output_image1_color[output_image1 == 2] = [200, 0, 0]  # Class 2: house
+    # compute average of output_image1_color and input1
+    output_image1_color = (output_image1_color*0.5 + image1*0.5)
+    output_image1 = Image.fromarray(output_image1_color.astype(np.uint8))
+    output_image2_color = np.zeros(
+        (output_image2.shape[0], output_image2.shape[1], 3))
+    output_image2_color[output_image2 == 0] = [0, 0, 0]  # Class 0: no change
+    output_image2_color[output_image2 == 1] = [0, 255, 0]  # Class 1: minor change
+    output_image2_color[output_image2 == 2] = [255, 255, 0]  # Class 2: major change
+    output_image2_color[output_image2 == 3] = [255, 0, 0]  # Class 3: destroyed
+    output_image2_color = output_image2_color*0.5 + image2*0.5
+    output_image2 = Image.fromarray(output_image2_color.astype(np.uint8))
+    return output_image1, output_image2
+
+
+# Define sample image pairs
+sample_images = [
+    ["assets/data/bata_1_pre.png", "assets/data/bata_1_post.png"],
+    ["assets/data/bata_2_pre.png", "assets/data/bata_2_post.png"],
+    ["assets/data/beirut_1_pre.png", "assets/data/beirut_1_post.png"]
+]
+
+with gr.Blocks() as demo:
+    gr.Markdown("## Infrastructure Damage Assessment")
+    # description
+    gr.Markdown(
+        "This is a demo for infrastructure damage assessment using satellite images.\
+        It contains two models: one for localization and the other for change detection. \
+        The localization model is used to segment the image into three classes: background (in black), road (in grey), and houses (in red). \
+        The change detection model is used to detect changes between two images.\
+        The output of the change detection model is colored as follows: no change (in black), minor change (in green), major change (in yellow), and destroyed (in red).\
+        The output of the localization model (on the left) is blended with the pre-disaster image to highlight the areas of interest.\
+        The output of the change detection model (on the right) is blended with the post-disaster image to highlight the changes.\
+        You can upload your own images or use the sample images provided below."
+    )
+    gr.Markdown(
+        "Note: the models run at half resolution for faster inference, \
+        so the output images will be less accurate than the full-resolution models.\
+        It still takes a few minutes to run the inference, so please be patient."
+    )
+    with gr.Row():  # Place images in the same row
+        with gr.Column(scale=1):  # First half-column
+            input_image1 = gr.Image(label="Pre-disaster Image", type="pil")
+        with gr.Column(scale=1):  # Second half-column
+            input_image2 = gr.Image(label="Post-disaster Image", type="pil")
+    with gr.Row():  # Row for output images
+        output_image1 = gr.Image(label="Roads and buildings localization", type="pil")
+        output_image2 = gr.Image(label="Change detection", type="pil")
+    submit_button = gr.Button("Run Inference")
+    examples = gr.Examples(examples=sample_images, inputs=[
+                           input_image1, input_image2])
+    submit_button.click(fn=infer, inputs=[input_image1, input_image2], outputs=[
+                        output_image1, output_image2])
+
+demo.launch()

assets/models/change_detection.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:02b62f43110535754fea7a6aedc9a2c1719b935830b220853d17afe8565297dd
+size 230737192

assets/models/localization.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:545413e05e89b0b3382b4aade9e8f2a5f8c825d5c4e91b77b6c4f57401ff3086
+size 295

requirements.txt

@@ -0,0 +1,7 @@
+gradio==5.21.0
+numpy==2.2.4
+onnxruntime==1.21.0
+pandas==2.2.3
+pillow==11.1.0
+scipy==1.15.2
+tqdm==4.67.1