pascal_finetune_dataset_save

import os
import numpy as np
import cv2
import pickle

data_dir = r"D:\dataset\rank_saliency_battle\PASCAL-S\train\train_gt"
img_dir = r"D:\dataset\rank_saliency_battle\PASCAL-S\train\train_imgs"
seg_maps_dir = r"D:\dataset\rank_saliency_battle\PASCAL-S\seg_masks"
gt_dir = r"D:\dataset\rank_saliency_battle\PASCAL-S\train\train_gt"

output_dir = r"D:\dataset\rank_saliency_battle\PASCAL-S"

def get_box(im):
    im_x = np.sum(im, axis=0)
    im_y = np.sum(im, axis=1)
    x1, y1, x2, y2 = 0, 0, len(im_x)-1, len(im_y)-1
    for i in range(len(im_x)):
        if im_x[i] == 0:
            x1 += 1
        else:
            break
    for i in range(len(im_x)-1, 0, -1):
        if im_x[i] == 0:
            x2 -= 1
        else:
            break
    for i in range(len(im_y)):
        if im_y[i] == 0:
            y1 += 1
        else:
            break
    for i in range(len(im_y)-1, 0, -1):
        if im_y[i] == 0:
            y2 -= 1
        else:
            break

    w = x2 - x1
    h = y2 - y1

    return x1, y1, w, h


def get_rank_box_seg(gt, gt_seg_map):
    img_w, img_h = gt.shape
    gt = gt.astype(np.float)
    gt_seg_map = gt_seg_map.astype(np.float)
    unique_mask = gt * (1 + gt_seg_map)
    instances = np.unique(unique_mask)[1:]
    np.random.shuffle(instances)

    gt_ranks = np.zeros(len(instances))
    boxes = np.zeros([len(instances), 4])
    segs = np.zeros([len(instances), img_w, img_h])
    for j, instance in enumerate(instances):
        instance_mask = unique_mask == instance
        rank = np.unique(gt[instance_mask])
        assert len(rank) == 1
        gt_ranks[j] = rank
        segs[j][instance_mask] = 1
        x, y, w, h = get_box(segs[j])
        boxes[j, :] = x, y, w, h

    return gt_ranks, boxes, segs


if __name__ == '__main__':
    files = os.listdir(data_dir)
    im_images = {}
    gt_boxes = {}
    gt_segs = {}
    gt_ranks = {}

    for j, file in enumerate(files):
        print(j, file)
        img = cv2.imread(os.path.join(img_dir, file[:-3] + 'jpg'))
        gt_seg_map = cv2.imread(os.path.join(seg_maps_dir, file[:-3] + 'png'))[:, :, 0]
        gt = cv2.imread(os.path.join(gt_dir, file[:-3] + 'png'))[:, :, 0]
        if img is None:
            print(file)
            print(data_dir)
            raise FileNotFoundError

        ranks, boxes, segs = get_rank_box_seg(gt, gt_seg_map)
        im_images[file[:-4]] = img
        gt_boxes[file[:-4]] = boxes
        gt_segs[file[:-4]] = segs
        gt_ranks[file[:-4]] = ranks

    data_set = {}
    data_set['train_imgs'] = im_images
    data_set['train_boxes'] = gt_boxes
    data_set['train_segs'] = gt_segs
    data_set['train_ranks'] = gt_ranks

    with open(os.path.join(output_dir, "pascal_dataset_finetune.pkl"), "wb") as f:
        pickle.dump(data_set, f)