Update dataframe creation

README.md

@@ -170,6 +170,6 @@ player --folder path_to_folder [--save] [--imshow]
 
 
 ```
-cvat2slowfast --miniscene path_to_mini_scenes --dataset dataset_name --classes path_to_classes_json [--old2new path_to_old2new_json]
+cvat2slowfast --miniscene path_to_mini_scenes --dataset dataset_name --classes path_to_classes_json [--old2new path_to_old2new_json] [--no_images]
 ```
 

src/kabr_tools/cvat2slowfast.py

@@ -11,7 +11,7 @@
 
 
 def cvat2slowfast(path_to_mini_scenes: str, path_to_new_dataset: str,
-                  label2number: dict, old2new: Optional[dict]) -> None:
+                  label2number: dict, old2new: Optional[dict], no_images: bool) -> None:
     """
     Convert CVAT annotations to the dataset in Charades format.
 
@@ -20,6 +20,7 @@ def cvat2slowfast(path_to_mini_scenes: str, path_to_new_dataset: str,
     path_to_new_dataset - str. Path to the folder to output dataset files.
     label2number - dict. Mapping of ethogram labels to integers.
     old2new - dict [optional]. Mapping of old ethogram labels to new ethogram labels.
+    no_images - bool. Flag to stop image output.
     """
     if not os.path.exists(path_to_new_dataset):
         os.makedirs(path_to_new_dataset)
