diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..682d61e --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +datumaro diff --git a/rotate_dataset.py b/rotate_dataset.py new file mode 100644 index 0000000..e1d97d0 --- /dev/null +++ b/rotate_dataset.py @@ -0,0 +1,89 @@ +import argparse +from copy import deepcopy +import datumaro as dm +from datumaro.components.transformer import ItemTransform +from datumaro.components.media import Image +from datumaro.components.annotation import Bbox +import numpy as np + + +class RotateTransform(ItemTransform): + def __init__(self, extractor, angle): + super().__init__(extractor) + self.angle = angle + # k is number of 90-degree CCW rotations + self.k = (angle // 90) % 4 + + def transform_item(self, item): + if not item.media or not isinstance(item.media, Image): + return item + + h, w = item.media.data.shape[:2] + # 1. Rotate Image + rotated_image = np.rot90(item.media.data, k=-self.k) # CCW to CW adjustment + + # 2. Rotate Bounding Boxes + new_annotations = [] + for ann in item.annotations: + if isinstance(ann, Bbox): + x, y, bw, bh = ann.points + # Calculate new coordinates based on rotation angle + if self.angle == 90: + new_bbox = [h - (y + bh), x, bh, bw] + elif self.angle == 180: + new_bbox = [w - (x + bw), h - (y + bh), bw, bh] + elif self.angle == 270: + new_bbox = [y, w - (x + bw), bh, bw] + else: + new_bbox = [x, y, bw, bh] + + new_annotations.append(ann.wrap(points=new_bbox)) + else: + new_annotations.append(ann) + + print(f"{item.id}_r{self.angle}") + return item.wrap( + id=f"{item.id}_r{self.angle}", + media=Image.from_numpy(rotated_image), + annotations=new_annotations, + ) + + +def main(): + parser = argparse.ArgumentParser(description="Rotate Datumaro datasets.") + parser.add_argument("--input", required=True, help="Path to input dataset") + parser.add_argument("--output", default="output_yolo", help="Output directory") + parser.add_argument( + "--angles", + nargs="+", + type=int, + default=[90, 180, 270], + help="Space separated angles (e.g. 90 180)", + ) + args = parser.parse_args() + + # 1. Load Original + dataset = dm.Dataset.import_from(args.input, "coco") + cloned_dataset = deepcopy(dataset) + + # 2. Generate list of rotated datasets from args + rotated_sets = [] + for angle in args.angles: + ds_rt = cloned_dataset.transform(RotateTransform, angle=angle) + rotated_sets.append(ds_rt) + cloned_dataset = deepcopy(dataset) + + # rotated_sets = [dataset.transform(RotateTransform, angle=a) for a in args.angles] + + # 3. Combine using splat operator on the array + combined_ds = dm.Dataset.from_extractors(dataset, *rotated_sets) + # combined_ds.transform("reindex", start=0) + + # 4. Export to Ultralytics YOLO + # combined_ds.export(args.output, format="coco", save_media=True) + combined_ds.export(args.output, format="coco", save_media=True) + print(f"✅ Success! Augmented dataset saved to: {args.output}") + + +if __name__ == "__main__": + main()