create new folders in dlc labeled-data when they do not exist

Author: darsnack

packages/cheese3d/cheese3d/backends/dlc.py

@@ -143,60 +143,61 @@ def fix_video_symlinks(self):
     def import_c3d_labels(self, videos: Dict[str, Path]):
         for name, path in videos.items():
             label_folder = self.project_path / "labeled-data" / name
-            if label_folder.exists():
-                images = reglob(r".*\.png", str(path))
-                for image in images:
-                    src_image = Path(image)
-                    dst_image = label_folder / src_image.name
-                    if dst_image.exists():
-                        os.remove(dst_image)
-                    relpath = os.path.relpath(src_image, label_folder)
-                    os.symlink(relpath, dst_image)
-                annotations_yaml = path / "annotations.yaml"
-                hdf = label_folder / f"CollectedData_{self.experimenter}.h5"
-                csv = label_folder / f"CollectedData_{self.experimenter}.csv"
-                if annotations_yaml.exists():
-                    with open(annotations_yaml, "r") as f:
-                        annotations = yaml.safe_load(f)
-                    data_dict = {}
-                    for kp, files in annotations.items():
-                        for file, coords in files.items():
-                            # create index for this row
-                            idx = ('labeled-data', str(label_folder.name), file)
-                            # get x and y coordinates (could be null/None)
-                            x_coord = coords[0][0]
-                            y_coord = coords[0][1]
-                            # create column keys for x and y
-                            x_col = (self.experimenter, kp, 'x')
-                            y_col = (self.experimenter, kp, 'y')
-                            # store in data dictionary
-                            if idx not in data_dict:
-                                data_dict[idx] = {}
-                            data_dict[idx][x_col] = x_coord
-                            data_dict[idx][y_col] = y_coord
-                    # convert to dataframe
-                    index = pd.MultiIndex.from_tuples(list(data_dict.keys()))
-                    columns = pd.MultiIndex.from_tuples(
-                        [(self.experimenter, bp, coord)
-                         for bp in annotations.keys()
-                         for coord in ['x', 'y']],
-                        names=["scorer", "bodyparts", "coords"]
-                    )
-                    # create empty df with the right structure
-                    df = pd.DataFrame(index=index, columns=columns)
-                    # fill in the values
-                    for idx, values in data_dict.items():
-                        for col, val in values.items():
-                            df.loc[idx, col] = val
-                    # convert to float (this will convert None/null to NaN)
-                    df = df.astype(float)
-                    # write dataframe to disk
-                    if hdf.exists():
-                        os.remove(hdf)
-                    df.to_hdf(hdf, key="df", mode="w")
-                    if csv.exists():
-                        os.remove(csv)
-                    df.to_csv(csv)
+            if not label_folder.exists():
+                label_folder.mkdir(exist_ok=True, parents=True)
+            images = reglob(r".*\.png", str(path))
+            for image in images:
+                src_image = Path(image)
+                dst_image = label_folder / src_image.name
+                if dst_image.exists():
+                    os.remove(dst_image)
+                relpath = os.path.relpath(src_image, label_folder)
+                os.symlink(relpath, dst_image)
+            annotations_yaml = path / "annotations.yaml"
+            hdf = label_folder / f"CollectedData_{self.experimenter}.h5"
+            csv = label_folder / f"CollectedData_{self.experimenter}.csv"
+            if annotations_yaml.exists():
+                with open(annotations_yaml, "r") as f:
+                    annotations = yaml.safe_load(f)
+                data_dict = {}
+                for kp, files in annotations.items():
+                    for file, coords in files.items():
+                        # create index for this row
+                        idx = ('labeled-data', str(label_folder.name), file)
+                        # get x and y coordinates (could be null/None)
+                        x_coord = coords[0][0]
+                        y_coord = coords[0][1]
+                        # create column keys for x and y
+                        x_col = (self.experimenter, kp, 'x')
+                        y_col = (self.experimenter, kp, 'y')
+                        # store in data dictionary
+                        if idx not in data_dict:
+                            data_dict[idx] = {}
+                        data_dict[idx][x_col] = x_coord
+                        data_dict[idx][y_col] = y_coord
+                # convert to dataframe
+                index = pd.MultiIndex.from_tuples(list(data_dict.keys()))
+                columns = pd.MultiIndex.from_tuples(
+                    [(self.experimenter, bp, coord)
+                     for bp in annotations.keys()
+                     for coord in ['x', 'y']],
+                    names=["scorer", "bodyparts", "coords"]
+                )
+                # create empty df with the right structure
+                df = pd.DataFrame(index=index, columns=columns)
+                # fill in the values
+                for idx, values in data_dict.items():
+                    for col, val in values.items():
+                        df.loc[idx, col] = val
+                # convert to float (this will convert None/null to NaN)
+                df = df.astype(float)
+                # write dataframe to disk
+                if hdf.exists():
+                    os.remove(hdf)
+                df.to_hdf(hdf, key="df", mode="w")
+                if csv.exists():
+                    os.remove(csv)
+                df.to_csv(csv)
 
     def export_c3d_labels(self, videos: Dict[str, Path]):
         for name, path in videos.items():