@@ -33,12 +34,13 @@ def cvat2slowfast(path_to_mini_scenes: str, path_to_new_dataset: str,
     with open(f"{path_to_new_dataset}/annotation/classes.json", "w") as file:
         json.dump(label2number, file)
 
-    headers = {"original_vido_id": [], "video_id": pd.Series(dtype="int"), "frame_id": pd.Series(dtype="int"),
-               "path": [], "labels": []}
+    headers = ["original_vido_id", "video_id", "frame_id", "path", "labels"]
+    charades_data = []
+
     charades_df = pd.DataFrame(data=headers)
     video_id = 1
     folder_name = 1
-    flag = False
+    flag = not no_images
 
     for i, folder in enumerate(natsorted(os.listdir(path_to_mini_scenes))):
         if os.path.exists(f"{path_to_mini_scenes}/{folder}/actions"):
@@ -127,13 +129,11 @@ def cvat2slowfast(path_to_mini_scenes: str, path_to_new_dataset: str,
                                 if flag:
                                     cv2.imwrite(f"{output_folder}/{adjusted_index}.jpg", frame)
 
-                                # TODO: Major slow down here. Add to a list rather than dataframe,
-                                #  and create dataframe at the end.
-                                charades_df.loc[len(charades_df.index)] = [f"{folder_code}",
-                                                                           video_id,
-                                                                           adjusted_index,
-                                                                           f"{folder_code}/{adjusted_index}.jpg",
-                                                                           str(label2number[behavior])]
+                                charades_data.append([f"{folder_code}",
+                                                        video_id,
+                                                        adjusted_index,
+                                                        f"{folder_code}/{adjusted_index}.jpg",
+                                                        str(label2number[behavior])])
 
                                 adjusted_index += 1
 
@@ -145,9 +145,11 @@ def cvat2slowfast(path_to_mini_scenes: str, path_to_new_dataset: str,
                     video_id += 1
 
                     if video_id % 10 == 0:
+                        charades_df = pd.DataFrame(charades_data, columns=headers)
                         charades_df.to_csv(
                             f"{path_to_new_dataset}/annotation/data.csv", sep=" ", index=False)
 
+    charades_df = pd.DataFrame(charades_data, columns=headers)
     charades_df.to_csv(
         f"{path_to_new_dataset}/annotation/data.csv", sep=" ", index=False)
 
@@ -178,6 +180,11 @@ def parse_args() -> argparse.Namespace:
         help="path to old to new ethogram labels json",
         required=False
     )
+    local_parser.add_argument(
+        "--no_images",
+        action="store_true",
+        help="flag to stop image output"
+    )
     return local_parser.parse_args()
 
 
@@ -194,8 +201,8 @@ def main() -> None:
     else:
         old2new = None
 
-    cvat2slowfast(args.miniscene, args.dataset, label2number, old2new)
+    cvat2slowfast(args.miniscene, args.dataset, label2number, old2new, args.no_images)
 
 
 if __name__ == "__main__":
-    main()
+    main()

src/kabr_tools/miniscene2behavior.py

@@ -19,12 +19,15 @@ def get_input_clip(cap: cv2.VideoCapture, cfg: CfgNode, keyframe_idx: int) -> li
     # https://github.com/facebookresearch/SlowFast/blob/bac7b672f40d44166a84e8c51d1a5ba367ace816/slowfast/visualization/ava_demo_precomputed_boxes.py
     seq_length = cfg.DATA.NUM_FRAMES * cfg.DATA.SAMPLING_RATE
     total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    assert keyframe_idx < total_frames, f"keyframe_idx: {keyframe_idx}" \
+        f" >= total_frames: {total_frames}"
     seq = get_sequence(
         keyframe_idx,
         seq_length // 2,
         cfg.DATA.SAMPLING_RATE,
         total_frames,
     )
+
     clip = []
     for frame_idx in seq:
         cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
@@ -124,29 +127,34 @@ def annotate_miniscene(cfg: CfgNode, model: torch.nn.Module,
 
     # find all tracks
     tracks = []
+    frames = {}
     for track in root.iterfind("track"):
         track_id = track.attrib["id"]
         tracks.append(track_id)
+        frames[track_id] = []
 
-    # find all frames
-    # TODO: rewrite - some tracks may have different frames
-    assert len(tracks) > 0, "No tracks found in track file"
-    frames = []
-    for box in track.iterfind("box"):
-        frames.append(int(box.attrib["frame"]))
+        # find all frames
+        for box in track.iterfind("box"):
+            frames[track_id].append(int(box.attrib["frame"]))
 
     # run model on miniscene
     for track in tracks:
         video_file = f"{miniscene_path}/{track}.mp4"
         cap = cv2.VideoCapture(video_file)
-        for frame in tqdm(frames, desc=f"{track} frames"):
-            inputs = get_input_clip(cap, cfg, frame)
+        index = 0
+        for frame in tqdm(frames[track], desc=f"{track} frames"):
+            try:
+                inputs = get_input_clip(cap, cfg, index)
+            except AssertionError as e:
+                print(e)
+                break
+            index += 1
 
             if cfg.NUM_GPUS:
                 # transfer the data to the current GPU device.
                 if isinstance(inputs, (list,)):
-                    for i in range(len(inputs)):
-                        inputs[i] = inputs[i].cuda(non_blocking=True)
+                    for i, input_clip in enumerate(inputs):
+                        inputs[i] = input_clip.cuda(non_blocking=True)
                 else:
                     inputs = inputs.cuda(non_blocking=True)
 
@@ -163,6 +171,7 @@ def annotate_miniscene(cfg: CfgNode, model: torch.nn.Module,
             if frame % 20 == 0:
                 pd.DataFrame(label_data).to_csv(
                     output_path, sep=" ", index=False)
+        cap.release()
     pd.DataFrame(label_data).to_csv(output_path, sep=" ", index=False)
 
 

tests/examples/DETECTOR1/DJI_tracks.xml

@@ -0,0 +1,49 @@
+<?xml version='1.0' encoding='UTF-8'?>
+<annotations>
+  <track id="0" label="Grevy" source="manual">
+    <box frame="0" outside="0" occluded="0" keyframe="1" xtl="2011.00" ytl="1883.00" xbr="2133.00" ybr="1978.00" z_order="0"/>
+    <box frame="1" outside="0" occluded="0" keyframe="0" xtl="2011.00" ytl="1883.00" xbr="2133.00" ybr="1978.00" z_order="0"/>
+    <box frame="2" outside="0" occluded="0" keyframe="0" xtl="2011.00" ytl="1883.00" xbr="2133.00" ybr="1978.00" z_order="0"/>
+    <box frame="3" outside="0" occluded="0" keyframe="0" xtl="2011.00" ytl="1883.00" xbr="2133.00" ybr="1978.00" z_order="0"/>
+    <box frame="4" outside="0" occluded="0" keyframe="0" xtl="2011.00" ytl="1883.00" xbr="2134.00" ybr="1978.00" z_order="0"/>
+    <box frame="5" outside="0" occluded="0" keyframe="0" xtl="2011.00" ytl="1883.00" xbr="2134.00" ybr="1978.00" z_order="0"/>
+    <box frame="6" outside="0" occluded="0" keyframe="0" xtl="2011.00" ytl="1883.00" xbr="2134.00" ybr="1978.00" z_order="0"/>
+    <box frame="7" outside="0" occluded="0" keyframe="0" xtl="2012.00" ytl="1883.00" xbr="2134.00" ybr="1978.00" z_order="0"/>
+    <box frame="8" outside="0" occluded="0" keyframe="0" xtl="2012.00" ytl="1883.00" xbr="2134.00" ybr="1978.00" z_order="0"/>
+    <box frame="9" outside="0" occluded="0" keyframe="0" xtl="2012.00" ytl="1883.00" xbr="2134.00" ybr="1978.00" z_order="0"/>
+    <box frame="10" outside="0" occluded="0" keyframe="0" xtl="2012.00" ytl="1882.00" xbr="2134.00" ybr="1978.00" z_order="0"/>
+    <box frame="21" outside="0" occluded="0" keyframe="0" xtl="2012.00" ytl="1882.00" xbr="2134.00" ybr="1978.00" z_order="0"/>
+    <box frame="22" outside="0" occluded="0" keyframe="0" xtl="2012.00" ytl="1882.00" xbr="2134.00" ybr="1978.00" z_order="0"/>
+    <box frame="23" outside="0" occluded="0" keyframe="0" xtl="2012.00" ytl="1882.00" xbr="2135.00" ybr="1978.00" z_order="0"/>
+    <box frame="24" outside="0" occluded="0" keyframe="0" xtl="2012.00" ytl="1882.00" xbr="2135.00" ybr="1977.00" z_order="0"/>
+    <box frame="25" outside="0" occluded="0" keyframe="0" xtl="2012.00" ytl="1882.00" xbr="2135.00" ybr="1977.00" z_order="0"/>
+    <box frame="26" outside="0" occluded="0" keyframe="0" xtl="2013.00" ytl="1882.00" xbr="2135.00" ybr="1977.00" z_order="0"/>
+    <box frame="27" outside="0" occluded="0" keyframe="0" xtl="2013.00" ytl="1882.00" xbr="2135.00" ybr="1977.00" z_order="0"/>
+    <box frame="28" outside="0" occluded="0" keyframe="0" xtl="2013.00" ytl="1882.00" xbr="2135.00" ybr="1977.00" z_order="0"/>
+    <box frame="29" outside="0" occluded="0" keyframe="0" xtl="2013.00" ytl="1882.00" xbr="2135.00" ybr="1977.00" z_order="0"/>
+    <box frame="30" outside="0" occluded="0" keyframe="0" xtl="2013.00" ytl="1882.00" xbr="2135.00" ybr="1977.00" z_order="0"/>
+  </track>
+  <track id="1" label="Grevy" source="manual">
+    <box frame="0" outside="0" occluded="0" keyframe="1" xtl="2513.00" ytl="1921.00" xbr="2615.00" ybr="1991.00" z_order="0"/>
+    <box frame="1" outside="0" occluded="0" keyframe="0" xtl="2513.00" ytl="1921.00" xbr="2615.00" ybr="1991.00" z_order="0"/>
+    <box frame="2" outside="0" occluded="0" keyframe="0" xtl="2512.00" ytl="1921.00" xbr="2614.00" ybr="1991.00" z_order="0"/>
+    <box frame="3" outside="0" occluded="0" keyframe="0" xtl="2512.00" ytl="1921.00" xbr="2613.00" ybr="1991.00" z_order="0"/>
+    <box frame="4" outside="0" occluded="0" keyframe="0" xtl="2511.00" ytl="1921.00" xbr="2613.00" ybr="1991.00" z_order="0"/>
+    <box frame="5" outside="0" occluded="0" keyframe="0" xtl="2510.00" ytl="1921.00" xbr="2612.00" ybr="1991.00" z_order="0"/>
+    <box frame="6" outside="0" occluded="0" keyframe="0" xtl="2510.00" ytl="1921.00" xbr="2612.00" ybr="1991.00" z_order="0"/>
+    <box frame="7" outside="0" occluded="0" keyframe="0" xtl="2509.00" ytl="1921.00" xbr="2611.00" ybr="1991.00" z_order="0"/>
+    <box frame="8" outside="0" occluded="0" keyframe="0" xtl="2508.00" ytl="1921.00" xbr="2610.00" ybr="1991.00" z_order="0"/>
+    <box frame="9" outside="0" occluded="0" keyframe="0" xtl="2508.00" ytl="1921.00" xbr="2610.00" ybr="1991.00" z_order="0"/>
+    <box frame="10" outside="0" occluded="0" keyframe="0" xtl="2507.00" ytl="1921.00" xbr="2609.00" ybr="1991.00" z_order="0"/>
+    <box frame="11" outside="0" occluded="0" keyframe="0" xtl="2507.00" ytl="1921.00" xbr="2608.00" ybr="1991.00" z_order="0"/>
+    <box frame="12" outside="0" occluded="0" keyframe="0" xtl="2506.00" ytl="1921.00" xbr="2608.00" ybr="1991.00" z_order="0"/>
+    <box frame="13" outside="0" occluded="0" keyframe="0" xtl="2505.00" ytl="1921.00" xbr="2607.00" ybr="1991.00" z_order="0"/>
+    <box frame="14" outside="0" occluded="0" keyframe="0" xtl="2505.00" ytl="1921.00" xbr="2607.00" ybr="1991.00" z_order="0"/>
+    <box frame="15" outside="0" occluded="0" keyframe="0" xtl="2504.00" ytl="1920.00" xbr="2606.00" ybr="1991.00" z_order="0"/>
+    <box frame="16" outside="0" occluded="0" keyframe="0" xtl="2503.00" ytl="1920.00" xbr="2605.00" ybr="1991.00" z_order="0"/>
+    <box frame="17" outside="0" occluded="0" keyframe="0" xtl="2503.00" ytl="1920.00" xbr="2605.00" ybr="1991.00" z_order="0"/>
+    <box frame="18" outside="0" occluded="0" keyframe="0" xtl="2502.00" ytl="1920.00" xbr="2604.00" ybr="1991.00" z_order="0"/>
+    <box frame="19" outside="0" occluded="0" keyframe="0" xtl="2502.00" ytl="1920.00" xbr="2603.00" ybr="1990.00" z_order="0"/>
+    <box frame="20" outside="0" occluded="0" keyframe="0" xtl="2501.00" ytl="1920.00" xbr="2603.00" ybr="1990.00" z_order="0"/>
+  </track>
+</annotations>

tests/test_cvat2slowfast.py

@@ -1,11 +1,20 @@
 import unittest
 import sys
 import os
+import json
+import pandas as pd
+import cv2
 from kabr_tools import cvat2slowfast
+from tests.test_tracks_extractor import (
+    scene_width,
+    scene_height
+)
 from tests.utils import (
-    get_behavior,
     del_dir,
-    del_file
+    del_file,
+    dir_exists,
+    file_exists,
+    get_behavior
 )
 
 
@@ -50,6 +59,40 @@ def test_run(self):
                     "--classes", self.classes]
         run()
 
+        # check output dirs
+        self.assertTrue(dir_exists(self.dataset))
+        self.assertTrue(dir_exists(f"{self.dataset}/annotation"))
+        self.assertTrue(dir_exists(f"{self.dataset}/dataset/image"))
+        self.assertTrue(file_exists(f"{self.dataset}/annotation/classes.json"))
+        self.assertTrue(file_exists(f"{self.dataset}/annotation/data.csv"))
+
+        # check classes.json
+        with open(f"{self.dataset}/annotation/classes.json", "r", encoding="utf-8") as f:
+            classes = json.load(f)
+        with open(self.classes, "r", encoding="utf-8") as f:
+            ethogram = json.load(f)
+        self.assertEqual(classes, ethogram)
+
+        # check data.csv
+        with open(f"{self.dataset}/annotation/data.csv", "r", encoding="utf-8") as f:
+            df = pd.read_csv(f, sep=" ")
+
+        video_id = 1
+        for i, row in df.iterrows():
+            self.assertEqual(row["original_vido_id"], f"Z{video_id:04d}")
+            self.assertEqual(row["video_id"], video_id)
+            self.assertEqual(row["frame_id"], i+1)
+            self.assertEqual(row["path"], f"Z{video_id:04d}/{i+1}.jpg")
+            self.assertEqual(row["labels"], 1)
+        self.assertEqual(i, 90)
+
+        # check dataset
+        for i in range(1, 92):
+            data_im = f"{self.dataset}/dataset/image/Z{video_id:04d}/{i}.jpg"
+            self.assertTrue(file_exists(data_im))
+            data_im = cv2.imread(data_im)
+            self.assertEqual(data_im.shape, (scene_height, scene_width, 3))
+
     def test_parse_arg_min(self):
         # parse arguments
         sys.argv = [self.tool,
@@ -65,24 +108,21 @@ def test_parse_arg_min(self):
 
         # check default argument values
         self.assertEqual(args.old2new, None)
-
-        # run cvat2slowfast
-        run()
+        self.assertTrue(not args.no_images)
 
     def test_parse_arg_full(self):
         # parse arguments
         sys.argv = ["cvat2slowfast.py",
                     "--miniscene", self.miniscene,
                     "--dataset", self.dataset,
                     "--classes", self.classes,
-                    "--old2new", self.old2new]
+                    "--old2new", self.old2new,
+                    "--no_images"]
         args = cvat2slowfast.parse_args()
 
         # check parsed argument values
         self.assertEqual(args.miniscene, self.miniscene)
         self.assertEqual(args.dataset, self.dataset)
         self.assertEqual(args.classes, self.classes)
         self.assertEqual(args.old2new, self.old2new)
-
-        # run cvat2slowfast
-        run()
+        self.assertTrue(args.no_images)