""" 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()