wjbmattingly commited on
Commit
d138ef9
·
verified ·
1 Parent(s): 79f236a

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +122 -0
app.py ADDED
@@ -0,0 +1,122 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Tuple
2
+ import gradio as gr
3
+ import supervision as sv
4
+ import numpy as np
5
+ from PIL import Image
6
+ from huggingface_hub import hf_hub_download
7
+ from ultralytics import YOLO
8
+
9
+ # Load the YOLO model from Hugging Face
10
+ model_path = hf_hub_download(
11
+ repo_id="cultural-heritage/medieval-manuscript-yolov11",
12
+ filename="medieval-yolov11n.pt"
13
+ )
14
+ # Load the YOLO model from local path
15
+ model = YOLO(model_path)
16
+
17
+ # Create annotators
18
+ LABEL_ANNOTATOR = sv.LabelAnnotator(text_color=sv.Color.BLACK)
19
+ BOX_ANNOTATOR = sv.BoxAnnotator()
20
+
21
+ def detect_and_annotate(
22
+ image: np.ndarray,
23
+ conf_threshold: float,
24
+ iou_threshold: float
25
+ ) -> np.ndarray:
26
+ # Perform inference
27
+ results = model.predict(
28
+ image,
29
+ conf=conf_threshold,
30
+ iou=iou_threshold
31
+ )[0]
32
+
33
+ # Convert results to supervision Detections
34
+ boxes = results.boxes.xyxy.cpu().numpy()
35
+ confidence = results.boxes.conf.cpu().numpy()
36
+ class_ids = results.boxes.cls.cpu().numpy().astype(int)
37
+
38
+ # Create Detections object
39
+ detections = sv.Detections(
40
+ xyxy=boxes,
41
+ confidence=confidence,
42
+ class_id=class_ids
43
+ )
44
+
45
+ # Create labels with confidence scores
46
+ labels = [
47
+ f"{results.names[class_id]} ({conf:.2f})"
48
+ for class_id, conf
49
+ in zip(class_ids, confidence)
50
+ ]
51
+
52
+ # Annotate image
53
+ annotated_image = image.copy()
54
+ annotated_image = BOX_ANNOTATOR.annotate(scene=annotated_image, detections=detections)
55
+ annotated_image = LABEL_ANNOTATOR.annotate(scene=annotated_image, detections=detections, labels=labels)
56
+
57
+ return annotated_image
58
+
59
+ # Create Gradio interface
60
+ with gr.Blocks() as demo:
61
+ gr.Markdown("# Medieval Manuscript Detection with YOLO")
62
+
63
+ with gr.Row():
64
+ with gr.Column():
65
+ input_image = gr.Image(
66
+ label="Input Image",
67
+ type='numpy'
68
+ )
69
+ with gr.Accordion("Detection Settings", open=True):
70
+ with gr.Row():
71
+ conf_threshold = gr.Slider(
72
+ label="Confidence Threshold",
73
+ minimum=0.0,
74
+ maximum=1.0,
75
+ step=0.05,
76
+ value=0.25,
77
+ )
78
+ iou_threshold = gr.Slider(
79
+ label="IoU Threshold",
80
+ minimum=0.0,
81
+ maximum=1.0,
82
+ step=0.05,
83
+ value=0.45,
84
+ info="Decrease for stricter detection, increase for more overlapping boxes"
85
+ )
86
+ with gr.Row():
87
+ clear_btn = gr.Button("Clear")
88
+ detect_btn = gr.Button("Detect", variant="primary")
89
+
90
+ with gr.Column():
91
+ output_image = gr.Image(
92
+ label="Detection Result",
93
+ type='numpy'
94
+ )
95
+
96
+ def process_image(
97
+ image: np.ndarray,
98
+ conf_threshold: float,
99
+ iou_threshold: float
100
+ ) -> Tuple[np.ndarray, np.ndarray]:
101
+ if image is None:
102
+ return None, None
103
+ annotated_image = detect_and_annotate(image, conf_threshold, iou_threshold)
104
+ return image, annotated_image
105
+
106
+ def clear():
107
+ return None, None
108
+
109
+ # Connect buttons to functions
110
+ detect_btn.click(
111
+ process_image,
112
+ inputs=[input_image, conf_threshold, iou_threshold],
113
+ outputs=[input_image, output_image]
114
+ )
115
+ clear_btn.click(
116
+ clear,
117
+ inputs=None,
118
+ outputs=[input_image, output_image]
119
+ )
120
+
121
+ if __name__ == "__main__":
122
+ demo.launch(debug=True, show_error=True)