metrics.py

import numpy as np
import matplotlib.pyplot as plt

import numpy as np 
def mask_to_boxes(mask):
    """ Convert a boolean (Height x Width) mask into a (N x 4) array of NON-OVERLAPPING bounding boxes
    surrounding "islands of truth" in the mask.  Boxes indicate the (Left, Top, Right, Bottom) bounds
    of each island, with Right and Bottom being NON-INCLUSIVE (ie they point to the indices AFTER the island).

    This algorithm (Downright Boxing) does not necessarily put separate connected components into
    separate boxes.

    You can "cut out" the island-masks with
        boxes = mask_to_boxes(mask)
        island_masks = [mask[t:b, l:r] for l, t, r, b in boxes]
    """
    max_ix = max(s+1 for s in mask.shape)   # Use this to represent background
    # These arrays will be used to carry the "box start" indices down and to the right.
    x_ixs = np.full(mask.shape, fill_value=max_ix)
    y_ixs = np.full(mask.shape, fill_value=max_ix)

    # Propagate the earliest x-index in each segment to the bottom-right corner of the segment
    for i in range(mask.shape[0]):
        x_fill_ix = max_ix
        for j in range(mask.shape[1]):
            above_cell_ix = x_ixs[i-1, j] if i>0 else max_ix
            still_active = mask[i, j] or ((x_fill_ix != max_ix) and (above_cell_ix != max_ix))
            x_fill_ix = min(x_fill_ix, j, above_cell_ix) if still_active else max_ix
            x_ixs[i, j] = x_fill_ix

    # Propagate the earliest y-index in each segment to the bottom-right corner of the segment
    for j in range(mask.shape[1]):
        y_fill_ix = max_ix
        for i in range(mask.shape[0]):
            left_cell_ix = y_ixs[i, j-1] if j>0 else max_ix
            still_active = mask[i, j] or ((y_fill_ix != max_ix) and (left_cell_ix != max_ix))
            y_fill_ix = min(y_fill_ix, i, left_cell_ix) if still_active else max_ix
            y_ixs[i, j] = y_fill_ix

    # Find the bottom-right corners of each segment
    new_xstops = np.diff((x_ixs != max_ix).astype(np.int32), axis=1, append=False)==-1
    new_ystops = np.diff((y_ixs != max_ix).astype(np.int32), axis=0, append=False)==-1
    corner_mask = new_xstops & new_ystops
    y_stops, x_stops = np.array(np.nonzero(corner_mask))

    # Extract the boxes, getting the top-right corners from the index arrays
    x_starts = x_ixs[y_stops, x_stops]
    y_starts = y_ixs[y_stops, x_stops]
    ltrb_boxes = np.hstack([x_starts[:, None], y_starts[:, None], x_stops[:, None]+1, y_stops[:, None]+1])
    return ltrb_boxes


def get_top_predictions(prediction = None, threshold = 0.8):
    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']]
        #print(top_pred['pred_above_thresh_labels'])
        #print(top_pred['pred_above_thresh_scores'])
        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 and 'shadow' not 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)
                    continue 
                else:
                    size_list.append(tmp_class_sizes['feet'])

                if 'shark' in tmp_pred.lower() and 'shadow' not 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 and len(shark_size_list) > 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)
                    continue
                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 and len(shark_weight_list) > 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'):
    if class_sizes[labels[0]] == None or class_sizes[labels[1]] == None:
        return 9999
    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):
    
    # Attempt to create bounding boxes for humans(surfers) detected.  
    # mask = predictions.pred_panoptic_seg.sem_seg[0]
    # print(np.unique(mask))
    # mask = mask >2000
    # bboxes = mask_to_boxes(mask)
    # bboxes = [box for box in bboxes if box[2]*box[3]>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_suspected': top_pred['shark_n'] > 0,
            'human_suspected': top_pred['human_n'] > 0,
            'shark_n': top_pred['shark_n'],
            #'human_n': max(top_pred['human_n'],len(bboxes)-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:
    if top_pred['shark_sighted'] > 0:
        shark_suspected = 'Shark Sighted !'
    elif top_pred['shark_suspected'] > 0:
        shark_suspected = 'Shark Suspected !'
    else:
        shark_suspected = 'No Sharks ...'

    if top_pred['human_sighted'] > 0:
        human_suspected = 'Human Sighted !'
    elif top_pred['human_suspected'] > 0:
        human_suspected = 'Human Suspected !'
    else:
        human_suspected = 'No Humans ...'

    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_confirmed'] else 'Low'

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

    # Create a list of strings to plot
    strings = [shark_suspected, human_suspected, 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()