merge_SocAndSalicon

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)