Yolov5 support for Rikai

rikai-yolov5 integrates Yolov5 implemented in PyTorch with Rikai. It is based on the packaged ultralytics/yolov5.

Notebooks

• Demo 1: Using Rikai to analyze an image from Jay Chou's Mojito.

Usage

There are two ways to use rikai-yolov5.

Set customized_flavor to yolov5 when logging the model, rikai will use rikai.contrib.yolov5.codegen.generate_udf instead of rikai.spark.sql.codegen.pytorch.generate_udf.

rikai.mlflow.pytorch.log_model(
    model,
    "model",
    OUTPUT_SCHEMA,
    pre_processing=pre,
    post_processing=post,
    registered_model_name=registered_model_name,
    customized_flavor="yolov5",
)

Another way is setting the flavor in Rikai SQL:

CREATE MODEL mlflow_yolov5_m
FLAVOR yolov5
OPTIONS (
  device='cpu'
)
USING 'mlflow:///{registered_model_name}';

Available Options

Here is a sample usage of the above options:

CREATE MODEL mlflow_yolov5_m
OPTIONS (
  device='cpu',
  iou_thres=0.5
)
USING 'mlflow:///{registered_model_name}';