adding icons and small tracking update

app.py

@@ -2,7 +2,7 @@ from webbrowser import get
 import gradio as gr
 import os
 import os.path as op
-from tools.files import download_from_url, create_unique_folder
+from tools.files import download_from_url, create_unique_folder, load_trash_icons
 
 import json
 from typing import Dict, List, Tuple
@@ -62,8 +62,11 @@ id_categories = {
 
 
 config_track = DotDict({
-    "confidence_threshold": 0.5,
-    "detection_threshold": 0.3,
+    "yolo_conf_thrld": 0.35,
+    "yolo_iou_thrld": 0.5,
+
+    "confidence_threshold": 0.004, # for the tracking part
+    "detection_threshold": 0.3, # for centernet
     "downsampling_factor": 4,
     "noise_covariances_path": "data/tracking_parameters",
     "output_shape": (960,544),
@@ -86,7 +89,7 @@ os.environ["VERBOSE"] = "False"
 URL_MODEL = "https://github.com/surfriderfoundationeurope/IA_Pau/releases/download/v0.1/yolov5.pt"
 FILE_MODEL = "yolov5.pt"
 model_path = download_from_url(URL_MODEL, FILE_MODEL, logger, "./models")
-model_yolo = load_model(model_path, config_track.device)
+model_yolo = load_model(model_path, config_track.device, config_track.yolo_conf_thrld, config_track.yolo_iou_thrld)
 
 
 logger.info('---Centernet model...')
@@ -113,6 +116,8 @@ video3_path = op.join("./data", FILE_DEMO3)
 JSON_FILE_PATH = "data/"
 
 
+labels2icons = load_trash_icons("./data/icons/")
+
 def track(args):
     device = torch.device("cpu")
 
@@ -206,14 +211,15 @@ def run_model(video_path, model_type, seconds, skip, tau, kappa, gps_file):
                                  progress_bar=True,
                                  preload=False,
                                  max_frame=config_track.max_length)
-    
+
     # Get GPS Data
     gps_data = get_filled_gps(gps_file,video_path)
 
     # Generate new video
     generate_video_with_annotations(reader, output_json, output_path,
                                     config_track.skip_frames, config_track.max_length,
-                                    config_track.downscale_output, logger,gps_data)
+                                    config_track.downscale_output, logger, gps_data=gps_data,
+                                    labels2icons=labels2icons)
     output_label = count_objects(output_json, id_categories)
 
 
@@ -225,7 +231,7 @@ def run_model(video_path, model_type, seconds, skip, tau, kappa, gps_file):
         with open(output_json_path) as json_file:
             predictions = json.load(json_file)
             trash_df = get_df_prediction(predictions, reader.fps)
-        if len(trash_df) != 0 :    
+        if len(trash_df) != 0 :
             # Get Trash prediction alongside GPS data
             trash_gps_df = get_trash_gps_df(trash_df,gps_data)
             trash_gps_geo_df = get_trash_gps_geo_df(trash_gps_df)
@@ -235,10 +241,10 @@ def run_model(video_path, model_type, seconds, skip, tau, kappa, gps_file):
             map_path = get_plastic_map(center_lat,center_long,trash_gps_geo_df,out_folder)
             html_content = codecs.open(map_path, 'r')
             map_html = html_content.read()
-            map_frame =  f"""<iframe style="width: 100%; height: 480px" name="result" allow="midi; geolocation; microphone; camera; 
-        display-capture; encrypted-media;" sandbox="allow-modals allow-forms 
-        allow-scripts allow-same-origin allow-popups 
-        allow-top-navigation-by-user-activation allow-downloads" allowfullscreen="" 
+            map_frame =  f"""<iframe style="width: 100%; height: 480px" name="result" allow="midi; geolocation; microphone; camera;
+        display-capture; encrypted-media;" sandbox="allow-modals allow-forms
+        allow-scripts allow-same-origin allow-popups
+        allow-top-navigation-by-user-activation allow-downloads" allowfullscreen=""
         allowpaymentrequest="" frameborder="0" srcdoc='{map_html}'></iframe>"""
 
     logger.info('---Surfnet End processing...')
@@ -275,7 +281,7 @@ def get_plastic_map(center_lat,center_long,trash_gps_gdf,out_folder)->str:
     Returns:
         map_html_path (str): full path to html map
     """
-    
+
     m = folium.Map([center_lat, center_long], zoom_start=16)
     locs = zip(trash_gps_gdf.geometry.y,trash_gps_gdf.geometry.x)
     labels = list(trash_gps_gdf['label'])

tools/files.py

@@ -1,11 +1,13 @@
 import os
 import os.path as op
+from pathlib import Path
 from urllib.request import urlretrieve
 import datetime
+import cv2
 
 
 def create_unique_folder(base_folder, filename):
-    """Creates a unique folder based on the filename and timestamp
+    """ Creates a unique folder based on the filename and timestamp
     """
     folder_name = op.splitext(op.basename(filename))[0] + "_out_"
     folder_name += datetime.datetime.now().strftime('%Y%m%d%H%M%S%f')
@@ -14,9 +16,9 @@ def create_unique_folder(base_folder, filename):
         os.mkdir(output_dir)
     return output_dir
 
+
 def download_from_url(url, filename, logger, folder="./data/"):
-    """
-    Download a file and place it in the corresponding folder if it does
+    """ Download a file and place it in the corresponding folder if it does
     not already exists
     """
     filepath = op.realpath(op.join(folder, filename))
@@ -26,3 +28,27 @@ def download_from_url(url, filename, logger, folder="./data/"):
     else:
         logger.info('---File already downloaded.')
     return filepath
+
+
+def load_trash_icons(folder_path):
+    """ loads all icons using cv2 format and returns a dict class -> opened icon
+    """
+    folder_path = Path(folder_path)
+    id_path = {
+        'Fragment': folder_path / "fragment.png",#'Fragment',    #'Sheet / tarp / plastic bag / fragment',
+        'Insulating': folder_path / "mousse.png",#'Insulating',  #'Insulating material',
+        'Bottle': folder_path / "bouteille.png",#'Bottle',      #'Bottle-shaped',
+        'Can': folder_path / "briquet.png",#'Can',         #'Can-shaped',
+        'Drum': folder_path / "contenant.png",#'Drum',
+        'Packaging': folder_path / "emballage.png",#'Packaging',   #'Other packaging',
+        'Tire': folder_path / "pneu.png",#'Tire',
+        'Fishing net': folder_path / "hamecon.png",#'Fishing net', #'Fishing net / cord',
+        'Easily namable': folder_path / "chaussure.png",#'Easily namable',
+        'Unclear': folder_path / "dechet.png"#'Unclear'
+    }
+    out_dict = {}
+    for idx, path in id_path.items():
+        img = cv2.imread(path.resolve().as_posix(), cv2.IMREAD_UNCHANGED)
+        resized_img = cv2.resize(img, (100,60), interpolation = cv2.INTER_AREA)
+        out_dict[idx] = resized_img
+    return out_dict

tracking/track_video.py

@@ -109,7 +109,7 @@ def track_video(reader, detections, args, engine, transition_variance, observati
     detections_for_frame, confs, labels = interpret_detection(detections_for_frame, args.downsampling_factor, is_yolo)
 
     max_distance = euclidean(reader.output_shape, np.array([0,0]))
-    delta = 0.05*max_distance
+    delta = 0.005*max_distance
 
     if display is not None and display.on:
 

tracking/utils.py

@@ -83,7 +83,45 @@ def get_detections_for_video(reader, detector, batch_size=16, device=None):
     return detections
 
 
-def generate_video_with_annotations(reader, output_detected, output_filename, skip_frames, maxframes, downscale, logger,gps_data=None):
+def overlay_transparent(background, overlay, x, y):
+    """ Overlays a transparent image over a background at topleft corner (x,y)
+    """
+    background_width = background.shape[1]
+    background_height = background.shape[0]
+
+    if x >= background_width or y >= background_height:
+        return background
+
+    h, w = overlay.shape[0], overlay.shape[1]
+
+    if x + w > background_width:
+        w = background_width - x
+        overlay = overlay[:, :w]
+
+    if y + h > background_height:
+        h = background_height - y
+        overlay = overlay[:h]
+
+    if overlay.shape[2] < 4:
+        overlay = np.concatenate(
+            [
+                overlay,
+                np.ones((overlay.shape[0], overlay.shape[1], 1), dtype = overlay.dtype) * 255
+            ],
+            axis = 2,
+        )
+
+    overlay_image = overlay[..., :3]
+    mask = overlay[..., 3:] / 255.0
+
+    background[y:y+h, x:x+w] = (1.0 - mask) * background[y:y+h, x:x+w] + mask * overlay_image
+    return background
+
+
+def generate_video_with_annotations(reader, output_detected, output_filename, skip_frames,
+                                    maxframes, downscale, logger, gps_data=None, labels2icons=None):
+    """ Generates output video at 24 fps, with optional gps_data
+    """
     fps = 24
     logger.info("---Intepreting json")
     results = defaultdict(list)
@@ -113,11 +151,17 @@ def generate_video_with_annotations(reader, output_detected, output_filename, sk
                                       '-vcodec': 'libx264',
                                       '-b': '5000000'})
 
-    font = cv2.FONT_HERSHEY_COMPLEX
+    font = cv2.FONT_HERSHEY_TRIPLEX
     for frame_nb, frame in enumerate(reader):
         detections_for_frame = results[frame_nb]
         for detection in detections_for_frame:
-            cv2.putText(frame, f'{detection[0]}/{detection[3]}', (int(detection[1]), int(detection[2])+5), font, 2, (0, 0, 255), 3, cv2.LINE_AA)
+            if labels2icons is None:
+                # write name of class
+                cv2.putText(frame, f'{detection[0]}/{detection[3]}', (int(detection[1]), int(detection[2])+5), font, 2, (0, 0, 255), 3, cv2.LINE_AA)
+            else:
+                # icons
+                overlay_transparent(frame, labels2icons[detection[3]], int(detection[1])+5, int(detection[2]))
+                cv2.putText(frame, f'{detection[0]}', (int(detection[1]+46+5), int(detection[2])+42), font, 1.2, (0, 0, 0), 2, cv2.LINE_AA)
 
             if gps_data is not None:
                 latitude  = gps_data[frame_nb//fps]['Latitude']
@@ -126,21 +170,10 @@ def generate_video_with_annotations(reader, output_detected, output_filename, sk
 
         frame = downscale_local_mean(frame, (downscale,downscale,1)).astype(np.uint8)
         writer.writeFrame(frame[:,:,::-1])
-        # moviepy version
-        # frames.append(frame[:,:,::-1])
-
-        #ret, frame, frame_nb = video.read()
-        #if frame_nb > maxframes:
-        #    break
 
     writer.close()
     reader.video.release()
 
-    # version with moviepy
-    #clip = ImageSequenceClip(sequence=frames, fps=fps)
-    #clip.write_videofile(output_filename, fps=fps)
-    #del frames
-
     logger.info("---finished writing video")