app.py

"""
Rebar Detection with YOLO Models - Gradio Application

This application provides a web interface for detecting rebars in GPR images
using YOLO-based ONNX models. It includes features for tiled processing
with configurable overlap and confidence thresholds.

Developer: Ahmed Elseicy
Email: ahmedmossadibrahim.elseicy@uvigo.gal
Date: July 30, 2025
"""
import gradio as gr
import numpy as np
import onnxruntime as ort
from PIL import Image, ImageDraw
import os

# --- Configuration ---

MODEL_DIR = "models"
EXAMPLE_DIR = "examples"
# Ensure the directories exist
if not os.path.exists(MODEL_DIR):
    os.makedirs(MODEL_DIR)
if not os.path.exists(EXAMPLE_DIR):
    os.makedirs(EXAMPLE_DIR)

# --- Model Loading ---


def get_available_models():
    if not os.path.exists(MODEL_DIR) or not os.listdir(MODEL_DIR):
        print(
            f"Warning: No models found in '{MODEL_DIR}'. Please upload your .onnx files.")
        return ["No models found"]
    # Strip the .onnx extension for a cleaner display name
    return [os.path.splitext(f)[0] for f in os.listdir(MODEL_DIR) if f.endswith(".onnx")]


AVAILABLE_MODELS = get_available_models()

# --- Helper Function ---


def create_blank_image(width=512, height=512):
    """Creates a blank white PIL image."""
    return Image.new('RGB', (width, height), 'white')

# --- Image Processing and Inference Logic ---


def slice_image(image, tile_size=(256, 256), overlap_ratio=0.2):
    """Slices an image into overlapping tiles."""
    img_w, img_h = image.size
    tile_w, tile_h = tile_size
    stride_w = int(tile_w * (1 - overlap_ratio))
    stride_h = int(tile_h * (1 - overlap_ratio))

    for y in range(0, img_h, stride_h):
        for x in range(0, img_w, stride_w):
            box = (x, y, x + tile_w, y + tile_h)
            if box[2] > img_w:
                box = (img_w - tile_w, box[1], img_w, box[3])
            if box[3] > img_h:
                box = (box[0], img_h - tile_h, box[2], img_h)

            yield image.crop(box), (box[0], box[1])

            if box[2] >= img_w:
                break
        if box[3] >= img_h:
            break


def run_yolo_inference(session, image_tile):
    """
    Runs inference using a YOLO ONNX model and returns processed detections.
    """
    # 1. Preprocess the image
    input_image = np.array(image_tile.resize(
        (256, 256)), dtype=np.float32) / 255.0
    input_image = np.expand_dims(input_image, axis=0)
    input_image = np.transpose(input_image, (0, 3, 1, 2))

    # 2. Run inference
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name
    result = session.run([output_name], {input_name: input_image})[0]

    # 3. Post-process the output
    detections = result[0].T
    boxes = []
    scores = []

    for row in detections:
        # For object detection, the row format is typically [cx, cy, w, h, class_confidence, ...]
        class_probs = row[4:]
        class_id = np.argmax(class_probs)
        confidence = class_probs[class_id]

        # Extract box and convert from [center_x, center_y, width, height] to [x1, y1, x2, y2]
        cx, cy, w, h = row[:4]
        x1 = cx - w / 2
        y1 = cy - h / 2
        x2 = cx + w / 2
        y2 = cy + h / 2

        boxes.append([x1, y1, x2, y2])
        scores.append(confidence)

    return np.array(boxes), np.array(scores)


def non_max_suppression(boxes, scores, iou_threshold):
    """Performs Non-Maximum Suppression to merge overlapping boxes."""
    if len(boxes) == 0:
        return []

    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]

    areas = (x2 - x1) * (y2 - y1)
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)

        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        w = np.maximum(0.0, xx2 - xx1)
        h = np.maximum(0.0, yy2 - yy1)

        intersection = w * h
        iou = intersection / (areas[i] + areas[order[1:]] - intersection)

        inds = np.where(iou <= iou_threshold)[0]
        order = order[inds + 1]

    return keep


