How can I write a custom metric in tensorflow using opencv functions underneath?

Viewed 47

I am working with the TrackNet GitHub repository and trying to rewrite a metric in a way that I can use it in tensor operations. The project simply tries to track a badminton ball in 3 consecutive frames, it uses a u-net shape network that takes a tensor with shape (2, 1, 288, 512) and outputs an image with the same shape segmenting a point on the image which points to the location of the ball on the image.

For the evaluation metric, I'm trying to locate the point on both prediction and ground truth using opencv.findContours. Let's have a look:

def confusion(y_pred, y_true, tol):
    """
    compute confusion matrix value
    TP: True positive
    TN: True negative
    FP2: False positive
    FN: False negative
    FP1: If the distance of the ball center between 
        ground truth and prediction is larger than tolerance

    param:
    y_pred --> predicted heat map
    y_true --> ground truth heat map
    tol --> acceptable tolerance of heat map circle center 
            between ground truth and prediction
    """
    
    batch_size = y_pred.shape[0]
    TP = TN = FP1 = FP2 = FN = 0
    for b in range(batch_size):
        h_pred = y_pred[b]*255
        h_pred = h_pred.astype('uint8')
        h_true = y_true[b]*255
        h_true = h_true.astype('uint8')
        if np.amax(h_pred)==0 and np.amax(h_true)==0:
            TN += 1
        elif np.amax(h_pred)>0 and np.amax(h_true)==0:
            FP2 += 1
        elif np.amax(h_pred)==0 and np.amax(h_true)>0:
            FN += 1
        elif np.amax(h_pred)>0 and np.amax(h_true)>0:
            # find center of ball for prediction
            contours, _ = cv2.findContours(h_pred[0].copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            bboxes = [cv2.boundingRect(cnt) for cnt in contours]

            areas = np.array([bbox[2] * bbox[3] for bbox in bboxes])
            target = bboxes[np.argmax(areas)]
            x, y, w, h = target
            (cx_pred, cy_pred) = (int(x+w/2), int(y+h/2))

            # find center of ball for ground truth
            contours, _ = cv2.findContours(h_true.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            bboxes = [cv2.boundingRect(cnt) for cnt in contours]

            areas = np.array([bbox[2] * bbox[3] for bbox in bboxes])
            target = bboxes[np.argmax(areas)]
            x, y, w, h = target
            (cx_true, cy_true) = (int(x+w/2), int(y+h/2))

            dist = ((cx_pred-cx_true)**2 + (cy_pred-cy_true)**2)**0.5
            if dist > tol:
                FP1 += 1
            else:
                TP += 1
    
    return (TP, TN, FP1, FP2, FN)

And after I get these values, I would use this function to calculate accuracy, precision and recall:

def compute_acc(evaluation):
    """
    Compute accuracy, precision, and recall

    parame:
    evaluation --> a tuple containing 5 variable(TP, TN, FP1, FP2, FN)
    """
    (TP, TN, FP1, FP2, FN) = evaluation
    try:
        accuracy = (TP + TN) / (TP + TN + FP1 + FP2 + FN) 
    except:
        accuracy = 0
    try:
        precision = TP / (TP + FP1 + FP2)
    except:
        precision = 0
    try:
        recall = TP / (TP + FN)
    except:
        recall = 0
    return accuracy, precision, recall

The inputs and outputs of these functions are arrays. However, to reimplement the functions in tensorflow, I have to work with tensor inputs and outputs. I wonder if anyone can come up with a tensor-friendly operation that helps me to use it in Keras model metrics?

0 Answers
Related