122 lines
3.6 KiB
Python
122 lines
3.6 KiB
Python
import os
|
|
from copy import deepcopy
|
|
import json
|
|
import re
|
|
from _ctypes import PyObj_FromPtr
|
|
|
|
import mmcv
|
|
import mmdet.core
|
|
|
|
|
|
class NoIndent(object):
|
|
""" Value wrapper. """
|
|
def __init__(self, value):
|
|
self.value = value
|
|
|
|
|
|
class MyEncoder(json.JSONEncoder):
|
|
FORMAT_SPEC = '@@{}@@'
|
|
regex = re.compile(FORMAT_SPEC.format(r'(\d+)'))
|
|
|
|
def __init__(self, **kwargs):
|
|
# Save copy of any keyword argument values needed for use here.
|
|
self.__sort_keys = kwargs.get('sort_keys', None)
|
|
super(MyEncoder, self).__init__(**kwargs)
|
|
|
|
def default(self, obj):
|
|
return (self.FORMAT_SPEC.format(id(obj)) if isinstance(obj, NoIndent)
|
|
else super(MyEncoder, self).default(obj))
|
|
|
|
def encode(self, obj):
|
|
format_spec = self.FORMAT_SPEC # Local var to expedite access.
|
|
json_repr = super(MyEncoder, self).encode(obj) # Default JSON.
|
|
|
|
# Replace any marked-up object ids in the JSON repr with the
|
|
# value returned from the json.dumps() of the corresponding
|
|
# wrapped Python object.
|
|
for match in self.regex.finditer(json_repr):
|
|
# see https://stackoverflow.com/a/15012814/355230
|
|
id = int(match.group(1))
|
|
no_indent = PyObj_FromPtr(id)
|
|
json_obj_repr = json.dumps(
|
|
no_indent.value, sort_keys=self.__sort_keys
|
|
)
|
|
|
|
# Replace the matched id string with json formatted representation
|
|
# of the corresponding Python object.
|
|
json_repr = json_repr.replace(
|
|
'"{}"'.format(format_spec.format(id)), json_obj_repr)
|
|
|
|
return json_repr
|
|
|
|
|
|
DAG_template = {
|
|
"comment": "DAG defined, use DAG generate data flow",
|
|
"green_mode": True,
|
|
"model_list": [
|
|
{
|
|
"model_runner": "",
|
|
"model_id": 244,
|
|
"model_init_params_file": ""
|
|
}
|
|
],
|
|
"DAG": {
|
|
"DAG0": [
|
|
NoIndent([["start_0"], ["0_0"]]),
|
|
NoIndent([["start_0", "0_0"], [None]])
|
|
]
|
|
},
|
|
"report_params": {
|
|
"0_0": {
|
|
"dbkey": "bbox",
|
|
"conf_thres": 0.3
|
|
}
|
|
}
|
|
}
|
|
|
|
|
|
init_params_template = {
|
|
"model_path": "",
|
|
"num_classes": 80,
|
|
"remapping_type": "COCO",
|
|
"input_shape": [640, 640],
|
|
"conf_threshold": 0.3,
|
|
"iou_threshold": 0.5,
|
|
"top_k_num": -1,
|
|
"post_process_type": "mm"
|
|
}
|
|
|
|
|
|
def main(cfg_path, onnx_path):
|
|
out_init_params = deepcopy(init_params_template)
|
|
cfg = mmcv.Config.fromfile(cfg_path)
|
|
w, h = cfg.test_pipeline[1].img_scale
|
|
dataset = cfg.train_dataset.dataset.type[:-7].upper()
|
|
num_classes = len(mmdet.core.get_classes(dataset.lower()))
|
|
out_init_params['num_classes'] = num_classes
|
|
out_init_params['input_shape'] = NoIndent([h, w])
|
|
out_init_params['iou_threshold'] = cfg.model.test_cfg.nms.iou_threshold
|
|
out_init_params['model_path'] = os.path.abspath(onnx_path)
|
|
out_init_params['remapping_type'] = dataset
|
|
out_init_params_path = os.path.splitext(cfg_path)[0] + "_init_params.json"
|
|
|
|
s = json.dumps(out_init_params, indent=4, cls=MyEncoder)
|
|
with open(out_init_params_path, 'w') as f:
|
|
f.write(s)
|
|
|
|
out_DAG = deepcopy(DAG_template)
|
|
out_DAG['model_list'][0]['model_runner'] = cfg.model.type.lower()
|
|
out_DAG['model_list'][0][
|
|
'model_init_params_file'
|
|
] = os.path.abspath(out_init_params_path)
|
|
out_DAG_path = os.path.splitext(cfg_path)[0] + "_DAG.json"
|
|
|
|
s = json.dumps(out_DAG, indent=4, cls=MyEncoder)
|
|
with open(out_DAG_path, 'w') as f:
|
|
f.write(s)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
import sys
|
|
main(sys.argv[1], sys.argv[2])
|