def run_word_segmentation_net()

in sagemaker/src/run_inference.py [0:0]

• 37 lines of code
• 4 McCabe index (conditional complexity)

def run_word_segmentation_net(nets, input_image, ctx, min_c=0.1, overlap_thres=0.1, topk=600):
    x = nets["body"][0](input_image)
    
    anchor_sizes = [[.1, .2], [.2, .3], [.2, .4], [.3, .4], [.3, .5], [.4, .6]]
    anchor_ratios = [[1, 3, 5], [1, 3, 5], [1, 6, 8], [1, 4, 7], [1, 6, 8], [1, 5, 7]]

    num_anchors = len(anchor_sizes)
    num_classes = 2
    
    # Run SSD
    default_anchors = []
    predicted_boxes = []
    predicted_classes = []

    for i in range(num_anchors):
        default_anchors.append(mx.contrib.ndarray.MultiBoxPrior(x, sizes=anchor_sizes[i], ratios=anchor_ratios[i]))
        predicted_boxes.append(_flatten_prediction(nets["box_preds"][i](x)))
        predicted_classes.append(_flatten_prediction(nets["class_preds"][i](x)))
        if i < len(nets["downsamples"]):
            x = nets["downsamples"][i](x)
        elif i == 3:
            x = mx.nd.Pooling(x, global_pool=True, pool_type='max', kernel=(4, 4))
    
    anchors = mx.nd.concat(*default_anchors, dim=1)
    box_preds = mx.nd.concat(*predicted_boxes, dim=1)
    class_preds = mx.nd.concat(*predicted_classes, dim=1)
    class_preds = mx.nd.reshape(class_preds, shape=(0, -1, num_classes + 1))
        
    # Do SSD prediction
    bb = np.zeros(shape=(13, 5))
    bb = mx.nd.array(bb)
    bb = bb.as_in_context(ctx)
    bb = bb.expand_dims(axis=0)
    
    class_preds = mx.nd.transpose(class_preds, axes=(0, 2, 1))
    box_target, box_mask, cls_target = mx.contrib.ndarray.MultiBoxTarget(
        anchors, bb, class_preds)

    cls_probs = mx.nd.SoftmaxActivation(class_preds, mode='channel')
    predicted_bb = mx.contrib.ndarray.MultiBoxDetection(*[cls_probs, box_preds, anchors], force_suppress=True, clip=False)
    predicted_bb = mx.contrib.ndarray.box_nms(predicted_bb, overlap_thresh=overlap_thres, valid_thresh=min_c, topk=topk)
    predicted_bb = predicted_bb.asnumpy()
    predicted_bb = predicted_bb[0, predicted_bb[0, :, 0] != -1]
    predicted_bb = predicted_bb[:, 2:]
    predicted_bb[:, 2] = predicted_bb[:, 2] - predicted_bb[:, 0]
    predicted_bb[:, 3] = predicted_bb[:, 3] - predicted_bb[:, 1]

    return predicted_bb