prepare_dataset(get_bbox...)

import os
import numpy as np
import cv2
import pickle

pascal_data_dir = '/data/lilong/rank_saliency/dataset/pascal_s/selected_imgs/1'
pascal_img_dir = '/data/lilong/rank_saliency/dataset/pascal_s/renamed_imgs'
pascal_img_gt_dir = '/data/lilong/rank_saliency/dataset/pascal_s/renamed_gt'

soc_data_dir = '/data/lilong/rank_saliency/dataset/soc_rank_saliency/train_val/gt'
soc_img_dir = '/data/lilong/rank_saliency/dataset/soc_rank_saliency/train_val/imgs'
soc_img_gt_dir = '/data/lilong/rank_saliency/dataset/soc_rank_saliency/train_val/gt'

output_dir = '/data/lilong/rank_saliency/dataset'


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 prepare_data(data_dir, img_dir, img_gt_dir):

    im_images = {}
    gt_boxes = {}
    gt_segs = {}
    gt_ranks = {}

    files = os.listdir(data_dir)
    for i, file in enumerate(files):
        print(i, file)
        img = cv2.imread(os.path.join(img_dir, file[:-3] + 'jpg'))
        img_gt = cv2.imread(os.path.join(img_gt_dir, file[:-3] + 'png'))[:, :, 0]

        ranks = np.unique(img_gt)[1:]
        boxes = np.zeros([len(ranks), 4])
        segs = np.zeros([len(ranks), img_gt.shape[0], img_gt.shape[1]])

        for j, rank in enumerate(ranks):
            segs[j][img_gt == rank] = 1
            x, y, w, h = get_box(segs[j])

            boxes[j, 0] = x
            boxes[j, 1] = y
            boxes[j, 2] = w
            boxes[j, 3] = h

        im_images[file[:-4]] = img
        gt_boxes[file[:-4]] = boxes
        gt_segs[file[:-4]] = segs
        gt_ranks[file[:-4]] = ranks

    return im_images, gt_boxes, gt_segs, gt_ranks


if __name__ == '__main__':

    im_images, gt_boxes, gt_segs, gt_ranks = prepare_data(soc_data_dir, soc_img_dir, soc_img_gt_dir)
    im_images1, gt_boxes1, gt_segs1, gt_ranks1 = prepare_data(pascal_data_dir, pascal_img_dir, pascal_img_gt_dir)
    im_images.update(im_images1)
    gt_boxes.update(gt_boxes1)
    gt_segs.update(gt_segs1)
    gt_ranks.update(gt_ranks1)

    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, 'dataset_pascal+soc_train.pkl'), 'wb') as f:
        pickle.dump(data_set, f)