adding cockpit view, code needs clean up

app.py

@@ -33,15 +33,63 @@ import pathlib
 import multiprocessing as mp
 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')
+if not os.path.exists('videos_example'):
+    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 add_border(frame, color = (255, 0, 0), thickness = 2):
+    # Add a red border to the image
+    relative = max(frame.shape[0],frame.shape[1])
+    top = int(relative*0.025)
+    bottom = int(relative*0.025)
+    left = int(relative*0.025)
+    right =  int(relative*0.025)
+    # Add the border to the image
+    bordered_image = cv2.copyMakeBorder(frame, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)
+    
+    return bordered_image 
+    
+def overlay_text_on_image(image, text_list, font=cv2.FONT_HERSHEY_SIMPLEX, font_size=0.5, font_thickness=1, margin=10, color=(255, 255, 255)):
+    relative = min(image.shape[0],image.shape[1])
+    y0, dy = margin, int(relative*0.1) # start y position and line gap
+    for i, line in enumerate(text_list):
+        y = y0 + i * dy
+        text_width, _ = cv2.getTextSize(line, font, font_size, font_thickness)[0]
+        cv2.putText(image, line, (image.shape[1] - text_width - margin, y), font, font_size, color, font_thickness, lineType=cv2.LINE_AA)
+    return image
+
+def draw_cockpit(frame, top_pred,cnt):
+    # Bullet points:
+    high_danger_color = (255,0,0)
+    low_danger_color = yellowgreen = (154,205,50)
+    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]
+    relative = max(frame.shape[0],frame.shape[1])
+    if top_pred['shark_sighted'] and top_pred['dangerous_dist'] and cnt%2 == 0:
+        relative = max(frame.shape[0],frame.shape[1])
+        frame  = add_border(frame, color=high_danger_color, thickness=int(relative*0.025))
+    elif top_pred['shark_sighted'] and not top_pred['dangerous_dist'] and cnt%2 == 0:
+         relative = max(frame.shape[0],frame.shape[1])
+         frame  = add_border(frame, color=low_danger_color, thickness=int(relative*0.025))
+    overlay_text_on_image(frame, strings, font=cv2.FONT_HERSHEY_SIMPLEX, font_size=1.5, font_thickness=3, margin=int(relative*0.05), color=(255, 255, 255))
+    return frame
+    
+    
+
 def process_video(input_video, out_fps = 'auto', skip_frames = 7):
     cap = cv2.VideoCapture(input_video)
 
@@ -70,10 +118,23 @@ def process_video(input_video, out_fps = 'auto', skip_frames = 7):
             top_pred = process_results_for_plot(predictions = result.numpy(),
                                                 classes = classes,
                                                 class_sizes = class_sizes_lower)
+            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            prediction_frame = cv2.cvtColor(display_frame, cv2.COLOR_BGR2RGB)
+            
+            
+            frame = cv2.resize(frame, (int(width*4), int(height*4)))
+            
+            if cnt*skip_frames %2==0 and top_pred['shark_sighted']: 
+                prediction_frame = cv2.resize(prediction_frame, (int(width*4), int(height*4)))
+                frame =prediction_frame
+             
+            if  top_pred['shark_sighted']: 
+                frame = draw_cockpit(frame, top_pred,cnt*skip_frames)
+            
             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, pred_dashbord
+            yield prediction_frame,frame , None, pred_dashbord
         cnt += 1
         iterating, frame = cap.read()
     
@@ -81,15 +142,15 @@ def process_video(input_video, out_fps = 'auto', skip_frames = 7):
     yield None, None, output_path, None
 
 with gr.Blocks(theme=theme) as demo:
-    with gr.Row():
+    with gr.Row().style(equal_height=True,height='50%'):
         input_video = gr.Video(label="Input")
+        processed_frames = gr.Image(label="Shark Engine")
         output_video = gr.Video(label="Output Video")
-    
-    with gr.Row():
-        original_frames = gr.Image(label="Original Frame")
         dashboard = gr.Image(label="Dashboard")
-        processed_frames = gr.Image(label="Shark Engine")
-
+    
+    with gr.Row(height='200%',width='200%'):
+        original_frames = gr.Image(label="Original Frame", width='200%', height='200%').style( width=1000)
+        
     with gr.Row():
         paths = sorted(pathlib.Path('videos_example/').rglob('*.mp4'))
         samples=[[path.as_posix()] for path in paths if 'raw_videos'  in str(path)]

inference.py

@@ -132,10 +132,11 @@ 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"
 
-snapshot_download(repo_id=REPO_ID, token= os.environ.get('SHARK_MODEL'),local_dir='model/')
+if not os.path.exists('model'):
+    REPO_ID = "SharkSpace/maskformer_model"
+    FILENAME = "mask2former"
+    snapshot_download(repo_id=REPO_ID, token= os.environ.get('SHARK_MODEL'),local_dir='model/')
 
 # Choose to use a config and initialize the detector
 config_file ='model/mask2former_swin-t-p4-w7-224_8xb2-lsj-50e_coco-panoptic/mask2former_swin-t-p4-w7-224_8xb2-lsj-50e_coco-panoptic.py'