diff --git a/tutorials/frontend/from_darknet.py b/tutorials/frontend/from_darknet.py index f3eb05f9d20b..e90c8bb63b32 100644 --- a/tutorials/frontend/from_darknet.py +++ b/tutorials/frontend/from_darknet.py @@ -45,6 +45,11 @@ import tvm.relay.testing.yolo_detection import tvm.relay.testing.darknet +###################################################################### +# Choose the model +# ----------------------- +# Models are: 'yolov2', 'yolov3' or 'yolov3-tiny' + # Model name MODEL_NAME = 'yolov3' @@ -124,6 +129,11 @@ # execute print("Running the test image...") +# detection +# thresholds +thresh = 0.5 +nms_thresh = 0.45 + m.run() # get outputs tvm_out = [] @@ -155,9 +165,22 @@ layer_out['classes'] = layer_attr[4] tvm_out.append(layer_out) +elif MODEL_NAME == 'yolov3-tiny': + for i in range(2): + layer_out = {} + layer_out['type'] = 'Yolo' + # Get the yolo layer attributes (n, out_c, out_h, out_w, classes, total) + layer_attr = m.get_output(i*4+3).asnumpy() + layer_out['biases'] = m.get_output(i*4+2).asnumpy() + layer_out['mask'] = m.get_output(i*4+1).asnumpy() + out_shape = (layer_attr[0], layer_attr[1]//layer_attr[0], + layer_attr[2], layer_attr[3]) + layer_out['output'] = m.get_output(i*4).asnumpy().reshape(out_shape) + layer_out['classes'] = layer_attr[4] + tvm_out.append(layer_out) + thresh = 0.560 + # do the detection and bring up the bounding boxes -thresh = 0.5 -nms_thresh = 0.45 img = tvm.relay.testing.darknet.load_image_color(img_path) _, im_h, im_w = img.shape dets = tvm.relay.testing.yolo_detection.fill_network_boxes((netw, neth), (im_w, im_h), thresh,