def main():
cate_fol_rp = '/disk5/yangle/BasicDataset/dataset/DAVIS-Semantic/MergeFiles/DAVIS/JPEGImages/480p/'
cate_gt_rp = '/disk5/yangle/BasicDataset/dataset/DAVIS-Semantic/MergeFiles/DAVIS/Annotations/480p/'
sep_cate_rp = '/disk5/yangle/DAVIS/result/mask/try/all-cate-e2e_mask_rcnn_X-101-64x4d-FPN_1xdemo/'
res_rp = '/disk5/yangle/DAVIS/result/mask/try/sel/'
if not os.path.exists(res_rp):
os.makedirs(res_rp)
# load model
model = UNet_deeplab(in_channels=4, feature_length=512)
model = model.cuda()
weights_fpath = '/disk5/yangle/DAVIS/result/TrainNet/ranking-0518/weights/ranking-test-weights-10-0.071-0.000-0.022-0.000.pth'
state = torch.load(weights_fpath)
model.load_state_dict(state['state_dict'])
model.eval()
cate_set = os.listdir(sep_cate_rp)
cate_set.sort()
for cate_name in cate_set[3:4]:
# compare with the first annotated image
img_ref = Image.open(cate_fol_rp + cate_name + '/00000.jpg')
mask_begin = Image.open(cate_gt_rp + cate_name + '/00000.png')
# only dispose test data
gt_set = os.listdir(cate_gt_rp + cate_name + '/')
if len(gt_set) > 1:
continue
_, object_num = mask_begin.getextrema()
for iobj in range(1, object_num + 1):
# only consider the i th object
mask_ref = mask_begin.point(lambda i: 255 if i == iobj else 0)
feature_ref = net_prediction(img_ref, mask_ref, model)
img_fol_rp = sep_cate_rp + cate_name + '/'
img_fol_set = os.listdir(img_fol_rp)
img_fol_set.sort()
# error occurs here #
for img_ord, img_fol_name in enumerate(img_fol_set):
img = Image.open(cate_fol_rp + cate_name + '/' + img_fol_name +
'.jpg')
mask_rp = img_fol_rp + img_fol_name + '/'
mask_set = os.listdir(mask_rp)
mask_set.sort()
distance_min = 100
for mask_name in mask_set:
print(mask_rp + mask_name)
mask = Image.open(mask_rp + mask_name)
# skip too tiny object
if not utils.HasObject(mask):
continue
feature_c = net_prediction(img, mask, model)
distance = torch.dist(feature_c, feature_ref, p=2)
distance = distance.data[0]
print('distance between two vectors is %f' % distance)
if distance < distance_min:
distance_min = distance
mask_sel = mask
name_sel = mask_name
feature_sel = feature_c
mask_save_file = res_rp + cate_name + img_fol_name + name_sel[:-4] + str(
iobj) + '.png'
mask_sel.save(mask_save_file)
feature_ref = feature_sel
# update feature
order = img_ord + 2.0
feature_ref = (
order - 1) / order * feature_ref + 1 / order * feature_sel
def main():
cate_fol_rp = '/disk5/yangle/BasicDataset/dataset/DAVIS-Semantic/JPEGImages/480p/'
cate_gt_rp = '/disk5/yangle/BasicDataset/dataset/DAVIS-Semantic/Annotations/480p/'
sep_cate_rp = '/disk5/yangle/DAVIS/result/RankNet/e2e_mask_rcnn_X-101-64x4d-FPN_1xdemo/'
res_rp = '/disk5/yangle/DAVIS/result/ZD/RankingNet/e2e_mask_rcnn_X-101-64x4d-FPN_1xdemo/'
SEL_DIST_TH = 20
if not os.path.exists(res_rp):
os.makedirs(res_rp)
# load model
model = UNet_deeplab(in_channels=4, feature_length=512)
model = model.cuda()
weights_fpath = 'ranking-weights-62-ZD.pth'
state = torch.load(weights_fpath)
model.load_state_dict(state['state_dict'])
model.eval()
cate_set = os.listdir(sep_cate_rp)
cate_set.sort()
for cate_name in cate_set:
# compare with the first annotated image
img_ref = Image.open(cate_fol_rp + cate_name + '/00000.jpg')
mask_begin = Image.open(cate_gt_rp + cate_name + '/00000.png')
# only dispose test data
gt_set = os.listdir(cate_gt_rp + cate_name + '/')
if len(gt_set) > 1:
continue
_, object_num = mask_begin.getextrema()
for iobj in range(1, object_num + 1):
obj_str = str(iobj)
obj_str = obj_str.zfill(2)
# only consider the i th object
mask_ref = mask_begin.point(lambda i: 255 if i == iobj else 0)
# print(img_ref.mode)
feature_ref = net_prediction(img_ref, mask_ref, model)
img_fol_rp = sep_cate_rp + cate_name + '/'
img_fol_set = os.listdir(img_fol_rp)
img_fol_set.sort()
for img_ord, img_fol_name in enumerate(img_fol_set):
img = Image.open(cate_fol_rp + cate_name + '/' + img_fol_name +
'.jpg')
cols_img, rows_img = img.size
mask_rp = img_fol_rp + img_fol_name + '/'
mask_set = os.listdir(mask_rp)
mask_set.sort()
if len(mask_set) == 0:
print('encounter empty folder')
continue
mask_save = np.zeros((rows_img, cols_img), dtype=np.uint8)
for iord, mask_name in enumerate(mask_set):
print(mask_rp + mask_name)
mask = Image.open(mask_rp + mask_name)
# skip too tiny object
if not utils.HasObject(mask):
# print('skip tiny mask %s', mask_name)
mask_set[iord] = []
continue
feature_c = net_prediction(img, mask, model)
distance = torch.dist(feature_c, feature_ref, p=2)
distance = distance.data[0]
print('distance between two vectors is %f' % distance)
if distance < SEL_DIST_TH:
mask_np = np.asarray(mask, dtype=np.uint8)
mask_save = np.where(mask_np == 255, 255, mask_save)
mask_save_file = res_rp + cate_name + img_fol_name + '_' + obj_str + '.png'
mask_mer = Image.fromarray(mask_save, mode='L')
mask_mer.save(mask_save_file)
if not utils.HasObject(mask_mer):
continue
feature_sel = net_prediction(img, mask_mer, model)
order = img_ord + 2.0
feature_ref = (
order - 1) / order * feature_ref + 1 / order * feature_sel
def main():
cate_fol_rp = '/disk5/yangle/BasicDataset/dataset/DAVIS-Semantic/ZIP_files/DAVIS-cha/JPEGImages/480p/'
cate_gt_rp = '/disk5/yangle/BasicDataset/dataset/DAVIS-Semantic/ZIP_files/DAVIS-cha/Annotations/480p/'
sep_cate_rp = '/disk5/yangle/DAVIS/result/fuse/'
res_rp = '/disk5/yangle/DAVIS/result/fuse_sig/'
if not os.path.exists(res_rp):
os.makedirs(res_rp)
# load model
model = tiramisu.FCDenseNet57(in_channels=8, n_classes=N_CLASSES)
model = model.cuda()
# model = torch.nn.DataParallel(model).cuda()
weights_fpath = 'cat_sig-weights-36.pth'
state = torch.load(weights_fpath)
model.load_state_dict(state['state_dict'])
model.eval()
cate_set = os.listdir(sep_cate_rp)
cate_set.sort()
for cate_name in cate_set:
# skip test or dev
ann_rp = cate_gt_rp + cate_name
if not os.path.exists(ann_rp):
continue
mask_begin = Image.open(ann_rp + '/00000.png')
_, object_num = mask_begin.getextrema()
for iobj in range(1, object_num + 1):
# only consider the i th object
mask_former = mask_begin.point(lambda i: 255 if i == iobj else 0)
img_fol_rp = sep_cate_rp + cate_name + '/'
img_fol_set = os.listdir(img_fol_rp)
img_fol_set.sort()
# start from image 00001.jpg
for ord_fol, img_fol_name in enumerate(img_fol_set[1:]):
img = Image.open(cate_fol_rp + cate_name + '/' + img_fol_name +
'.jpg')
mask_rp = img_fol_rp + img_fol_name + '/'
mask_set = os.listdir(mask_rp)
mask_set.sort()
comask_set = []
for mask_name in mask_set:
# only consider the i th object
if iobj != int(mask_name[-6:-4]):
continue
comask = Image.open(mask_rp + mask_name)
# skip too tiny object
if not utils.HasObject(comask):
continue
comask_set.append(comask)
output = net_prediction(img, comask, mask_former, model)
pred_mask = output[4]
pred_mask = pred_mask.squeeze()
pred_mask = F.sigmoid(pred_mask)
pred_mask = pred_mask.data.cpu()
save_path = res_rp + cate_name + img_fol_name + '_' + mask_name
print(save_path)
torchvision.utils.save_image(pred_mask, save_path)
# collect former mask (as union) for next iteration
if len(comask_set) > 0:
mask_former = utils.collect_former_mask(comask_set)
def main():
cate_fol_rp = '/disk5/yangle/BasicDataset/dataset/DAVIS-Semantic/MergeFiles/DAVIS/JPEGImages/480p/'
cate_gt_rp = '/disk5/yangle/BasicDataset/dataset/DAVIS-Semantic/MergeFiles/DAVIS/Annotations/480p/'
sep_cate_rp = '/disk5/yangle/DAVIS/result/mask/step2_rearrange/'
res_rp = '/disk5/yangle/DAVIS/result/mask/step3_objectmask/'
if not os.path.exists(res_rp):
os.makedirs(res_rp)
# load model
model = UNet_deeplab(in_channels=4, feature_length=512)
model = model.cuda()
weights_fpath = 'ranking-weights-62-ZD.pth'
state = torch.load(weights_fpath)
model.load_state_dict(state['state_dict'])
model.eval()
cate_set = os.listdir(sep_cate_rp)
cate_set.sort()
for cate_name in cate_set:
# only dispose test data
gt_rp = cate_gt_rp + cate_name + '/'
if not os.path.exists(gt_rp):
continue
# compare with the first annotated image
img_ref = Image.open(cate_fol_rp + cate_name + '/00000.jpg')
mask_begin = Image.open(cate_gt_rp + cate_name + '/00000.png')
_, object_num = mask_begin.getextrema()
# srote the reference feature
obj_ref_fea = []
for iobj in range(1, object_num + 1):
# only consider the i th object
mask_ref = mask_begin.point(lambda i: 255 if i == iobj else 0)
feature_ref = net_prediction(img_ref, mask_ref, model)
obj_ref_fea.append(feature_ref)
img_fol_rp = sep_cate_rp + cate_name + '/'
img_fol_set = os.listdir(img_fol_rp)
img_fol_set.sort()
for img_ord, img_fol_name in enumerate(img_fol_set):
img = Image.open(cate_fol_rp + cate_name + '/' + img_fol_name +
'.jpg')
mask_rp = img_fol_rp + img_fol_name + '/'
mask_set = os.listdir(mask_rp)
mask_set.sort()
mask_num = len(mask_set)
# record distance for each mask
dist_doc = 100 * np.ones((object_num, mask_num))
for jobj in range(object_num):
feature_ref = obj_ref_fea[jobj]
for kmask, mask_name in enumerate(mask_set):
# print(mask_rp + mask_name)
mask = Image.open(mask_rp + mask_name)
# skip too tiny object
if not utils.HasObject(mask):
continue
feature_c = net_prediction(img, mask, model)
distance = torch.dist(feature_c, feature_ref, p=2)
dist_doc[jobj, kmask] = distance.data[0]
for jobj in range(object_num):
# minimul distance
if not dist_doc.size:
continue
row_min, col_min = np.unravel_index(
np.argmin(dist_doc, axis=None), dist_doc.shape)
str_obj = str(row_min + 1)
str_obj = str_obj.zfill(2)
mask_save_file = res_rp + cate_name + img_fol_name + '_' + str_obj + '.png'
print(mask_save_file)
shutil.copyfile(mask_rp + mask_set[col_min], mask_save_file)
# disable the column and the row
dist_doc[:, col_min] = 100
dist_doc[row_min, :] = 100
def main():
cate_fol_rp = '/disk5/yangle/BasicDataset/dataset/DAVIS-Semantic/MergeFiles/DAVIS/JPEGImages/480p/'
cate_gt_rp = '/disk5/yangle/BasicDataset/dataset/DAVIS-Semantic/MergeFiles/DAVIS/Annotations/480p/'
sep_cate_rp = '/disk5/yangle/DAVIS/result/mask/step2_rearrange/'
res_rp = '/disk5/yangle/DAVIS/result/mask/step3_objectmask/'
if not os.path.exists(res_rp):
os.makedirs(res_rp)
# load model
model = UNet_deeplab(in_channels=4, feature_length=512)
model = model.cuda()
weights_fpath = 'ranking-weights-62-ZD.pth'
state = torch.load(weights_fpath)
model.load_state_dict(state['state_dict'])
model.eval()
cate_set = os.listdir(sep_cate_rp)
cate_set.sort()
for cate_name in cate_set[54:]:
# only dispose test data
gt_rp = cate_gt_rp + cate_name + '/'
if not os.path.exists(gt_rp):
continue
# compare with the first annotated image
img_ref = Image.open(cate_fol_rp + cate_name + '/00000.jpg')
mask_begin = Image.open(cate_gt_rp + cate_name + '/00000.png')
_, object_num = mask_begin.getextrema()
for iobj in range(1, object_num + 1):
obj_str = str(iobj)
obj_str = obj_str.zfill(2)
post_fix = '_' + obj_str + '.png'
# only consider the i th object
mask_ref = mask_begin.point(lambda i: 255 if i == iobj else 0)
feature_ref = net_prediction(img_ref, mask_ref, model)
img_fol_rp = sep_cate_rp + cate_name + '/'
img_fol_set = os.listdir(img_fol_rp)
img_fol_set.sort()
for img_ord, img_fol_name in enumerate(img_fol_set):
img = Image.open(cate_fol_rp + cate_name + '/' + img_fol_name +
'.jpg')
mask_rp = img_fol_rp + img_fol_name + '/'
# select mask_set, only dispose current object
mask_set = os.listdir(mask_rp)
mask_set_s = [f for f in mask_set if f.endswith(post_fix)]
# this object contains no masks
if len(mask_set_s) == 0:
continue
mask_set_s.sort()
distance_min = 100
for mask_name in mask_set_s:
print(mask_rp + mask_name)
mask = Image.open(mask_rp + mask_name)
# skip too tiny object
if not utils.HasObject(mask):
continue
feature_c = net_prediction(img, mask, model)
distance = torch.dist(feature_c, feature_ref, p=2)
distance = distance.data[0]
print('distance between two vectors is %f' % distance)
if distance < distance_min:
distance_min = distance
mask_sel = mask
feature_sel = feature_c
# do not select the object
if distance_min == 100:
continue
save_rp = res_rp + cate_name + '/' + img_fol_name + '/'
if not os.path.exists(save_rp):
os.makedirs(save_rp)
mask_save_file = save_rp + obj_str + '.png'
mask_sel.save(mask_save_file)
def main():
cate_fol_rp = '/disk5/yangle/BasicDataset/dataset/DAVIS-Semantic/JPEGImages/480p/'
cate_gt_rp = '/disk5/yangle/BasicDataset/dataset/DAVIS-Semantic/Annotations/480p/'
sep_cate_rp = '/disk5/yangle/DAVIS/result/RankNet/e2e_mask_rcnn_X-101-64x4d-FPN_1xdemo/'
res_rp = '/disk5/yangle/DAVIS/result/ZD/RankingNet/e2e_mask_rcnn_X-101-64x4d-FPN_1xdemo/'
if not os.path.exists(res_rp):
os.makedirs(res_rp)
# load model
model = UNet_deeplab(in_channels=4, feature_length=512)
model = model.cuda()
weights_fpath = '/disk5/zhangdong/DAVIS/result/TrainNet2/ranking-test/weights/ranking-test-weights-62-0.135-0.000-0.384-0.000.pth'
# weights_fpath = '/disk5/yangle/DAVIS/result/TrainNet/ranking-0518/weights/ranking-test-weights-10-0.071-0.000-0.022-0.000.pth'
state = torch.load(weights_fpath)
model.load_state_dict(state['state_dict'])
model.eval()
distance_set = []
ratio_set = []
cate_set = os.listdir(sep_cate_rp)
cate_set.sort()
for cate_name in cate_set:
# only dispose train val data
gt_set = os.listdir(cate_gt_rp + cate_name + '/')
if len(gt_set) == 1:
continue
img_fol_rp = sep_cate_rp + cate_name + '/'
img_fol_set = os.listdir(img_fol_rp)
img_fol_set.sort()
for img_ord, img_fol_name in enumerate(img_fol_set):
# sample every ten frames
if img_ord % 10 != 0:
continue
img = Image.open(cate_fol_rp + cate_name + '/' + img_fol_name +
'.jpg')
gt = Image.open(cate_gt_rp + cate_name + '/' + img_fol_name +
'.png')
_, object_num = gt.getextrema()
for iobj in range(1, object_num + 1):
obj_str = str(iobj)
obj_str = obj_str.zfill(2)
# only consider the i th object
mask_ref = gt.point(lambda i: 255 if i == iobj else 0)
if not utils.HasObject(mask_ref):
continue
feature_ref = net_prediction(img, mask_ref, model)
mask_rp = img_fol_rp + img_fol_name + '/'
mask_set = os.listdir(mask_rp)
mask_set.sort()
if len(mask_set) == 0:
print('encounter empty folder')
continue
for iord, mask_name in enumerate(mask_set):
print(mask_rp + mask_name)
mask = Image.open(mask_rp + mask_name)
# skip too tiny object
if not utils.HasObject(mask):
# print('skip tiny mask %s', mask_name)
mask_set[iord] = []
continue
# feature vector diatance
feature_c = net_prediction(img, mask, model)
distance = torch.dist(feature_c, feature_ref, p=2)
distance = distance.data[0]
distance_set.append(distance)
# segmentation ratio
mask_ref_np = np.asarray(mask_ref, dtype=np.uint8)
ground_truth_bool = mask_ref_np == 255
mask_np = np.asarray(mask, dtype=np.uint8)
mask_bool = mask_np == 255
obj_value = np.sum(
(mask_bool & ground_truth_bool)) / np.sum(
mask_bool, dtype=np.float32)
ratio_set.append(obj_value)
out = open('data_zd.pkl', 'wb')
pickle.dump([distance_set, ratio_set], out)
out.close()