From 1523fa12d801df237a76a4ac1cdf59823e16a994 Mon Sep 17 00:00:00 2001 From: Galasnow <854932917@qq.com> Date: Fri, 9 Aug 2024 04:33:04 +0800 Subject: [PATCH] Add `allow_background_images=True` in split_dota.py (#15037) Co-authored-by: Glenn Jocher --- ultralytics/data/split_dota.py | 25 +++++++++++++------------ 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/ultralytics/data/split_dota.py b/ultralytics/data/split_dota.py index f0a1630c..435de769 100644 --- a/ultralytics/data/split_dota.py +++ b/ultralytics/data/split_dota.py @@ -144,7 +144,7 @@ def get_window_obj(anno, windows, iof_thr=0.7): return [np.zeros((0, 9), dtype=np.float32) for _ in range(len(windows))] # window_anns -def crop_and_save(anno, windows, window_objs, im_dir, lb_dir): +def crop_and_save(anno, windows, window_objs, im_dir, lb_dir, allow_background_images=True): """ Crop images and save new labels. @@ -154,6 +154,7 @@ def crop_and_save(anno, windows, window_objs, im_dir, lb_dir): window_objs (list): A list of labels inside each window. im_dir (str): The output directory path of images. lb_dir (str): The output directory path of labels. + allow_background_images (bool): Whether to include background images without labels. Notes: The directory structure assumed for the DOTA dataset: @@ -173,19 +174,19 @@ def crop_and_save(anno, windows, window_objs, im_dir, lb_dir): patch_im = im[y_start:y_stop, x_start:x_stop] ph, pw = patch_im.shape[:2] - cv2.imwrite(str(Path(im_dir) / f"{new_name}.jpg"), patch_im) label = window_objs[i] - if len(label) == 0: - continue - label[:, 1::2] -= x_start - label[:, 2::2] -= y_start - label[:, 1::2] /= pw - label[:, 2::2] /= ph + if len(label) or allow_background_images: + cv2.imwrite(str(Path(im_dir) / f"{new_name}.jpg"), patch_im) + if len(label): + label[:, 1::2] -= x_start + label[:, 2::2] -= y_start + label[:, 1::2] /= pw + label[:, 2::2] /= ph - with open(Path(lb_dir) / f"{new_name}.txt", "w") as f: - for lb in label: - formatted_coords = ["{:.6g}".format(coord) for coord in lb[1:]] - f.write(f"{int(lb[0])} {' '.join(formatted_coords)}\n") + with open(Path(lb_dir) / f"{new_name}.txt", "w") as f: + for lb in label: + formatted_coords = ["{:.6g}".format(coord) for coord in lb[1:]] + f.write(f"{int(lb[0])} {' '.join(formatted_coords)}\n") def split_images_and_labels(data_root, save_dir, split="train", crop_sizes=(1024,), gaps=(200,)):