Spaces:
Sleeping
Sleeping
Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,49 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import gradio as gr
|
2 |
+
from ultralytics import YOLO
|
3 |
+
from PIL import Image, ImageDraw
|
4 |
+
|
5 |
+
# Load YOLO model (ensure best.pt exists in the working directory)
|
6 |
+
YOLO_MODEL_PATH = "best (3).pt"
|
7 |
+
model = YOLO(YOLO_MODEL_PATH, task='detect').to("cpu")
|
8 |
+
|
9 |
+
def detect_text_lines(image):
|
10 |
+
"""Detects text lines and draws bounding boxes on the image."""
|
11 |
+
image = Image.fromarray(image)
|
12 |
+
original_image = image.copy()
|
13 |
+
|
14 |
+
# Run YOLO text detection
|
15 |
+
results = model.predict(image, conf=0.1, iou=0.2, device="cpu")
|
16 |
+
detected_boxes = results[0].boxes.xyxy.tolist()
|
17 |
+
detected_boxes = [list(map(int, box)) for box in detected_boxes]
|
18 |
+
|
19 |
+
# Draw bounding boxes on the image
|
20 |
+
draw = ImageDraw.Draw(original_image)
|
21 |
+
for idx, (x1, y1, x2, y2) in enumerate(detected_boxes):
|
22 |
+
draw.rectangle([x1, y1, x2, y2], outline="blue", width=2)
|
23 |
+
draw.text((x1, y1 - 10), f"Line {idx}", fill="blue")
|
24 |
+
|
25 |
+
return original_image
|
26 |
+
|
27 |
+
# Gradio UI
|
28 |
+
with gr.Blocks() as iface:
|
29 |
+
gr.Markdown("# 📜 Text Line Detection with YOLO")
|
30 |
+
gr.Markdown("## 📷 Upload an image to detect text lines")
|
31 |
+
|
32 |
+
with gr.Row():
|
33 |
+
with gr.Column(scale=1):
|
34 |
+
gr.Markdown("### 📤 Upload Image")
|
35 |
+
image_input = gr.Image(type="numpy", label="Upload an image")
|
36 |
+
|
37 |
+
with gr.Column(scale=1):
|
38 |
+
gr.Markdown("### 🖼 Annotated Image with Bounding Boxes")
|
39 |
+
output_annotated = gr.Image(type="pil", label="Detected Text Lines")
|
40 |
+
|
41 |
+
image_input.upload(
|
42 |
+
detect_text_lines,
|
43 |
+
inputs=image_input,
|
44 |
+
outputs=output_annotated
|
45 |
+
)
|
46 |
+
|
47 |
+
# 🚀 Ensure the app runs properly in Hugging Face Spaces
|
48 |
+
if __name__ == "__main__":
|
49 |
+
iface.launch(server_name="0.0.0.0", server_port=7860)
|