def get_ROI_position(img, net, opt, keepsize=True):
mask = run_segment(img, net, size=360, gpu_id=opt.gpu_id)
mask = impro.mask_threshold(mask, opt.mask_extend, opt.mask_threshold)
if keepsize:
mask = impro.resize_like(mask, img)
x, y, halfsize, area = impro.boundingSquare(mask, 1)
return mask, x, y, halfsize, area
mask_path = os.path.join(output_dir, 'mask')
util.makedirs(mask_path)
mask_names = os.listdir(mask_dir)
img_names = os.listdir(img_dir)
mask_names.sort()
img_names.sort()
print('Find images:', len(img_names))
cnt = 0
for fold in range(FOLD_NUM):
for img_name, mask_name in zip(img_names, mask_names):
try:
img = impro.imread(os.path.join(img_dir, img_name))
mask = impro.imread(os.path.join(mask_dir, mask_name), 'gray')
mask = impro.resize_like(mask, img)
x, y, size, area = impro.boundingSquare(mask, 1.5)
if area > 100:
if Bounding:
img = impro.resize(
img[y - size:y + size, x - size:x + size], OUT_SIZE)
mask = impro.resize(
mask[y - size:y + size, x - size:x + size], OUT_SIZE)
img_mosaic = mosaic.addmosaic_random(img, mask)
if HD:
cv2.imwrite(
os.path.join(train_A_path, '%05d' % cnt + '.jpg'),
img_mosaic)
cv2.imwrite(
os.path.join(train_B_path, '%05d' % cnt + '.jpg'), img)
if 'network' in opt.mod:
net = loadmodel.bisenet(opt, 'roi')
print('Find images:', len(imgpaths))
starttime = datetime.datetime.now()
filecnt = 0
savecnt = opt.start
for fold in range(opt.fold):
for i in range(len(imgpaths)):
filecnt += 1
try:
# load image and get mask
img = impro.imread(imgpaths[i])
if 'drawn' in opt.mod:
mask_drawn = impro.imread(maskpaths[i], 'gray')
mask_drawn = impro.resize_like(mask_drawn, img)
mask = mask_drawn
if 'irregular' in opt.mod:
mask_irr = impro.imread(irrpaths[random.randint(0, 12000 - 1)],
'gray')
mask_irr = data.random_transform_single(
mask_irr, (img.shape[0], img.shape[1]))
mask = mask_irr
if 'network' in opt.mod:
mask_net = runmodel.get_ROI_position(img,
net,
opt,
keepsize=True)[0]
if opt.use_gpu != -1:
torch.cuda.empty_cache()
if not opt.all_mosaic_area: