support to show all frames

MEMTrack/src/GenerateVideo.py

@@ -3,32 +3,49 @@ import os
 import argparse
 from natsort import natsorted
 
-def create_video(data_dir, image_dir, video_name,fps):
+def create_video(data_dir, image_dir, video_name,fps, all_image_path=None, all_images=False):
     # choose codec according to format needed
     fourcc = cv2.VideoWriter_fourcc(*'mp4v') 
     #print(data_dir)
-    img_sample = cv2.imread(os.path.join(image_dir,"0.png"))
-    #print(img_sample.shape)
+    plotted_images = natsorted(os.listdir(image_dir))
+    images = plotted_images
+    if all_images:
+        all_images_list = natsorted(os.listdir(all_image_path))
+        images = all_images_list
+        
+    img_sample = cv2.imread(os.path.join(image_dir,images[0]))
+    print(images[0])
+    print(img_sample.shape)
     height, width, channels = img_sample.shape
     
     video = cv2.VideoWriter(data_dir + video_name + ".mp4", fourcc, fps, (width, height))
     print("Image dir: ", image_dir)
-    print("Number of frames: ", len(os.listdir(image_dir)))
+    print("Number of frames: ", len(images))
     print("Video output fps: ", fps)
-    for frame in natsorted(os.listdir(image_dir)):
-        #print(frame)
-        img = cv2.imread(os.path.join(image_dir, frame))
-        video.write(img)
+
+    for frame in natsorted(images):
+        if frame in plotted_images:
+            img = cv2.imread(os.path.join(image_dir, frame))
+            video.write(img)
+        else:
+            img = cv2.imread(os.path.join(all_image_path, frame))
+            video.write(img)
+        
     video.release()
 
-def gen_tracking_video(video_num, fps=60, data_path=None, custom_test_dir=None):
+def gen_tracking_video(video_num, fps=60, data_path=None, custom_test_dir=None, all_images=False):
     if custom_test_dir:
         data_dir = custom_test_dir
     else:
         data_dir = data_path + f"/data_video{video_num}_feature_optical_flow_median_back_2pyr_18win_background_img/"
-    image_dir = data_dir + "/test/tracklets-filtered/"
-    video_name = f'video{video_num}-tracklets-filtered-{fps}'
-    create_video(data_dir, image_dir, video_name,fps)
+    image_dir_plotted = data_dir + "/test/tracklets-filtered/"
+    if all_images:
+        video_name = f'video{video_num}-tracklets-filtered-{fps}-all-frames'
+    else:
+        video_name = f'video{video_num}-tracklets-filtered-{fps}'
+    all_img_path =  data_dir + f"/test/images/"
+    
+    create_video(data_dir, image_dir_plotted, video_name,fps, all_image_path=all_img_path, all_images=all_images)
     return os.path.join(data_dir, video_name)+ ".mp4"
 
 if __name__ == "__main__":

MEMTrack/src/TrackingAnalysis.py

@@ -298,7 +298,7 @@ def analyse_tracking(video_num, min_track_length=60, custom_test_dir=None, data_
     if plot :
         for image_id in image_id_filtered:
             #print(image_id)
-            newname = save_path + str(image_id) + '.png'
+            newname = save_path + str(image_id) + '.tif'
             det_img = cv2.imread(os.path.join(img_path,str(image_id))+".tif")
             det_img_gt_only = det_img.copy()
             det_img_p_only = det_img.copy()