Add `allow_background_images=True` in split_dota.py (#15037)

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
pull/15050/head^2
Galasnow 4 months ago committed by GitHub
parent 54a0494e2d
commit 1523fa12d8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
  1. 27
      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
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")
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")
def split_images_and_labels(data_root, save_dir, split="train", crop_sizes=(1024,), gaps=(200,)):

Loading…
Cancel
Save