Using Your Own Custom Models

In the current release of Portal, TensorFlow and Darknet models are supported. In the event where you have another model that is built from a different machine learning library, you may consider creating a custom model class.
Datature's Engine is compatible with any type of model provided that they inherit from the BaseModel class architecture.

Creating a Custom Model

The steps to creating a custom model can be done in several steps:
From src/engine/server/models, create your own custom model module (e.g. example_model.py)
Import the following modules:

from server.services.errors import Errors, PortalError
    from server.services.hashing import get_hash
    from server.models.abstract.BaseModel import BaseModel

Create your own custom model model class (such as ExampleModel), which inherits from BaseModel.

class ExampleModel(BaseModel):{

}

Within your custom model class, define the following functions:

  • load_label_map(self)
    • Converts your label map into the following dictionary format and then saves it into self.label_map:
self._label_map_ = {
    '1':{
            'id': 1,
            'name': 'apple',
        },
    '2':{
            'id': 2,
            'name': 'pear',
        }
}
  • register(self)
    • Checks if all critical files needed for loading and prediction are inside self.directory
    • Set the height (self.height) and width (self.width) of the image that the model should receive.
    • Load the label map with the function load_label_map()
    • Set the model key (self.key) to be the hash of the model directory (self.directory) with the function:
from server.services.hashing import get_hash
self._key_ = get_hash(self._directory_)
* return (self._key_, self) as a tuple.
  • load(self)
    • Load the model into a variable and save that variable into self.model.
loaded_model = load_the_model(<model_path>)
self._model_ = loaded_model
  • predict(self, image_array)
    • Perform inference on the image array.
    • Return the inference as a dictionary of
{
    "detection_masks":  <ndarray of shape [Instances, Height, Width]
                        representing the prediction masks,
                        or None if this is not a segmentation model>,

    "detection_boxes":  <ndarray of shape [Instances, 4]
                        representing the bounding boxes,
                        in the form (Ymin, Xmin, Ymax, Xmax)>,

    "detection_scores": <ndarray of shape [Instances, 1]
                        representing the confidence>,

    "detection_classes": <ndarray of shape [Instances, 1],
                        representing label ids>,
}

🚧

The segmentation masks are in the form of image masks, not bounding box masks.

📘

You may also define other functions, but these functions are the basic necessity.
For models that are hosted on another platform/deployed on the cloud, you may also look at Datature's API Model Script as a template for your development.

In the Model Factory src/engine/server/models/abstract/Model.py, import and add your your custom model class into the model_class dictionary.

from server.models.example_model import ExampleModel

# Inside Model function:
model_class = {
    "tensorflow": TensorflowModel,
    "darknet": DarknetModel,
    "example": ExampleModel, # <<--------- add here
}

The Engine is now configured to accept a new type of model. Next, we configure the App.
In the model file src/app/src/components/annotations/model.tsx:

  • In FormData -> modelType add your custom model string
export type FormData = {
  type: string;
  name: string;
  description: string;
  directory: string;
  modelKey: string;
  projectSecret: string;
  modelType: "tensorflow" | "darknet" | "example" | ""; //<----Add here
};
  • In Model -> render() -> modelTypes, add your custom model
const modelTypes = {
  tensorflow: "TensorFlow 2.0",
  darknet: "DarkNet (YOLO v3, YOLO v4)",
  example: "Example Model", //<-------------Add here
};
  • In Model -> render() -> registerModelForm, add a new Menu.Item in the Menu component
<Menu>
  <Menu.Item
      shouldDismissPopover={false}
      text={modelTypes.tensorflow}
      onClick={() => {
      const event = {
          target: { name: "modelType", value: "tensorflow" },
      };
      this.handleChangeForm(event);
      }}
  />
  <Menu.Item
      shouldDismissPopover={false}
      text={modelTypes.darknet}
      onClick={() => {
      const event = {
          target: { name: "modelType", value: "darknet" },
      };
      this.handleChangeForm(event);
      }}
  />
  </Menu>
  //----------------Add below----------------------//
  <Menu.Item
    shouldDismissPopover={false}
    text={modelTypes.example}
    onClick={() => {
      const event = { target: { name: "modelType", value: "example" } };
      this.handleChangeForm(event);
    }}
  />
</Menu>

Now restart Portal and you should be able to see the changes!