HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

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

Browse source 
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Galasnow 2024-08-09 04:33:04 +08:00 • committed by GitHub
parent 54a0494e2d
commit 1523fa12d8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 13 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

25
ultralytics/data/split_dota.py
View file

@ -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,)):