def detect_rebars(model_name, input_image, overlap_ratio, confidence_threshold, iou_threshold):
    """Main function to orchestrate the detection process."""
    if model_name is None or model_name == "No models found" or input_image is None:
        return create_blank_image(), "Please select a model and upload an image."

    try:
        # Add the .onnx extension back to the model name for file path
        model_path = os.path.join(MODEL_DIR, model_name + ".onnx")
        session = ort.InferenceSession(model_path)
    except Exception as e:
        return create_blank_image(), f"Error loading model: {e}"

    # Convert input image to RGB to ensure drawing works correctly
    original_image = Image.fromarray(input_image).convert("RGB")
    all_boxes = []
    all_scores = []

    for tile, (x_offset, y_offset) in slice_image(original_image, overlap_ratio=overlap_ratio):
        try:
            boxes_on_tile, scores_on_tile = run_yolo_inference(session, tile)
            for box, score in zip(boxes_on_tile, scores_on_tile):
                if score >= confidence_threshold:
                    x1, y1, x2, y2 = box
                    all_boxes.append(
                        [x1 + x_offset, y1 + y_offset, x2 + x_offset, y2 + y_offset])
                    all_scores.append(score)
        except Exception as e:
            return create_blank_image(), f"An error occurred during inference: {e}."

    if not all_boxes:
        return original_image, "Detection complete. No rebars found."

    # Apply Non-Maximum Suppression to all collected boxes
    final_indices = non_max_suppression(
        np.array(all_boxes), np.array(all_scores), iou_threshold)

    # Create a copy to draw on
    stitched_image = original_image.copy()
    draw = ImageDraw.Draw(stitched_image)

    for i in final_indices:
        box = all_boxes[i]
        draw.rectangle(box, outline="red", width=3)

    status_message = f"Detection complete. Found {len(final_indices)} rebars."
    return stitched_image, status_message


# --- Gradio Web Interface ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Rebar Detection using YOLO Models")
    gr.Markdown(
        """
        Select a model, upload a GPR image, set processing parameters, and the model will predict rebar locations.
        **Note:** This is a prototype implementation running on a vCPU. The image slicing is currently based on pixels. For practical applications, slicing should be performed by distance (meters) in the horizontal direction.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            model_selector = gr.Dropdown(
                label="Select Model", choices=AVAILABLE_MODELS, value=None)
            image_input = gr.Image(type="numpy", label="Upload GPR Image")

            # Add example images for users to test
            gr.Examples(
                examples=os.path.join(os.path.dirname(__file__), EXAMPLE_DIR),
                inputs=image_input,
                label="Example Images"
            )

            overlap_slider = gr.Slider(
                minimum=0.0, maximum=0.9, step=0.05, value=0.2, label="Overlap Ratio")
            confidence_slider = gr.Slider(
                minimum=0.0, maximum=1.0, step=0.05, value=0.25, label="Confidence Threshold")
            iou_slider = gr.Slider(
                minimum=0.0, maximum=1.0, step=0.05, value=0.45, label="IoU Threshold (for NMS)")
            submit_btn = gr.Button("Detect Rebars", variant="primary")

        with gr.Column(scale=2):
            status_output = gr.Textbox(label="Status", interactive=False)
            # Set height to "auto" to prevent shrinking with wide images
            image_output = gr.Image(
                type="pil", label="Detection Result", height="auto")


    submit_btn.click(
        fn=detect_rebars,
        inputs=[model_selector, image_input,
                overlap_slider, confidence_slider, iou_slider],
        outputs=[image_output, status_output]
    )

    
    attribution_info = """
    ## 📜 Paper Information
    This Space is based on the research presented in our paper for IWAGPR25:
    ```bibtex
    @inproceedings{elseicy2025rebar,
        title     = {Preliminary Study on Automating Rebar Detection in Reinforced Concrete Structures Using YOLOv11 and GPR Data},
        author    = {Elseicy, Ahmed and Solla, Mercedes and Novo, Alexandre},
        year      = {2025},
        month     = {July},
        booktitle = {2025 13th International Workshop on Advanced Ground Penetrating Radar (IWAGPR)},
        publisher = {IEEE},
        pages     = {323--328},
        isbn      = {979-8-3315-2335-0},
        issn      = {2687-7899}
    }
    ```
    ## 💾 Dataset Reference

    The full models and the dataset used in the project are published in Zenodo [DOI: 10.5281/zenodo.16638791](https://doi.org/10.5281/zenodo.16638791).

    ## 💰 Funding Acknowledgement

    This research and development were made possible through the OVERSIGHT project (PID2022-138526OB-I00) funded by MICIU/AEI/10.13039/501100011033/FEDER, UE.
    Grant PREP2022-000030 for the training of predoctoral researchers funded by MICIU/
    AEI/10.13039/501100011033 and by FSE+. 
    M. Solla acknowledges the Grant RYC2019–026604–I funded by MICIU/
    AEI/10.13039/501100011033 and by “ESF Investing in
    your future”.
"""
    with gr.Accordion("Show Publication Info, Dataset & Funding Details", open=True):
        gr.Markdown(attribution_info)

if __name__ == "__main__":
    demo.launch()