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!
Updated about 1 year ago