add dashboard

app.py

@@ -24,6 +24,10 @@ import glob
 from inference import inference_frame,inference_frame_serial
 from inference import inference_frame_par_ready
 from inference import process_frame
+from inference import classes
+from inference import class_sizes_lower
+from metrics import process_results_for_plot
+from metrics import prediction_dashboard
 import os
 import pathlib
 import multiprocessing as mp
@@ -33,6 +37,11 @@ from time import time
 REPO_ID='SharkSpace/videos_examples'
 snapshot_download(repo_id=REPO_ID, token=os.environ.get('SHARK_MODEL'),repo_type='dataset',local_dir='videos_example')
 
+theme = gr.themes.Soft(
+    primary_hue="sky",
+    neutral_hue="slate",
+)
+
 def process_video(input_video, out_fps = 'auto', skip_frames = 7):
     cap = cv2.VideoCapture(input_video)
 
@@ -55,26 +64,31 @@ def process_video(input_video, out_fps = 'auto', skip_frames = 7):
     while iterating:
         if (cnt % skip_frames) == 0:
             # flip frame vertically
-            display_frame = inference_frame_serial(frame)
+            display_frame, result = inference_frame_serial(frame)
             video.write(cv2.cvtColor(display_frame, cv2.COLOR_BGR2RGB))
-            print('sending frame')
+            #print(result)
+            top_pred = process_results_for_plot(predictions = result.numpy(),
+                                                classes = classes,
+                                                class_sizes = class_sizes_lower)
+            pred_dashbord = prediction_dashboard(top_pred = top_pred)
+            #print('sending frame')
             print(cnt)
-            yield cv2.cvtColor(display_frame, cv2.COLOR_BGR2RGB), cv2.cvtColor(frame, cv2.COLOR_BGR2RGB), None
+            yield cv2.cvtColor(display_frame, cv2.COLOR_BGR2RGB), cv2.cvtColor(frame, cv2.COLOR_BGR2RGB), None, pred_dashbord
         cnt += 1
         iterating, frame = cap.read()
     
     video.release()
-    yield None, None, output_path
+    yield None, None, output_path, None
 
-with gr.Blocks() as demo:
+with gr.Blocks(theme=theme) as demo:
     with gr.Row():
         input_video = gr.Video(label="Input")
         output_video = gr.Video(label="Output Video")
     
     with gr.Row():
-        processed_frames = gr.Image(label="Live Frame")
-        # graphs = gr.Image(label="Graphs")
         original_frames = gr.Image(label="Original Frame")
+        dashboard = gr.Image(label="Dashboard")
+        processed_frames = gr.Image(label="Shark Engine")
 
     with gr.Row():
         paths = sorted(pathlib.Path('videos_example/').rglob('*.mp4'))
@@ -82,7 +96,7 @@ with gr.Blocks() as demo:
         examples = gr.Examples(samples, inputs=input_video)
         process_video_btn = gr.Button("Process Video")
 
-    process_video_btn.click(process_video, input_video, [processed_frames, original_frames, output_video])
+    process_video_btn.click(process_video, input_video, [processed_frames, original_frames, output_video, dashboard])
 
 demo.queue()
 if os.getenv('SYSTEM') == 'spaces':

inference.py

@@ -15,7 +15,7 @@ from huggingface_hub import hf_hub_download
 from huggingface_hub import snapshot_download
 from time import time
 
