wjm55 commited on
Commit
132787d
·
1 Parent(s): 5ab39af
Files changed (5) hide show
  1. Dockerfile +17 -0
  2. app.py +67 -0
  3. gradio-app.py +145 -0
  4. requirements.txt +5 -0
  5. 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())