def get_predictor(cls):
''' load trained model'''
with cls.lock:
# check if model is already loaded
if cls.predictor:
return cls.predictor
# create a mask r-cnn model
mask_rcnn_model = ResNetFPNModel()
try:
model_dir = os.environ['SM_MODEL_DIR']
except KeyError:
model_dir = '/opt/ml/model'
try:
cls.pretrained_model = os.environ['PRETRAINED_MODEL']
except KeyError:
pass
# file path to previoulsy trained mask r-cnn model
latest_trained_model = ""
model_search_path = os.path.join(model_dir, "model-*.index")
for model_file in glob.glob(model_search_path):
if model_file > latest_trained_model:
latest_trained_model = model_file
trained_model = latest_trained_model[:-6]
print(f'Using model: {trained_model}')
# fixed resnet50 backbone weights
cfg.BACKBONE.WEIGHTS = os.path.join(cls.pretrained_model)
cfg.MODE_FPN = True
cfg.MODE_MASK = True
cfg.TEST.RESULT_SCORE_THRESH = cfg.TEST.RESULT_SCORE_THRESH_VIS
finalize_configs(is_training=False)
# Create an inference model
# PredictConfig takes a model, input tensors and output tensors
input_tensors = mask_rcnn_model.get_inference_tensor_names()[0]
output_tensors = mask_rcnn_model.get_inference_tensor_names()[1]
cls.predictor = OfflinePredictor(
PredictConfig(model=mask_rcnn_model,
session_init=get_model_loader(trained_model),
input_names=input_tensors,
output_names=output_tensors))
return cls.predictor
def init_predictor():
register_coco(cfg.DATA.BASEDIR)
MODEL = ResNetFPNModel()
finalize_configs(is_training=False)
predcfg = PredictConfig(
model=MODEL,
#session_init=SmartInit("/home/jetson/Documents/trained_model/500000_17/checkpoint"),
session_init=SmartInit(
"/home/jetson/Documents/trained_model/255000_04.01/checkpoint"),
input_names=MODEL.get_inference_tensor_names()[0],
output_names=MODEL.get_inference_tensor_names()[1])
predictor = OfflinePredictor(predcfg)
return predictor
def evaluate_rcnn(model_name, paper_arxiv_id, cfg_list, model_file):
evaluator = COCOEvaluator(
root=COCO_ROOT, model_name=model_name, paper_arxiv_id=paper_arxiv_id
)
category_id_to_coco_id = {
v: k for k, v in COCODetection.COCO_id_to_category_id.items()
}
cfg.update_config_from_args(cfg_list) # TODO backup/restore config
finalize_configs(False)
MODEL = ResNetFPNModel() if cfg.MODE_FPN else ResNetC4Model()
predcfg = PredictConfig(
model=MODEL,
session_init=SmartInit(model_file),
input_names=MODEL.get_inference_tensor_names()[0],
output_names=MODEL.get_inference_tensor_names()[1],
)
predictor = OfflinePredictor(predcfg)
def xyxy_to_xywh(box):
box[2] -= box[0]
box[3] -= box[1]
return box
df = get_eval_dataflow("coco_val2017")
df.reset_state()
for img, img_id in tqdm.tqdm(df, total=len(df)):
results = predict_image(img, predictor)
res = [
{
"image_id": img_id,
"category_id": category_id_to_coco_id.get(
int(r.class_id), int(r.class_id)
),
"bbox": xyxy_to_xywh([round(float(x), 4) for x in r.box]),
"score": round(float(r.score), 3),
}
for r in results
]
evaluator.add(res)
if evaluator.cache_exists:
break
evaluator.save()
parser.add_argument(
'--benchmark',
action='store_true',
help="Benchmark the speed of the model + postprocessing")
parser.add_argument(
'--config',
help="A list of KEY=VALUE to overwrite those defined in config.py",
nargs='+')
parser.add_argument('--compact', help='Save a model to .pb')
parser.add_argument('--serving', help='Save a model to serving file')
args = parser.parse_args()
if args.config:
cfg.update_args(args.config)
register_coco(cfg.DATA.BASEDIR) # add COCO datasets to the registry
MODEL = ResNetFPNModel() if cfg.MODE_FPN else ResNetC4Model()
if not tf.test.is_gpu_available():
from tensorflow.python.framework import test_util
assert get_tf_version_tuple() >= (1, 7) and test_util.IsMklEnabled(), \
"Inference requires either GPU support or MKL support!"
assert args.load
finalize_configs(is_training=False)
if args.predict or args.visualize:
cfg.TEST.RESULT_SCORE_THRESH = cfg.TEST.RESULT_SCORE_THRESH_VIS
if args.visualize:
do_visualize(MODEL, args.load)
else:
predcfg = PredictConfig(
# add green rectangle arround original picture that with failure
height, width, channels = img.shape
cv2.rectangle(img, (0, 0), (width, height),
color=(100, 220, 80),
thickness=5)
viz = np.concatenate((img, final), axis=1)
cv2.imwrite(
"/home/jetson/tensorpack/examples/FasterRCNN/static/images/output.png",
viz)
logger.info("Inference output written to output.png")
if __name__ == '__main__':
register_coco(cfg.DATA.BASEDIR)
MODEL = ResNetFPNModel()
finalize_configs(is_training=False)
predcfg = PredictConfig(
model=MODEL,
session_init=SmartInit(
"/home/jetson/Documents/trained_model/500000_17/checkpoint"),
input_names=MODEL.get_inference_tensor_names()[0],
output_names=MODEL.get_inference_tensor_names()[1])
predictor = OfflinePredictor(predcfg)
do_predict(
predictor,
"/home/jetson/tensorpack/examples/FasterRCNN/static/images/original.jpg"
) # this line can be commented out, but the FIRST reference after service start will take longer
def get_predictor(cls):
"""load trained model"""
with cls.lock:
# check if model is already loaded
if cls.predictor:
return cls.predictor
# create a mask r-cnn model
mask_rcnn_model = ResNetFPNModel()
try:
model_dir = os.environ["SM_MODEL_DIR"]
except KeyError:
model_dir = "/opt/ml/model"
try:
resnet_arch = os.environ["RESNET_ARCH"]
except KeyError:
resnet_arch = "resnet50"
# file path to previoulsy trained mask r-cnn model
latest_trained_model = ""
model_search_path = os.path.join(model_dir, "model-*.index")
for model_file in glob.glob(model_search_path):
if model_file > latest_trained_model:
latest_trained_model = model_file
trained_model = latest_trained_model[:-6]
print(f"Using model: {trained_model}")
cfg.MODE_FPN = True
cfg.MODE_MASK = True
if resnet_arch == "resnet101":
cfg.BACKBONE.RESNET_NUM_BLOCKS = [3, 4, 23, 3]
else:
cfg.BACKBONE.RESNET_NUM_BLOCKS = [3, 4, 6, 3]
cfg_prefix = "CONFIG__"
for key, value in dict(os.environ).items():
if key.startswith(cfg_prefix):
attr_name = key[len(cfg_prefix) :]
attr_name = attr_name.replace("__", ".")
value = eval(value)
print(f"update config: {attr_name}={value}")
nested_var = cfg
attr_list = attr_name.split(".")
for attr in attr_list[0:-1]:
nested_var = getattr(nested_var, attr)
setattr(nested_var, attr_list[-1], value)
cfg.TEST.RESULT_SCORE_THRESH = cfg.TEST.RESULT_SCORE_THRESH_VIS
cfg.DATA.BASEDIR = "/data"
cfg.DATA.TRAIN = "coco_train2017"
cfg.DATA.VAL = "coco_val2017"
register_coco(cfg.DATA.BASEDIR)
finalize_configs(is_training=False)
# Create an inference model
# PredictConfig takes a model, input tensors and output tensors
input_tensors = mask_rcnn_model.get_inference_tensor_names()[0]
output_tensors = mask_rcnn_model.get_inference_tensor_names()[1]
cls.predictor = OfflinePredictor(
PredictConfig(
model=mask_rcnn_model,
session_init=get_model_loader(trained_model),
input_names=input_tensors,
output_names=output_tensors,
)
)
return cls.predictor