[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

Author: seanshi-scale

.pylintrc

@@ -24,3 +24,6 @@ reports=no
 
 [tool.pylint.FORMAT]
 max-line-length=79
+
+[tool.pylint.MASTER]
+extension-pkg-whitelist=pydantic

nucleus/test_launch_integration.py

@@ -0,0 +1,311 @@
+import io
+from typing import Any, Callable, Dict, List, Optional, Type
+
+from PIL import Image, ImageDraw
+from pydantic import BaseModel, Extra, ValidationError
+
+# From scaleapi/server/src/lib/select/api/types.ts
+# These classes specify how user models must pass output to Launch + Nucleus.
+
+
+class PointModel(BaseModel, extra=Extra.forbid):
+    x: float
+    y: float
+
+
+class BoxGeometryModel(BaseModel, extra=Extra.forbid):
+    x: float
+    y: float
+    width: float
+    height: float
+
+
+class BoxAnnotationModel(BaseModel, extra=Extra.forbid):
+    geometry: BoxGeometryModel
+    type: str
+    label: Optional[str] = None
+    confidence: Optional[float] = None
+    classPdf: Optional[Dict[str, float]] = None
+    metadata: Optional[Dict[str, Any]] = None
+
+
+class NoneGeometryModel(BaseModel, extra=Extra.forbid):
+    pass
+
+
+class CategoryAnnotationModel(BaseModel, extra=Extra.forbid):
+    geometry: NoneGeometryModel
+    type: str
+    label: Optional[str] = None
+    confidence: Optional[float] = None
+    classPdf: Optional[Dict[str, float]] = None
+    metadata: Optional[Dict[str, Any]] = None
+
+
+class LineGeometryModel(BaseModel, extra=Extra.forbid):
+    vertices: List[PointModel]
+
+
+class LineAnnotationModel(BaseModel, extra=Extra.forbid):
+    geometry: LineGeometryModel
+    type: str
+    label: Optional[str] = None
+    confidence: Optional[float] = None
+    classPdf: Optional[Dict[str, float]] = None
+    metadata: Optional[Dict[str, Any]] = None
+
+
+class PolygonGeometryModel(BaseModel, extra=Extra.forbid):
+    vertices: List[PointModel]
+
+
+class PolygonAnnotationModel(BaseModel, extra=Extra.forbid):
+    geometry: PolygonGeometryModel
+    type: str
+    label: Optional[str] = None
+    confidence: Optional[float] = None
+    classPdf: Optional[Dict[str, float]] = None
+    metadata: Optional[Dict[str, Any]] = None
+
+
+def verify_output(
+    annotation_list: List[Dict[str, Any]],
+    model: Type[BaseModel],
+    annotation_type: str,
+):
+    for annotation in annotation_list:
+        try:
+            model.parse_obj(annotation)
+        except ValidationError as e:
+            raise ValueError("Failed validation") from e
+        if annotation["type"] != annotation_type:
+            raise ValueError(
+                f"Bounding box type {annotation['type']} should equal {annotation_type}"
+            )
+
+
+def verify_box_output(bbox_list):
+    annotation_type = "box"
+    return verify_output(
+        bbox_list,
+        BoxAnnotationModel,
+        annotation_type,
+    )
+
+
+def verify_category_output(category_list):
+    """I think the annotation needs to be a list with a single element in the Launch+Nucleus sfn."""
+    annotation_type = "category"
+    return verify_output(
+        category_list, CategoryAnnotationModel, annotation_type
+    )
+
+
+def verify_line_output(line_list):
+    annotation_type = "line"
+    return verify_output(
+        line_list,
+        LineAnnotationModel,
+        annotation_type,
+    )
+
+
+def verify_polygon_output(polygon_list):
+    annotation_type = "polygon"
+    return verify_output(
+        polygon_list,
+        PolygonAnnotationModel,
+        annotation_type,
+    )
+
+
+def _run_model(
+    input_bytes: bytes,
+    load_predict_fn: Callable,
+    load_model_fn: Optional[Callable],
+    model: Optional[Any],
+):
+    if not (model is None) ^ (load_model_fn is None):
+        raise ValueError(
+            "Exactly one of `model` and `load_model_fn` must not be None."
+        )
+
+    if load_model_fn:
+        model = load_model_fn()
+
+    predict_fn = load_predict_fn(model)
+    return predict_fn(input_bytes)
+
+
+_FILL_COLOR = (0, 255, 0, 50)
+_OUTLINE_COLOR = (0, 255, 0, 255)
+
+
+def visualize_box_launch_bundle(
+    img_file: str,
+    load_predict_fn: Callable,
+    load_model_fn: Callable = None,
+    model: Any = None,
+    show_image: bool = False,
+    max_annotations: int = 5,
+) -> Image:
+    """
+    Run this function locally to visualize what your Launch bundle will do on a local image
+    Intended to verify that your Launch bundle returns annotations in the correct format, as well as sanity check
+    any coordinate systems used for the image.
+    Will display the image in a separate window if show_image == True.
+    Returns the image as well.
+
+    Parameters:
+        img_file: The path to a local image file.
+        load_predict_fn: The load_predict_fn as part of your Launch bundle
+        load_model_fn: The load_model_fn as part of your Launch bundle
+        model: The model as part of your Launch bundle. Note: exactly one of load_model_fn and model must be specified
+        show_image: Whether to automatically pop up the image + predictions in a separate window. Can be useful in a
+          script.
+        max_annotations: How many annotations you want to draw
+
+    Returns:
+        Image: The image with annotations drawn on top.
+    """
+    # Basically do the same thing as what Launch does but locally
+
+    with open(img_file, "rb") as f:
+        img_bytes = f.read()
+
+    output = _run_model(img_bytes, load_predict_fn, load_model_fn, model)
+    verify_box_output(output)
+
+    image = Image.open(io.BytesIO(img_bytes))
+    draw = ImageDraw.Draw(image, "RGBA")
+    for bbox in output[:max_annotations]:
+        geo = bbox["geometry"]
+        x, y, w, h = geo["x"], geo["y"], geo["width"], geo["height"]
+        draw.rectangle(
+            [(x, y), (x + w, y + h)], outline=_OUTLINE_COLOR, fill=_FILL_COLOR
+        )
+
+    if show_image:
+        image.show()
+
+    return image
+
+
+def run_category_launch_bundle(
+    img_file: str,
+    load_predict_fn: Callable,
+    load_model_fn: Callable = None,
+    model: Any = None,
+):
+    """
+    Run this function locally to test if your image categorization model returns a format consumable by Launch + Nucleus
+    Parameters:
+        img_file: The path to a local image file.
+        load_predict_fn: The load_predict_fn as part of your Launch bundle
+        load_model_fn: The load_model_fn as part of your Launch bundle
+        model: The model as part of your Launch bundle. Note: exactly one of load_model_fn and model must be specified
+    Returns:
+        The raw output (as a json) of your categorization model.
+    """
+    with open(img_file, "rb") as f:
+        img_bytes = f.read()
+
+    output = _run_model(img_bytes, load_predict_fn, load_model_fn, model)
+    verify_category_output(output)
+    return output
+
+
+def visualize_line_launch_bundle(
+    img_file: str,
+    load_predict_fn: Callable,
+    load_model_fn: Callable = None,
+    model: Any = None,
+    show_image: bool = False,
+    max_annotations: int = 5,
+) -> Image:
+    """
+    Run this function locally to visualize what your Launch bundle will do on a local image
+    Intended to verify that your Launch bundle returns annotations in the correct format, as well as sanity check
+    any coordinate systems used for the image.
+    Will display the image in a separate window if show_image == True.
+    Returns the image as well.
+
+    Parameters:
+        img_file: The path to a local image file.
+        load_predict_fn: The load_predict_fn as part of your Launch bundle
+        load_model_fn: The load_model_fn as part of your Launch bundle
+        model: The model as part of your Launch bundle. Note: exactly one of load_model_fn and model must be specified
+        show_image: Whether to automatically pop up the image + predictions in a separate window. Can be useful in a
+          script.
+        max_annotations: How many annotations you want to draw
+
+    Returns:
+        Image: The image with annotations drawn on top.
+    """
+    # Basically do the same thing as what Launch does but locally
+
+    with open(img_file, "rb") as f:
+        img_bytes = f.read()
+
+    output = _run_model(img_bytes, load_predict_fn, load_model_fn, model)
+    verify_line_output(output)
+
+    image = Image.open(io.BytesIO(img_bytes))
+    draw = ImageDraw.Draw(image, "RGBA")
+    for bbox in output[:max_annotations]:
+        geo = bbox["geometry"]
+        vertices = [(v["x"], v["y"]) for v in geo["vertices"]]
+        draw.line(vertices, fill=_OUTLINE_COLOR)
+
+    if show_image:
+        image.show()
+
+    return image
+
+
+def visualize_polygon_launch_bundle(
+    img_file: str,
+    load_predict_fn: Callable,
+    load_model_fn: Callable = None,
+    model: Any = None,
+    show_image: bool = False,
+    max_annotations: int = 5,
+) -> Image:
+    """
+    Run this function locally to visualize what your Launch bundle will do on a local image
+    Intended to verify that your Launch bundle returns annotations in the correct format, as well as sanity check
+    any coordinate systems used for the image.
+    Will display the image in a separate window if show_image == True.
+    Returns the image as well.
+
+    Parameters:
+        img_file: The path to a local image file.
+        load_predict_fn: The load_predict_fn as part of your Launch bundle
+        load_model_fn: The load_model_fn as part of your Launch bundle
+        model: The model as part of your Launch bundle. Note: exactly one of load_model_fn and model must be specified
+        show_image: Whether to automatically pop up the image + predictions in a separate window. Can be useful in a
+          script.
+        max_annotations: How many annotations you want to draw
+
+    Returns:
+        Image: The image with annotations drawn on top.
+    """
+    # Basically do the same thing as what Launch does but locally
+
+    with open(img_file, "rb") as f:
+        img_bytes = f.read()
+
+    output = _run_model(img_bytes, load_predict_fn, load_model_fn, model)
+    verify_polygon_output(output)
+
+    image = Image.open(io.BytesIO(img_bytes))
+    draw = ImageDraw.Draw(image, "RGBA")
+    for bbox in output[:max_annotations]:
+        geo = bbox["geometry"]
+        vertices = [(v["x"], v["y"]) for v in geo["vertices"]]
+        draw.polygon(vertices, outline=_OUTLINE_COLOR, fill=_FILL_COLOR)
+
+    if show_image:
+        image.show()
+
+    return image