fixing border and adding logo

app.py

@@ -33,7 +33,7 @@ import pathlib
 import multiprocessing as mp
 from time import time
 
-if not os.path.exists('videos_example'):
+if not os.path.exists('videos_example') and not os.getenv('SYSTEM') == 'spaces':
     REPO_ID='SharkSpace/videos_examples'
     snapshot_download(repo_id=REPO_ID, token=os.environ.get('SHARK_MODEL'),repo_type='dataset',local_dir='videos_example')
 
@@ -65,6 +65,46 @@ def overlay_text_on_image(image, text_list, font=cv2.FONT_HERSHEY_SIMPLEX, font_
         cv2.putText(image, line, (image.shape[1] - text_width - margin, y), font, font_size, color, font_thickness, lineType=cv2.LINE_AA)
     return image
 
+def overlay_logo(frame,logo, position=(10, 10)):
+    """
+    Overlay a transparent logo (with alpha channel) on a frame.
+
+    Parameters:
+    - frame: The main image/frame to overlay the logo on.
+    - logo_path: Path to the logo image.
+    - position: (x, y) tuple indicating where the logo starts (top left corner).
+    """
+    # Load the logo and its alpha channel
+    alpha_channel = np.ones(logo.shape[:2], dtype=logo.dtype)
+    print(logo.min(),logo.max())
+    logo = np.dstack((logo, alpha_channel))
+    
+    indexes = logo[:,:,1]>150
+    logo[indexes,3] = 0
+    l_channels = cv2.split(logo)
+    if len(l_channels) != 4:
+        raise ValueError("Logo doesn't have an alpha channel!")
+    l_b, l_g, l_r, l_alpha = l_channels
+    cv2.imwrite('l_alpha.png',l_alpha*255)
+    # Extract regions of interest (ROI) from both images
+    roi = frame[position[1]:position[1]+logo.shape[0], position[0]:position[0]+logo.shape[1]]
+
+    # Blend the logo using the alpha channel
+    for channel in range(0, 3):
+        roi[:, :, channel] = (l_alpha ) * l_channels[channel] + (1.0 - l_alpha ) * roi[:, :, channel]
+
+    return frame
+    
+    
+def add_danger_symbol_from_image(frame, top_pred):  
+    relative = max(frame.shape[0],frame.shape[1])
+    if top_pred['shark_sighted'] and top_pred['dangerous_dist']:
+        # Add the danger symbol
+        danger_symbol = cv2.imread('static/danger_symbol.jpeg')
+        danger_symbol = cv2.resize(danger_symbol, (int(relative*0.1), int(relative*0.1)), interpolation = cv2.INTER_AREA)[:,:,::-1]
+        frame = overlay_logo(frame,danger_symbol, position=(int(relative*0.05), int(relative*0.05)))
+    return frame
+
 def draw_cockpit(frame, top_pred,cnt):
     # Bullet points:
     high_danger_color = (255,0,0)
@@ -80,17 +120,21 @@ def draw_cockpit(frame, top_pred,cnt):
     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))
+        frame = add_danger_symbol_from_image(frame, top_pred)
     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))
+         frame = add_danger_symbol_from_image(frame, top_pred)
+    else: 
+        
+        frame  = add_border(frame, color=(0,0,0), thickness=int(relative*0.025))
+        
     overlay_text_on_image(frame, strings, font=cv2.FONT_HERSHEY_SIMPLEX, font_size=relative*0.0007, font_thickness=1, margin=int(relative*0.05), color=(255, 255, 255))
     return frame
     
     
 
-def process_video(input_video, out_fps = 'auto', skip_frames = 7):
+def process_video(input_video,out_fps = 'auto', skip_frames = 7):
     cap = cv2.VideoCapture(input_video)
 
     output_path = "output.mp4"
@@ -110,6 +154,8 @@ def process_video(input_video, out_fps = 'auto', skip_frames = 7):
     cnt = 0
     
     while iterating:
+        print('overall count ', cnt)
+        
         if (cnt % skip_frames) == 0:
             print('starting Frame: ', cnt)
             # flip frame vertically
@@ -124,38 +170,39 @@ def process_video(input_video, out_fps = 'auto', skip_frames = 7):
             
             
             #frame = cv2.resize(frame, (int(width), int(height)))
+            video.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
             
             if cnt*skip_frames %2==0 and top_pred['shark_sighted']: 
-                #prediction_frame = cv2.resize(prediction_frame, (int(width), int(height)))
+                prediction_frame = cv2.resize(prediction_frame, (int(width), int(height)))
                 frame =prediction_frame
-             
+                
             if  top_pred['shark_sighted']: 
                 frame = draw_cockpit(frame, top_pred,cnt*skip_frames)
-            video.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
+                video.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
+                
+            
             pred_dashbord = prediction_dashboard(top_pred = top_pred)
             #print('sending frame')
             print('finalizing frame:',cnt)
             print(pred_dashbord.shape)
             print(frame.shape)
             print(prediction_frame.shape)
-            yield prediction_frame,frame , None, pred_dashbord
-        print('overall count ', cnt)
+            yield frame , None
+        
         cnt += 1
         iterating, frame = cap.read()
     
     video.release()
-    yield None, None, output_path, None
+    yield None,  output_path
 
 with gr.Blocks(theme=theme) as demo:
-    with gr.Row().style(equal_height=True,height='25%'):
+    with gr.Row().style(equal_height=True):
         input_video = gr.Video(label="Input")
-        processed_frames = gr.Image(label="Shark Engine")
-        output_video = gr.Video(label="Output Video")
-        dashboard = gr.Image(label="Dashboard")
-    
-    with gr.Row():
         
-        original_frames = gr.Image(label="Original Frame").style( height=650)
+        original_frames = gr.Image(label="Processed Frame").style( height=650)
+        #processed_frames = gr.Image(label="Shark Engine")
+        output_video = gr.Video(label="Output Video")
+        #dashboard = gr.Image(label="Events")
         
     with gr.Row():
         paths = sorted(pathlib.Path('videos_example/').rglob('*.mp4'))
@@ -163,8 +210,8 @@ with gr.Blocks(theme=theme) 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, dashboard])
-
+    #process_video_btn.click(process_video, input_video, [processed_frames, original_frames, output_video, dashboard])
+    process_video_btn.click(process_video, input_video, [ original_frames, output_video])
 demo.queue()
 if os.getenv('SYSTEM') == 'spaces':
     demo.launch(width='40%',auth=(os.environ.get('SHARK_USERNAME'), os.environ.get('SHARK_PASSWORD')))

inference.py

@@ -133,7 +133,7 @@ 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]
 
 
-if not os.path.exists('model'):
+if not os.path.exists('model') and not os.getenv('SYSTEM') == 'spaces':
     REPO_ID = "SharkSpace/maskformer_model"
     FILENAME = "mask2former"
     snapshot_download(repo_id=REPO_ID, token= os.environ.get('SHARK_MODEL'),local_dir='model/')

metrics.py

@@ -119,6 +119,8 @@ def get_min_distance_shark_person(top_pred, class_sizes = None, dangerous_distan
             '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]]),