Spaces:
Running
Running
| import gradio as gr | |
| from ultralytics import YOLO | |
| from PIL import Image, ImageDraw | |
| YOLO_MODEL_PATH = "best.pt" | |
| model = YOLO(YOLO_MODEL_PATH, task='detect').to("cpu") | |
| def merge_boxes_into_lines(boxes, y_threshold=10): | |
| if len(boxes) == 0: | |
| return [] | |
| boxes = sorted(boxes, key=lambda b: b[1]) | |
| merged_lines = [] | |
| current_line = list(boxes[0]) | |
| for i in range(1, len(boxes)): | |
| x1, y1, x2, y2 = boxes[i] | |
| if abs(y1 - current_line[1]) < y_threshold: | |
| current_line[0] = min(current_line[0], x1) | |
| current_line[2] = max(current_line[2], x2) | |
| current_line[3] = max(current_line[3], y2) | |
| else: | |
| merged_lines.append(current_line) | |
| current_line = list(boxes[i]) | |
| merged_lines.append(current_line) | |
| return merged_lines | |
| def detect_and_crop_lines(image): | |
| image = Image.fromarray(image) | |
| original_image = image.copy() | |
| results = model.predict(image, conf=0.3, iou=0.5, device="cpu") | |
| detected_boxes = results[0].boxes.xyxy.tolist() | |
| detected_boxes = [list(map(int, box)) for box in detected_boxes] | |
| merged_boxes = merge_boxes_into_lines(detected_boxes) | |
| draw = ImageDraw.Draw(original_image) | |
| cropped_lines = [] | |
| for idx, (x1, y1, x2, y2) in enumerate(merged_boxes): | |
| draw.rectangle([x1, y1, x2, y2], outline="blue", width=2) | |
| draw.text((x1, y1 - 10), f"Line {idx}", fill="blue") | |
| cropped_line = image.crop((x1, y1, x2, y2)) | |
| # Resize the cropped line to a smaller fixed width while maintaining aspect ratio | |
| fixed_width = 200 | |
| aspect_ratio = cropped_line.height / cropped_line.width | |
| new_height = int(fixed_width * aspect_ratio) | |
| cropped_line = cropped_line.resize((fixed_width, new_height), Image.LANCZOS) | |
| cropped_lines.append(cropped_line) | |
| return original_image, cropped_lines | |
| with gr.Blocks() as iface: | |
| gr.Markdown("# Text Line Detection") | |
| gr.Markdown("## Input your custom image for text line detection") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| gr.Markdown("### Upload Image") | |
| image_input = gr.Image(type="numpy", label="Upload an image") | |
| with gr.Column(scale=1): | |
| gr.Markdown("### Annotated Image with Detected Lines") | |
| output_annotated = gr.Image(type="pil", label="Detected Text Lines") | |
| gr.Markdown("### Cropped Text Lines (Small Fixed-Size Previews)") | |
| cropped_output = gr.Gallery(label="Detected Text Lines", columns=1, preview=True) | |
| def process_and_display(image): | |
| annotated_img, cropped_imgs = detect_and_crop_lines(image) | |
| return annotated_img, cropped_imgs | |
| image_input.upload( | |
| process_and_display, | |
| inputs=image_input, | |
| outputs=[output_annotated, cropped_output] | |
| ) | |
| iface.launch() | |