File size: 1,497 Bytes
2e7bc51 |
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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
import cv2
from os.path import join
import os
import numpy as np
from tqdm import tqdm
dir = './datasets/Open_Images/'
mode = 'test'
image_dir = join(dir, 'images', mode)
seg_dir = join(dir, 'segs', mode)
files = os.listdir(seg_dir)
data_dict = {}
for file in tqdm(files):
seg_path = join(seg_dir, file)
image_path = join(image_dir, file.split('_')[0] + '.jpg')
seg = cv2.imread(seg_path)
image = cv2.imread(image_path)
seg = cv2.resize(seg, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
seg = seg[:, :, 0]
# obtain contours point set： contours
contours = cv2.findContours(seg, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
contours = contours[0] if len(contours) == 2 else contours[1]
if len(contours) > 1:
cntr = np.vstack(contours)
elif len(contours) == 1:
cntr = contours[0]
else:
continue
if len(cntr) < 2:
continue
hs, he = np.min(cntr[:, :, 1]), np.max(cntr[:, :, 1])
ws, we = np.min(cntr[:, :, 0]), np.max(cntr[:, :, 0])
h, w = seg.shape
if (he - hs) % 2 == 1 and (he + 1) <= h:
he = he + 1
if (he - hs) % 2 == 1 and (hs - 1) >= 0:
hs = hs - 1
if (we - ws) % 2 == 1 and (we + 1) <= w:
we = we + 1
if (we - ws) % 2 == 1 and (ws - 1) >= 0:
ws = ws - 1
if he - hs < 2 or we - ws < 2:
continue
data_dict[file] = [cntr, hs, he, ws, we]
np.save(join(dir, 'segs', f'{mode}_bbox_dict.npy'), data_dict) |