import os
import cv2
import numpy as np
import shutil
soc_root_dir = '/disk3/lilong/salient_instance/SOC6K/'
#Sets = ['ValSet', 'TrainSet', 'TestSet'] # 600, 3600, 600
#Sets = ['ValSet', 'TrainSet']
Sets = ['TestSet']
soc_all_Imgs_dir = [os.path.join(soc_root_dir, '{}/Instance_gt/gt'.format(Set)) for Set in Sets]
fixationDir = '/disk3/lilong/coco_2014/fixation_maps_org'
imgDir = '/disk3/lilong/salient_instance/SOC6K'
#outputDir = '/disk3/lilong/soc_rank_saliency/train_val'
outputDir = '/disk3/lilong/soc_rank_saliency/test'
if not os.path.exists(outputDir):
os.mkdir(outputDir)
fix_lists = os.listdir(fixationDir)
soc_img_list = [[os.path.join(soc_imgs_dir, soc_img_name) for soc_img_name in os.listdir(soc_imgs_dir)]
for soc_imgs_dir in soc_all_Imgs_dir]
soc_img_list = [img for img_list in soc_img_list for img in img_list]
for soc_img_dir in soc_img_list:
soc_img_name = soc_img_dir.split('/')[-1]
if soc_img_name.split('_')[0] == 'COCO':
fix_img_path = os.path.join(fixationDir, '{}.png'.format(soc_img_name[:-4]))
fix_img = cv2.imread(fix_img_path)[:, :, 0]
assert fix_img is not None
labelMap = np.zeros(fix_img.shape)
soc_img = cv2.imread(soc_img_dir)
if soc_img.sum() == 0:
continue
soc_img[soc_img != 255] = 0
soc_img[soc_img == 255] = 1
soc_img_schannel = soc_img[:, :, 0]*4 + soc_img[:, :, 1]*2 + soc_img[:, :, 2]
instances = np.unique(soc_img_schannel)[1:].tolist()
for instance in instances:
labelMask = soc_img_schannel == instance
fix_density = (fix_img[labelMask] == 255).sum()
soc_img_schannel[labelMask] = fix_density
instances_fixDensity = np.unique(soc_img_schannel)[1:].tolist()
if len(instances_fixDensity) == 1:
continue
for i in range(0, len(instances_fixDensity)):
rank = int(255./len(instances_fixDensity)*(i+1))
labelMask = soc_img_schannel == instances_fixDensity[i]
soc_img_schannel[labelMask] = rank
out_dir = os.path.join(outputDir, 'gt/{}.png'.format(soc_img_name[:-4]))
cv2.imwrite(out_dir, soc_img_schannel)
img_dir = os.path.join(imgDir, '{}/Imgs/{}.jpg'.format(soc_img_dir.split('/')[-4],
soc_img_dir.split('/')[-1].split('.')[0]))
img_out_dir = os.path.join(outputDir, 'imgs/{}.jpg'.format(soc_img_name.split('.')[0]))
shutil.copyfile(img_dir, img_out_dir)