Skip to content

Commit 185a060

Browse files
[Launch] Add a utility for users to test their Launch+Nucleus bundles locally (#362)
* add function to verify * add tests * oops * rename * comment * return the image so that it can display in a jupyter notebook * cloudpickle doesn't allow config * refactor some names * refactor out into helper fn * just use pydantic * verify category output * add untested fns for checking/visualizing category + line * polygon heh * test for categories * tests++++++ and make the fill look nicer :D * clean up the comments
1 parent 6108121 commit 185a060

File tree

3 files changed

+728
-0
lines changed

3 files changed

+728
-0
lines changed

.pylintrc

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,3 +24,6 @@ reports=no
2424

2525
[tool.pylint.FORMAT]
2626
max-line-length=79
27+
28+
[tool.pylint.MASTER]
29+
extension-pkg-whitelist=pydantic

nucleus/test_launch_integration.py

Lines changed: 311 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,311 @@
1+
import io
2+
from typing import Any, Callable, Dict, List, Optional, Type
3+
4+
from PIL import Image, ImageDraw
5+
from pydantic import BaseModel, Extra, ValidationError
6+
7+
# From scaleapi/server/src/lib/select/api/types.ts
8+
# These classes specify how user models must pass output to Launch + Nucleus.
9+
10+
11+
class PointModel(BaseModel, extra=Extra.forbid):
12+
x: float
13+
y: float
14+
15+
16+
class BoxGeometryModel(BaseModel, extra=Extra.forbid):
17+
x: float
18+
y: float
19+
width: float
20+
height: float
21+
22+
23+
class BoxAnnotationModel(BaseModel, extra=Extra.forbid):
24+
geometry: BoxGeometryModel
25+
type: str
26+
label: Optional[str] = None
27+
confidence: Optional[float] = None
28+
classPdf: Optional[Dict[str, float]] = None
29+
metadata: Optional[Dict[str, Any]] = None
30+
31+
32+
class NoneGeometryModel(BaseModel, extra=Extra.forbid):
33+
pass
34+
35+
36+
class CategoryAnnotationModel(BaseModel, extra=Extra.forbid):
37+
geometry: NoneGeometryModel
38+
type: str
39+
label: Optional[str] = None
40+
confidence: Optional[float] = None
41+
classPdf: Optional[Dict[str, float]] = None
42+
metadata: Optional[Dict[str, Any]] = None
43+
44+
45+
class LineGeometryModel(BaseModel, extra=Extra.forbid):
46+
vertices: List[PointModel]
47+
48+
49+
class LineAnnotationModel(BaseModel, extra=Extra.forbid):
50+
geometry: LineGeometryModel
51+
type: str
52+
label: Optional[str] = None
53+
confidence: Optional[float] = None
54+
classPdf: Optional[Dict[str, float]] = None
55+
metadata: Optional[Dict[str, Any]] = None
56+
57+
58+
class PolygonGeometryModel(BaseModel, extra=Extra.forbid):
59+
vertices: List[PointModel]
60+
61+
62+
class PolygonAnnotationModel(BaseModel, extra=Extra.forbid):
63+
geometry: PolygonGeometryModel
64+
type: str
65+
label: Optional[str] = None
66+
confidence: Optional[float] = None
67+
classPdf: Optional[Dict[str, float]] = None
68+
metadata: Optional[Dict[str, Any]] = None
69+
70+
71+
def verify_output(
72+
annotation_list: List[Dict[str, Any]],
73+
model: Type[BaseModel],
74+
annotation_type: str,
75+
):
76+
for annotation in annotation_list:
77+
try:
78+
model.parse_obj(annotation)
79+
except ValidationError as e:
80+
raise ValueError("Failed validation") from e
81+
if annotation["type"] != annotation_type:
82+
raise ValueError(
83+
f"Bounding box type {annotation['type']} should equal {annotation_type}"
84+
)
85+
86+
87+
def verify_box_output(bbox_list):
88+
annotation_type = "box"
89+
return verify_output(
90+
bbox_list,
91+
BoxAnnotationModel,
92+
annotation_type,
93+
)
94+
95+
96+
def verify_category_output(category_list):
97+
"""I think the annotation needs to be a list with a single element in the Launch+Nucleus sfn."""
98+
annotation_type = "category"
99+
return verify_output(
100+
category_list, CategoryAnnotationModel, annotation_type
101+
)
102+
103+
104+
def verify_line_output(line_list):
105+
annotation_type = "line"
106+
return verify_output(
107+
line_list,
108+
LineAnnotationModel,
109+
annotation_type,
110+
)
111+
112+
113+
def verify_polygon_output(polygon_list):
114+
annotation_type = "polygon"
115+
return verify_output(
116+
polygon_list,
117+
PolygonAnnotationModel,
118+
annotation_type,
119+
)
120+
121+
122+
def _run_model(
123+
input_bytes: bytes,
124+
load_predict_fn: Callable,
125+
load_model_fn: Optional[Callable],
126+
model: Optional[Any],
127+
):
128+
if not (model is None) ^ (load_model_fn is None):
129+
raise ValueError(
130+
"Exactly one of `model` and `load_model_fn` must not be None."
131+
)
132+
133+
if load_model_fn:
134+
model = load_model_fn()
135+
136+
predict_fn = load_predict_fn(model)
137+
return predict_fn(input_bytes)
138+
139+
140+
_FILL_COLOR = (0, 255, 0, 50)
141+
_OUTLINE_COLOR = (0, 255, 0, 255)
142+
143+
144+
def visualize_box_launch_bundle(
145+
img_file: str,
146+
load_predict_fn: Callable,
147+
load_model_fn: Callable = None,
148+
model: Any = None,
149+
show_image: bool = False,
150+
max_annotations: int = 5,
151+
) -> Image:
152+
"""
153+
Run this function locally to visualize what your Launch bundle will do on a local image
154+
Intended to verify that your Launch bundle returns annotations in the correct format, as well as sanity check
155+
any coordinate systems used for the image.
156+
Will display the image in a separate window if show_image == True.
157+
Returns the image as well.
158+
159+
Parameters:
160+
img_file: The path to a local image file.
161+
load_predict_fn: The load_predict_fn as part of your Launch bundle
162+
load_model_fn: The load_model_fn as part of your Launch bundle
163+
model: The model as part of your Launch bundle. Note: exactly one of load_model_fn and model must be specified
164+
show_image: Whether to automatically pop up the image + predictions in a separate window. Can be useful in a
165+
script.
166+
max_annotations: How many annotations you want to draw
167+
168+
Returns:
169+
Image: The image with annotations drawn on top.
170+
"""
171+
# Basically do the same thing as what Launch does but locally
172+
173+
with open(img_file, "rb") as f:
174+
img_bytes = f.read()
175+
176+
output = _run_model(img_bytes, load_predict_fn, load_model_fn, model)
177+
verify_box_output(output)
178+
179+
image = Image.open(io.BytesIO(img_bytes))
180+
draw = ImageDraw.Draw(image, "RGBA")
181+
for bbox in output[:max_annotations]:
182+
geo = bbox["geometry"]
183+
x, y, w, h = geo["x"], geo["y"], geo["width"], geo["height"]
184+
draw.rectangle(
185+
[(x, y), (x + w, y + h)], outline=_OUTLINE_COLOR, fill=_FILL_COLOR
186+
)
187+
188+
if show_image:
189+
image.show()
190+
191+
return image
192+
193+
194+
def run_category_launch_bundle(
195+
img_file: str,
196+
load_predict_fn: Callable,
197+
load_model_fn: Callable = None,
198+
model: Any = None,
199+
):
200+
"""
201+
Run this function locally to test if your image categorization model returns a format consumable by Launch + Nucleus
202+
Parameters:
203+
img_file: The path to a local image file.
204+
load_predict_fn: The load_predict_fn as part of your Launch bundle
205+
load_model_fn: The load_model_fn as part of your Launch bundle
206+
model: The model as part of your Launch bundle. Note: exactly one of load_model_fn and model must be specified
207+
Returns:
208+
The raw output (as a json) of your categorization model.
209+
"""
210+
with open(img_file, "rb") as f:
211+
img_bytes = f.read()
212+
213+
output = _run_model(img_bytes, load_predict_fn, load_model_fn, model)
214+
verify_category_output(output)
215+
return output
216+
217+
218+
def visualize_line_launch_bundle(
219+
img_file: str,
220+
load_predict_fn: Callable,
221+
load_model_fn: Callable = None,
222+
model: Any = None,
223+
show_image: bool = False,
224+
max_annotations: int = 5,
225+
) -> Image:
226+
"""
227+
Run this function locally to visualize what your Launch bundle will do on a local image
228+
Intended to verify that your Launch bundle returns annotations in the correct format, as well as sanity check
229+
any coordinate systems used for the image.
230+
Will display the image in a separate window if show_image == True.
231+
Returns the image as well.
232+
233+
Parameters:
234+
img_file: The path to a local image file.
235+
load_predict_fn: The load_predict_fn as part of your Launch bundle
236+
load_model_fn: The load_model_fn as part of your Launch bundle
237+
model: The model as part of your Launch bundle. Note: exactly one of load_model_fn and model must be specified
238+
show_image: Whether to automatically pop up the image + predictions in a separate window. Can be useful in a
239+
script.
240+
max_annotations: How many annotations you want to draw
241+
242+
Returns:
243+
Image: The image with annotations drawn on top.
244+
"""
245+
# Basically do the same thing as what Launch does but locally
246+
247+
with open(img_file, "rb") as f:
248+
img_bytes = f.read()
249+
250+
output = _run_model(img_bytes, load_predict_fn, load_model_fn, model)
251+
verify_line_output(output)
252+
253+
image = Image.open(io.BytesIO(img_bytes))
254+
draw = ImageDraw.Draw(image, "RGBA")
255+
for bbox in output[:max_annotations]:
256+
geo = bbox["geometry"]
257+
vertices = [(v["x"], v["y"]) for v in geo["vertices"]]
258+
draw.line(vertices, fill=_OUTLINE_COLOR)
259+
260+
if show_image:
261+
image.show()
262+
263+
return image
264+
265+
266+
def visualize_polygon_launch_bundle(
267+
img_file: str,
268+
load_predict_fn: Callable,
269+
load_model_fn: Callable = None,
270+
model: Any = None,
271+
show_image: bool = False,
272+
max_annotations: int = 5,
273+
) -> Image:
274+
"""
275+
Run this function locally to visualize what your Launch bundle will do on a local image
276+
Intended to verify that your Launch bundle returns annotations in the correct format, as well as sanity check
277+
any coordinate systems used for the image.
278+
Will display the image in a separate window if show_image == True.
279+
Returns the image as well.
280+
281+
Parameters:
282+
img_file: The path to a local image file.
283+
load_predict_fn: The load_predict_fn as part of your Launch bundle
284+
load_model_fn: The load_model_fn as part of your Launch bundle
285+
model: The model as part of your Launch bundle. Note: exactly one of load_model_fn and model must be specified
286+
show_image: Whether to automatically pop up the image + predictions in a separate window. Can be useful in a
287+
script.
288+
max_annotations: How many annotations you want to draw
289+
290+
Returns:
291+
Image: The image with annotations drawn on top.
292+
"""
293+
# Basically do the same thing as what Launch does but locally
294+
295+
with open(img_file, "rb") as f:
296+
img_bytes = f.read()
297+
298+
output = _run_model(img_bytes, load_predict_fn, load_model_fn, model)
299+
verify_polygon_output(output)
300+
301+
image = Image.open(io.BytesIO(img_bytes))
302+
draw = ImageDraw.Draw(image, "RGBA")
303+
for bbox in output[:max_annotations]:
304+
geo = bbox["geometry"]
305+
vertices = [(v["x"], v["y"]) for v in geo["vertices"]]
306+
draw.polygon(vertices, outline=_OUTLINE_COLOR, fill=_FILL_COLOR)
307+
308+
if show_image:
309+
image.show()
310+
311+
return image

0 commit comments

Comments
 (0)