make_bg.py
904 Bytes
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
import albumentations as A
import cv2
import os
from tqdm import tqdm
transform = A.Compose([
A.RandomResizedCrop(height=400, width=400, scale=(0.1, 0.2)),
A.HorizontalFlip(p=0.5),
A.VerticalFlip(p=0.5),
])
root = '/data1/lxl/data/ocr/doc_remove_stamp/'
save_path = '/data1/lxl/data/ocr/crop_bg/'
date_dirs = sorted(os.listdir(root))
img_list = list()
img_names = sorted(os.listdir(root))
for img_name in img_names:
img_path = os.path.join(root, img_name)
img_list.append(img_path)
print(f'src img number: {len(img_list)}')
cnt = 0
for crt_img in tqdm(img_list):
try:
img = cv2.imread(crt_img)
for _ in range(5):
transformed = transform(image=img)
transformed_img = transformed['image']
cv2.imwrite(os.path.join(save_path, f'{cnt :06d}.jpg'), transformed_img)
cnt += 1
except Exception:
continue