diff --git a/pylabel/exporter.py b/pylabel/exporter.py index 887af7e..cb96a06 100644 --- a/pylabel/exporter.py +++ b/pylabel/exporter.py @@ -608,14 +608,10 @@ def ExportToYoloV5( unique_images = yolo_dataset["img_filename"].unique() output_file_paths = [] pbar = tqdm(desc="Exporting YOLO files...", total=len(unique_images)) - for img_filename in unique_images: - df_single_img_annots = yolo_dataset.loc[ - yolo_dataset.img_filename == img_filename - ] - + for img_filename, df_single_img_annots in yolo_dataset.groupby("img_filename"): basename, _ = os.path.splitext(img_filename) annot_txt_file = basename + ".txt" - # Use the value of the split collumn to create a directory + # Use the value of the split column to create a directory # The values should be train, val, test or '' if use_splits: split_dir = df_single_img_annots.iloc[0].split @@ -628,7 +624,7 @@ def ExportToYoloV5( ).mkdir(parents=True, exist_ok=True) # If segmentation = false then output bounding boxes - if segmentation == False: + if not segmentation: columns = [ "cat_id", "center_x_scaled",