-classes= ['Beach',
+classes = ['Beach',
  'Sea',
  'Wave',
  'Rock',
@@ -68,6 +68,69 @@ classes= ['Beach',
  'Tiger shark',
  'Bull shark']*3
 
+class_sizes = {'Beach': None,
+               'Sea': None,
+               'Wave': None,
+               'Rock': None,
+               'Breaking wave': None,
+               'Reflection of the sea': None,
+               'Foam': None,
+               'Algae': None,
+               'Vegetation': None,
+               'Watermark': None,
+               'Bird': {'feet':[1, 3], 'meter': [0.3, 0.9], 'kg': [0.5, 1.5], 'pounds': [1, 3]},
+               'Ship': {'feet':[10, 100], 'meter': [3, 30], 'kg': [1000, 100000], 'pounds': [2200, 220000]},
+               'Boat': {'feet':[10, 45], 'meter': [3, 15], 'kg': [750, 80000], 'pounds': [1500, 160000]},
+               'Car': {'feet':[10, 20], 'meter': [3, 6], 'kg': [1000, 2000], 'pounds': [2200, 4400]},
+               'Kayak': {'feet':[10, 20], 'meter': [3, 6], 'kg': [50, 300], 'pounds': [100, 600]},
+               "Shark's line": None,
+               'Dock': None,
+               'Dog': {'feet':[1, 3], 'meter': [0.3, 0.9], 'kg': [10, 50], 'pounds': [20, 100]},
+               'Unidentifiable shade': None,
+               'Bird shadow': None,
+               'Boat shadow': None,
+               'Kayal shade': None,
+               'Surfer shadow': None,
+               'Shark shadow': None,
+               'Surfboard shadow': None,
+               'Crocodile': {'feet':[10, 20], 'meter': [3, 6], 'kg': [410, 1000], 'pounds': [900, 2200]},
+               'Sea cow':  {'feet':[9,12], 'meter': [3, 4], 'kg': [400, 590], 'pounds': [900, 1300]},
+               'Stingray': {'feet':[2, 7.5], 'meter': [0.6, 2.5], 'kg': [100, 300], 'pounds': [220, 770]},
+               'Person': {'feet':[5, 7], 'meter': [1.5, 2.1], 'kg': [50, 150], 'pounds': [110, 300]},
+               'Ocean': None,
+               'Surfer': {'feet':[5, 7], 'meter': [1.5, 2.1], 'kg': [50, 150], 'pounds': [110, 300]},
+               'Surfer': {'feet':[5, 7], 'meter': [1.5, 2.1], 'kg': [50, 150], 'pounds': [110, 300]},
+               'Fish':  {'feet':[1, 3], 'meter': [0.3, 0.9], 'kg': [20, 150], 'pounds': [40, 300]},
+               'Killer whale': {'feet':[10, 20], 'meter': [3, 6], 'kg': [3600, 5400], 'pounds': [8000, 12000]},
+               'Whale': {'feet':[15, 30], 'meter': [4.5, 10], 'kg': [2500, 80000], 'pounds': [55000, 176000]},
+               'Dolphin': {'feet':[6.6, 13.1], 'meter': [2, 4], 'kg': [150, 650], 'pounds': [330, 1430]},
+               'Miscellaneous': None,
+               'Unidentifiable shark': {'feet': [2, 15], 'meter': [0.6, 4.5], 'kg': [50, 1000], 'pounds': [110, 2200]},
+               'Carpet shark': {'feet': [4, 10], 'meter': [1.25, 3], 'kg': [50, 1000], 'pounds': [110, 2200]}, # Prob incorrect
+               'Dusty shark': {'feet': [9, 14], 'meter': [3, 4.25], 'kg': [160, 180], 'pounds': [350, 400]},
+               'Blue shark': {'feet': [7.9, 12.5], 'meter': [2.4, 3], 'kg': [60, 120], 'pounds': [130, 260]}, 
+               'Great white shark': {'feet': [13.1, 20], 'meter': [4, 6], 'kg': [680, 1800], 'pounds': [1500, 4000]},
+               'Copper shark': {'feet': [7.2, 10.8], 'meter': [2.2, 3.3], 'kg': [130, 300], 'pounds': [290, 660]},
+               'Nurse shark': {'feet': [7.9, 9.8], 'meter': [2.4, 3], 'kg': [90, 115], 'pounds': [200, 250]},
+               'Silky shark': {'feet': [6.6, 8.2], 'meter': [2, 2.5], 'kg': [300, 380], 'pounds': [660, 840]},
+               'Leopard shark': {'feet': [3.9, 4.9], 'meter': [1.2, 1.5], 'kg': [11, 20], 'pounds': [22, 44]},
+               'Shortfin mako shark': {'feet': [10.5, 12], 'meter': [3.2, 3.6], 'kg': [60, 135], 'pounds': [130, 300]},
+               'Hammerhead shark': {'feet': [4.9, 20], 'meter': [1.5, 6.1], 'kg': [230, 450], 'pounds': [500, 1000]},
+               'Oceanic whitetip shark': {'feet': [5.9, 9.8], 'meter': [1.8, 3], 'kg': [36, 170], 'pounds': [80, 375]},
+               'Blacktip shark': {'feet': [4.9, 6.6], 'meter': [1.5, 2], 'kg': [40, 100], 'pounds': [90, 220]},
+               'Tiger shark': {'feet': [9.8, 18], 'meter': [3, 5.5], 'kg': [385, 635], 'pounds': [850, 1400]},
+               'Bull shark': {'feet': [7.9, 11.2], 'meter': [2.4, 3.4], 'kg': [200, 315], 'pounds': [440, 690]},
+}
+
+class_sizes_lower = {k.lower(): v for k, v in class_sizes.items()}
+   
+classes_is_shark = [1 if 'shark' in x.lower() else 0 for x in classes]
+classes_is_human = [1 if 'person' or 'surfer' in x.lower() else 0 for x in classes]
+classes_is_unknown = [1 if 'unidentifiable' in x.lower() else 0 for x in classes]
+
+classes_is_shark_id = [i for i, x in enumerate(classes_is_shark) if x == 1]
+classes_is_human_id = [i for i, x in enumerate(classes_is_human) if x == 1]
+classes_is_unknown_id = [i for i, x in enumerate(classes_is_unknown) if x == 1]
 
 REPO_ID = "SharkSpace/maskformer_model"
 FILENAME = "mask2former"
@@ -101,7 +164,7 @@ palette = visualizer.dataset_meta.get('palette', None)
 print(len(classes))
 print(len(palette))
 def inference_frame_serial(image, visualize = True):
-    start = time()
+    #start = time()
     result = inference_detector(model, image)
     #print(f'inference time: {time()-start}')
     # show the results

metrics.py

@@ -0,0 +1,211 @@
+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()