Spaces:
Running
Running
| """ | |
| 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() | |