Spaces:
Running
Running
wjm55
commited on
Commit
·
132787d
1
Parent(s):
5ab39af
init
Browse files- Dockerfile +17 -0
- app.py +67 -0
- gradio-app.py +145 -0
- requirements.txt +5 -0
- test.py +24 -0
Dockerfile
ADDED
@@ -0,0 +1,17 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Read the doc: https://huggingface.co/docs/hub/spaces-sdks-docker
|
2 |
+
# you will also find guides on how best to write your Dockerfile
|
3 |
+
|
4 |
+
|
5 |
+
FROM python:3.10
|
6 |
+
RUN apt-get update && apt-get install -y libgl1-mesa-glx
|
7 |
+
RUN useradd -m -u 1000 user
|
8 |
+
USER user
|
9 |
+
ENV PATH="/home/user/.local/bin:$PATH"
|
10 |
+
|
11 |
+
WORKDIR /app
|
12 |
+
|
13 |
+
COPY --chown=user ./requirements.txt requirements.txt
|
14 |
+
RUN pip install --no-cache-dir --upgrade -r requirements.txt
|
15 |
+
|
16 |
+
COPY --chown=user . /app
|
17 |
+
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]
|
app.py
ADDED
@@ -0,0 +1,67 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from fastapi import FastAPI, UploadFile
|
2 |
+
from ultralytics import YOLO
|
3 |
+
import io
|
4 |
+
from PIL import Image
|
5 |
+
import numpy as np
|
6 |
+
import os
|
7 |
+
from huggingface_hub import hf_hub_download
|
8 |
+
from ultralytics import YOLO
|
9 |
+
import requests
|
10 |
+
import supervision as sv
|
11 |
+
|
12 |
+
|
13 |
+
def init_model(model_id: str):
|
14 |
+
# Define models
|
15 |
+
MODEL_OPTIONS = {
|
16 |
+
"YOLOv11-Nano": "medieval-yolov11n.pt",
|
17 |
+
"YOLOv11-Small": "medieval-yolov11s.pt",
|
18 |
+
"YOLOv11-Medium": "medieval-yolov11m.pt",
|
19 |
+
"YOLOv11-Large": "medieval-yolov11l.pt",
|
20 |
+
"YOLOv11-XLarge": "medieval-yolov11x.pt"
|
21 |
+
}
|
22 |
+
if model_id in MODEL_OPTIONS:
|
23 |
+
model_path = hf_hub_download(
|
24 |
+
repo_id="biglam/medieval-manuscript-yolov11",
|
25 |
+
filename=MODEL_OPTIONS[model_id]
|
26 |
+
)
|
27 |
+
return YOLO(model_path)
|
28 |
+
else:
|
29 |
+
raise ValueError(f"Model {model_id} not found")
|
30 |
+
|
31 |
+
app = FastAPI()
|
32 |
+
|
33 |
+
|
34 |
+
@app.post("/predict")
|
35 |
+
async def predict(image: UploadFile,
|
36 |
+
model_id: str = "YOLOv11-XLarge",
|
37 |
+
conf: float = 0.25,
|
38 |
+
iou: float = 0.7
|
39 |
+
):
|
40 |
+
# Initialize model at startup
|
41 |
+
model = init_model(model_id)
|
42 |
+
|
43 |
+
# Download and open image from URL
|
44 |
+
image = Image.open(image.file)
|
45 |
+
|
46 |
+
# Run inference with the PIL Image
|
47 |
+
results = model.predict(source=image, conf=conf, iou=iou)
|
48 |
+
|
49 |
+
# Extract detection results
|
50 |
+
result = results[0]
|
51 |
+
# print(result)
|
52 |
+
detections = []
|
53 |
+
|
54 |
+
for box in result.boxes:
|
55 |
+
detection = {
|
56 |
+
"class": result.names[int(box.cls[0])],
|
57 |
+
"confidence": float(box.conf[0]),
|
58 |
+
"bbox": box.xyxy[0].tolist() # Convert bbox tensor to list
|
59 |
+
}
|
60 |
+
detections.append(detection)
|
61 |
+
print(detections)
|
62 |
+
|
63 |
+
return {"detections": detections}
|
64 |
+
|
65 |
+
if __name__ == "__main__":
|
66 |
+
import uvicorn
|
67 |
+
uvicorn.run(app, host="0.0.0.0", port=7860)
|
gradio-app.py
ADDED
@@ -0,0 +1,145 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from typing import Tuple, Dict
|
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 |
+
# Define models
|
10 |
+
MODEL_OPTIONS = {
|
11 |
+
"YOLOv11-Nano": "medieval-yolov11n.pt",
|
12 |
+
"YOLOv11-Small": "medieval-yolov11s.pt",
|
13 |
+
"YOLOv11-Medium": "medieval-yolov11m.pt",
|
14 |
+
"YOLOv11-Large": "medieval-yolov11l.pt",
|
15 |
+
"YOLOv11-XLarge": "medieval-yolov11x.pt"
|
16 |
+
}
|
17 |
+
|
18 |
+
# Dictionary to store loaded models
|
19 |
+
models: Dict[str, YOLO] = {}
|
20 |
+
|
21 |
+
# Load all models
|
22 |
+
for name, model_file in MODEL_OPTIONS.items():
|
23 |
+
model_path = hf_hub_download(
|
24 |
+
repo_id="biglam/medieval-manuscript-yolov11",
|
25 |
+
filename=model_file
|
26 |
+
)
|
27 |
+
models[name] = YOLO(model_path)
|
28 |
+
|
29 |
+
# Create annotators
|
30 |
+
LABEL_ANNOTATOR = sv.LabelAnnotator(text_color=sv.Color.BLACK)
|
31 |
+
BOX_ANNOTATOR = sv.BoxAnnotator()
|
32 |
+
|
33 |
+
def detect_and_annotate(
|
34 |
+
image: np.ndarray,
|
35 |
+
model_name: str,
|
36 |
+
conf_threshold: float,
|
37 |
+
iou_threshold: float
|
38 |
+
) -> np.ndarray:
|
39 |
+
# Get the selected model
|
40 |
+
model = models[model_name]
|
41 |
+
|
42 |
+
# Perform inference
|
43 |
+
results = model.predict(
|
44 |
+
image,
|
45 |
+
conf=conf_threshold,
|
46 |
+
iou=iou_threshold
|
47 |
+
)[0]
|
48 |
+
|
49 |
+
# Convert results to supervision Detections
|
50 |
+
boxes = results.boxes.xyxy.cpu().numpy()
|
51 |
+
confidence = results.boxes.conf.cpu().numpy()
|
52 |
+
class_ids = results.boxes.cls.cpu().numpy().astype(int)
|
53 |
+
|
54 |
+
# Create Detections object
|
55 |
+
detections = sv.Detections(
|
56 |
+
xyxy=boxes,
|
57 |
+
confidence=confidence,
|
58 |
+
class_id=class_ids
|
59 |
+
)
|
60 |
+
|
61 |
+
# Create labels with confidence scores
|
62 |
+
labels = [
|
63 |
+
f"{results.names[class_id]} ({conf:.2f})"
|
64 |
+
for class_id, conf
|
65 |
+
in zip(class_ids, confidence)
|
66 |
+
]
|
67 |
+
|
68 |
+
# Annotate image
|
69 |
+
annotated_image = image.copy()
|
70 |
+
annotated_image = BOX_ANNOTATOR.annotate(scene=annotated_image, detections=detections)
|
71 |
+
annotated_image = LABEL_ANNOTATOR.annotate(scene=annotated_image, detections=detections, labels=labels)
|
72 |
+
|
73 |
+
return annotated_image
|
74 |
+
|
75 |
+
# Create Gradio interface
|
76 |
+
with gr.Blocks() as demo:
|
77 |
+
gr.Markdown("# Medieval Manuscript Detection with YOLO")
|
78 |
+
|
79 |
+
with gr.Row():
|
80 |
+
with gr.Column():
|
81 |
+
input_image = gr.Image(
|
82 |
+
label="Input Image",
|
83 |
+
type='numpy'
|
84 |
+
)
|
85 |
+
with gr.Accordion("Detection Settings", open=True):
|
86 |
+
model_selector = gr.Dropdown(
|
87 |
+
choices=list(MODEL_OPTIONS.keys()),
|
88 |
+
value=list(MODEL_OPTIONS.keys())[0],
|
89 |
+
label="Model",
|
90 |
+
info="Select YOLO model variant"
|
91 |
+
)
|
92 |
+
with gr.Row():
|
93 |
+
conf_threshold = gr.Slider(
|
94 |
+
label="Confidence Threshold",
|
95 |
+
minimum=0.0,
|
96 |
+
maximum=1.0,
|
97 |
+
step=0.05,
|
98 |
+
value=0.25,
|
99 |
+
)
|
100 |
+
iou_threshold = gr.Slider(
|
101 |
+
label="IoU Threshold",
|
102 |
+
minimum=0.0,
|
103 |
+
maximum=1.0,
|
104 |
+
step=0.05,
|
105 |
+
value=0.45,
|
106 |
+
info="Decrease for stricter detection, increase for more overlapping boxes"
|
107 |
+
)
|
108 |
+
with gr.Row():
|
109 |
+
clear_btn = gr.Button("Clear")
|
110 |
+
detect_btn = gr.Button("Detect", variant="primary")
|
111 |
+
|
112 |
+
with gr.Column():
|
113 |
+
output_image = gr.Image(
|
114 |
+
label="Detection Result",
|
115 |
+
type='numpy'
|
116 |
+
)
|
117 |
+
|
118 |
+
def process_image(
|
119 |
+
image: np.ndarray,
|
120 |
+
model_name: str,
|
121 |
+
conf_threshold: float,
|
122 |
+
iou_threshold: float
|
123 |
+
) -> Tuple[np.ndarray, np.ndarray]:
|
124 |
+
if image is None:
|
125 |
+
return None, None
|
126 |
+
annotated_image = detect_and_annotate(image, model_name, conf_threshold, iou_threshold)
|
127 |
+
return image, annotated_image
|
128 |
+
|
129 |
+
def clear():
|
130 |
+
return None, None
|
131 |
+
|
132 |
+
# Connect buttons to functions
|
133 |
+
detect_btn.click(
|
134 |
+
process_image,
|
135 |
+
inputs=[input_image, model_selector, conf_threshold, iou_threshold],
|
136 |
+
outputs=[input_image, output_image]
|
137 |
+
)
|
138 |
+
clear_btn.click(
|
139 |
+
clear,
|
140 |
+
inputs=None,
|
141 |
+
outputs=[input_image, output_image]
|
142 |
+
)
|
143 |
+
|
144 |
+
if __name__ == "__main__":
|
145 |
+
demo.launch(debug=True, show_error=True)
|
requirements.txt
ADDED
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
ultralytics
|
2 |
+
huggingface-hub
|
3 |
+
supervision
|
4 |
+
requests
|
5 |
+
fastapi
|
test.py
ADDED
@@ -0,0 +1,24 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import requests
|
2 |
+
|
3 |
+
# Path to your image
|
4 |
+
image_path = "/Users/wjm55/yale/weaviate-test/yolov11_output/valid/images/Paris_BnF_Velins_611_00003.jpg"
|
5 |
+
|
6 |
+
# Open the image file
|
7 |
+
with open(image_path, 'rb') as f:
|
8 |
+
# Create the files parameter for the POST request
|
9 |
+
files = {'image': f}
|
10 |
+
|
11 |
+
# Optional parameters (using defaults from the API)
|
12 |
+
params = {
|
13 |
+
'model_id': 'YOLOv11-XLarge', # default model
|
14 |
+
'conf': 0.25, # confidence threshold
|
15 |
+
'iou': 0.7 # IoU threshold
|
16 |
+
}
|
17 |
+
|
18 |
+
# Send POST request to the endpoint
|
19 |
+
response = requests.post('http://localhost:7860/predict',
|
20 |
+
files=files,
|
21 |
+
params=params)
|
22 |
+
|
23 |
+
# Print the results
|
24 |
+
print(response.json())
|