First Commit
This commit is contained in:
1
requirements.txt
Normal file
1
requirements.txt
Normal file
@@ -0,0 +1 @@
|
||||
datumaro
|
||||
89
rotate_dataset.py
Normal file
89
rotate_dataset.py
Normal file
@@ -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()
|
||||
Reference in New Issue
Block a user