metrics.py

import numpy as np
import matplotlib.pyplot as plt


def get_top_predictions(prediction = None, threshold = 0.7):
    if prediction is None:
        return None, None
    else:
        sorted_scores_ids = prediction.pred_instances.scores.argsort()[::-1]
        sorted_scores = prediction.pred_instances.scores[sorted_scores_ids]
        sorted_predictions = prediction.pred_instances.labels[sorted_scores_ids]
        return {'pred_above_thresh': sorted_predictions[sorted_scores > threshold], 
                'pred_above_thresh_id': sorted_scores_ids[sorted_scores > threshold],
                'pred_above_thresh_scores': sorted_scores[sorted_scores > threshold],
                'pred_above_thresh_bboxes': prediction.pred_instances['bboxes'][sorted_scores_ids][sorted_scores > threshold]}
    
def add_class_labels(top_pred = {}, class_labels = None):
    if class_labels == None:
        print('No class labels provided, returning original dictionary')
        return top_pred
    else:
        top_pred['pred_above_thresh_labels'] = [class_labels[x].lower() for x in top_pred['pred_above_thresh']]
        top_pred['any_detection'] = len(top_pred['pred_above_thresh_labels']) > 0
        if top_pred['any_detection']:
            # Get shark / human / unknown vectors
            top_pred['is_shark'] = np.array([1 if 'shark' in x else 0 for x in top_pred['pred_above_thresh_labels']])
            top_pred['is_human'] = np.array([1 if 'person' in x else 1 if 'surfer' in x else 0 for x in top_pred['pred_above_thresh_labels']])
            top_pred['is_unknown'] = np.array([1 if 'unidentifiable' in x else 0 for x in top_pred['pred_above_thresh_labels']])
            # Get shark / human / unknown  numbers of detections
            top_pred['shark_n'] = np.sum(top_pred['is_shark'])
            top_pred['human_n'] = np.sum(top_pred['is_human'])
            top_pred['unknown_n'] = np.sum(top_pred['is_unknown'])
        else:
            # Get shark / human / unknown vectors
            top_pred['is_shark'] = None
            top_pred['is_human'] = None
            top_pred['is_unknown'] = None
            # Get shark / human / unknown  numbers of detections
            top_pred['shark_n'] = 0
            top_pred['human_n'] = 0
            top_pred['unknown_n'] = 0
        return top_pred

def add_class_sizes(top_pred = {}, class_sizes = None):
    size_list = []
    shark_size_list = []
    if top_pred['any_detection']:
        for tmp_pred in top_pred['pred_above_thresh_labels']:
                tmp_class_sizes = class_sizes[tmp_pred.lower()]
                if tmp_class_sizes == None:
                    size_list.append(None)
                else:
                    size_list.append(tmp_class_sizes['feet'])

                if 'shark' in tmp_pred.lower():
                    shark_size_list.append(np.mean(tmp_class_sizes['feet']))

        top_pred['pred_above_thresh_sizes'] = size_list

        if top_pred['shark_n'] > 0:
            top_pred['biggest_shark_size'] = np.max(shark_size_list)
        else:
            top_pred['biggest_shark_size'] = None
    else:
        top_pred['pred_above_thresh_sizes'] = None
        top_pred['biggest_shark_size'] = None
    return top_pred

def add_class_weights(top_pred = {}, class_weights = None):
    weight_list = []
    shark_weight_list = []
    if top_pred['any_detection']:
        for tmp_pred in top_pred['pred_above_thresh_labels']:
                tmp_class_weights = class_weights[tmp_pred.lower()]
                if tmp_class_weights == None:
                    weight_list.append(None)
                else:
                    weight_list.append(tmp_class_weights['pounds'])

                if 'shark' in tmp_pred.lower():
                    shark_weight_list.append(np.mean(tmp_class_weights['pounds']))

        top_pred['pred_above_thresh_weights'] = weight_list

        if top_pred['shark_n'] > 0:
            top_pred['biggest_shark_weight'] = np.max(shark_weight_list)
        else:
            top_pred['biggest_shark_weight'] = None
    else:
        top_pred['pred_above_thresh_weights'] = None
        top_pred['biggest_shark_weight'] = None
    return top_pred

# Sizes
def get_min_distance_shark_person(top_pred, class_sizes = None, dangerous_distance = 100):
    min_dist = 99999
    dist_calculated = False
    # Calculate distance for every pairing of human and shark
    # and accumulate the min distance
    for i, tmp_shark in enumerate(top_pred['is_shark']):
        for j, tmp_person in enumerate(top_pred['is_human']):
            if tmp_shark == 1 and tmp_person == 1:
                dist_calculated = True
                #print(top_pred['pred_above_thresh_bboxes'][i])
                #print(top_pred['pred_above_thresh_bboxes'][j])
                tmp_dist_feed = _calculate_dist_estimate(top_pred['pred_above_thresh_bboxes'][i], 
                                                         top_pred['pred_above_thresh_bboxes'][j], 
                                                         [top_pred['pred_above_thresh_labels'][i], top_pred['pred_above_thresh_labels'][j]],
                                                         class_sizes,
                                                         measurement = 'feet')
                #print(tmp_dist_feed)
                min_dist = min(min_dist, tmp_dist_feed)
            else:
                pass
    return {'min_dist': str(round(min_dist,1)) + ' feet' if dist_calculated else '', 
            'any_dist_calculated': dist_calculated, 
            'dangerous_dist': min_dist < dangerous_distance}

