Spaces:
Running
Running
File size: 1,406 Bytes
033bd8b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
import cv2
import shutil
from tqdm import tqdm
from pathlib import Path
max_size = 1024
input_dataset_path = r'.\iHarmony4\HAdobe5k'
output_path = f'{input_dataset_path}_resized{max_size}'
input_dataset_path = Path(input_dataset_path)
output_path = Path(output_path)
assert not output_path.exists()
output_path.mkdir()
for subfolder in ['composite_images', 'masks', 'real_images']:
(output_path / subfolder).mkdir()
for annotation_path in input_dataset_path.glob('*.txt'):
shutil.copy(annotation_path, output_path / annotation_path.name)
images_list = sorted(input_dataset_path.rglob('*.jpg'))
images_list.extend(sorted(input_dataset_path.rglob('*.png')))
for x in tqdm(images_list):
image = cv2.imread(str(x), cv2.IMREAD_UNCHANGED)
new_path = output_path / x.relative_to(input_dataset_path)
if max(image.shape[:2]) <= max_size:
shutil.copy(x, new_path)
continue
new_width = max_size
new_height = max_size
scale = max_size / max(image.shape[:2])
if image.shape[0] > image.shape[1]:
new_width = int(round(scale * image.shape[1]))
else:
new_height = int(round(scale * image.shape[0]))
image = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_LANCZOS4)
if x.suffix == '.jpg':
cv2.imwrite(str(new_path), image, [cv2.IMWRITE_JPEG_QUALITY, 90])
else:
cv2.imwrite(str(new_path), image) |