Norakneath commited on
Commit
e6db2ab
·
verified ·
1 Parent(s): 7ab587d

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +49 -0
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)