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)