def _calculate_dist_estimate(bbox1, bbox2, labels, class_sizes = None, measurement = 'feet'):
    class_feet_size_mean = np.array([class_sizes[labels[0]][measurement][0], 
                                     class_sizes[labels[1]][measurement][0]]).mean()
    box_pixel_size_mean = np.array([np.linalg.norm(bbox1[[0, 1]] - bbox1[[2, 3]]), 
                                    np.linalg.norm(bbox2[[0, 1]] - bbox2[[2, 3]])]).mean()
    
    # Calculate the max size of the two boxes
    box_center_1 = np.array([(bbox1[2] - bbox1[0])/2 + bbox1[0], 
                             (bbox1[3] - bbox1[1])/2 + bbox1[1]])
    box_center_2 = np.array([(bbox2[2] - bbox2[0])/2 + bbox2[0], 
                             (bbox2[3] - bbox2[1])/2 + bbox2[1]])
    
    # Return ratio distance
    return np.linalg.norm(box_center_1 - box_center_2) / box_pixel_size_mean * class_feet_size_mean

# bboxes info!
# 1 x1 (left, lower pixel number)
# 2 y1 (top , lower pixel number)
# 3 x2 (right, higher pixel number)
# 4 y2 (bottom, higher pixel number)

def process_results_for_plot(predictions = None, threshold = 0.5, classes = None,
                             class_sizes = None, dangerous_distance = 100):
    
    top_pred = get_top_predictions(predictions, threshold = threshold)
    top_pred = add_class_labels(top_pred, class_labels = classes)
    top_pred = add_class_sizes(top_pred, class_sizes = class_sizes)
    top_pred = add_class_weights(top_pred, class_weights = class_sizes)
    
    if len(top_pred['pred_above_thresh']) > 0:
        min_dist = get_min_distance_shark_person(top_pred, class_sizes = class_sizes)
    else:
        min_dist = {'any_dist_calculated': False,
                    'min_dist': '',
                    'dangerous_dist': False}

    return {'min_dist_str': min_dist['min_dist'], 
            'shark_sighted': top_pred['shark_n'] > 0,
            'human_sighted': top_pred['human_n'] > 0,
            'shark_n': top_pred['shark_n'],
            'human_n': top_pred['human_n'],
            'human_and_shark': (top_pred['shark_n'] > 0) and (top_pred['human_n'] > 0),
            'dangerous_dist': min_dist['dangerous_dist'],
            'dist_calculated': min_dist['any_dist_calculated'],
            'biggest_shark_size': '' if top_pred['biggest_shark_size'] == None else str(round(top_pred['biggest_shark_size'],1)) + ' feet',
            'biggest_shark_weight': '' if top_pred['biggest_shark_weight'] == None else str(round(top_pred['biggest_shark_weight'],1)) + ' pounds',
            }

def prediction_dashboard(top_pred = None):
    # Bullet points:
    shark_sighted = 'Shark Detected: ' + str(top_pred['shark_sighted'])
    human_sighted = 'Number of Humans: ' + str(top_pred['human_n'])

    shark_size_estimate = 'Biggest shark size: ' + str(top_pred['biggest_shark_size'])
    shark_weight_estimate = 'Biggest shark weight: ' + str(top_pred['biggest_shark_weight'])

    danger_level = 'Danger Level: ' 
    danger_level += 'High' if top_pred['dangerous_dist'] else 'Low'

    danger_color = 'orangered' if top_pred['dangerous_dist'] else 'yellowgreen'

    # Create a list of strings to plot
    strings = [shark_sighted, human_sighted, shark_size_estimate, shark_weight_estimate, danger_level]

    # Create a figure and axis
    fig, ax = plt.subplots()
    fig.set_facecolor((35/255,40/255,54/255))

    # Hide axes
    ax.axis('off')

    # Position for starting to place text, starting from top
    y_pos = 0.7

    # Iterate through list and place each item as text on the plot
    for s in strings:
        if 'danger' in s.lower():
            ax.text(0.05, y_pos, s, transform=ax.transAxes, fontsize=16, color=danger_color)
        else:
            ax.text(0.05, y_pos, s, transform=ax.transAxes, fontsize=16, color=(0, 204/255, 153/255))
        y_pos -= 0.1  # move down for next item

    # plt.tight_layout()
    # If we haven't already shown or saved the plot, then we need to
    # draw the figure first...
    fig.canvas.draw();

    # Now we can save it to a numpy array.
    data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
    data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
    plt.close()
    #plt.savefig('tmp.png', format='png')
    return data #plt.show()