Spaces:
Running
Running
| import gradio as gr | |
| from ultralytics import YOLO | |
| from PIL import Image, ImageDraw | |
| import pytesseract | |
| import subprocess | |
| # Ensure Tesseract OCR is installed and detected | |
| TESSERACT_PATH = "/usr/bin/tesseract" | |
| pytesseract.pytesseract.tesseract_cmd = TESSERACT_PATH | |
| def check_tesseract(): | |
| """Check if Tesseract is installed and print its version.""" | |
| try: | |
| tesseract_version = subprocess.check_output([TESSERACT_PATH, "--version"]).decode("utf-8").split("\n")[0] | |
| print(f"Tesseract Version: {tesseract_version}") | |
| return True | |
| except Exception as e: | |
| print(f"Tesseract not found: {e}") | |
| return False | |
| # Load YOLO model (ensure best.pt exists in the working directory) | |
| YOLO_MODEL_PATH = "best.pt" | |
| model = YOLO(YOLO_MODEL_PATH, task='detect').to("cpu") | |
| def merge_boxes_into_lines(boxes, y_threshold=10): | |
| """Merge bounding boxes if they belong to the same text row.""" | |
| if len(boxes) == 0: | |
| return [] | |
| boxes = sorted(boxes, key=lambda b: b[1]) # Sort by y-axis (top position) | |
| 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: # Close enough to the previous line | |
| current_line[0] = min(current_line[0], x1) # Extend left boundary | |
| current_line[2] = max(current_line[2], x2) # Extend right boundary | |
| current_line[3] = max(current_line[3], y2) # Extend bottom boundary | |
| else: | |
| merged_lines.append(current_line) | |
| current_line = list(boxes[i]) | |
| merged_lines.append(current_line) | |
| return merged_lines | |
| def detect_and_ocr(image): | |
| """Detects text lines, draws bounding boxes, and runs OCR if available.""" | |
| 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) | |
| extracted_text_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)) | |
| if check_tesseract(): # If Tesseract is installed, run OCR | |
| try: | |
| ocr_text = pytesseract.image_to_string(cropped_line, lang="khm+eng").strip() | |
| if ocr_text: | |
| extracted_text_lines.append(ocr_text) | |
| except Exception as e: | |
| print(f"OCR failed for line {idx}: {e}") | |
| full_text = "\n".join(extracted_text_lines) if extracted_text_lines else "⚠️ OCR not available. Showing detected lines only." | |
| return original_image, full_text | |
| # Gradio UI | |
| with gr.Blocks() as iface: | |
| gr.Markdown("# 📜 Text Line Detection with Khmer OCR") | |
| gr.Markdown("## 📷 Upload an image to detect text lines and extract Khmer text") | |
| 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 Bounding Boxes") | |
| output_annotated = gr.Image(type="pil", label="Detected Text Lines") | |
| gr.Markdown("### 📝 Extracted Text (OCR Result)") | |
| output_text = gr.Textbox(label="Extracted Text", lines=10) | |
| image_input.upload( | |
| detect_and_ocr, | |
| inputs=image_input, | |
| outputs=[output_annotated, output_text] | |
| ) | |
| # 🚀 Ensure the app runs properly in Hugging Face Spaces | |
| if __name__ == "__main__": | |
| iface.launch(server_name="0.0.0.0", server_port=7860) | |