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/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
type='float',
help='learning rate')
parser.add_option('-g',
'--gpu',
action='store_true',
dest='gpu',
default=True,
help='use cuda')
# parser.add_option('-c', '--load', dest='load',
# default=False, help='load file model')
(options, args) = parser.parse_args()
deeplab_caffe2pytorch = 'train_iter_20000.caffemodel.pth'
# vgg16_model = 'vgg16-397923af.pth'
net = UNet_deeplab(3, 1)
net.cuda()
model = torch.load(deeplab_caffe2pytorch)
# vgg16_model = torch.load(vgg16_model)
print('load model:', deeplab_caffe2pytorch)
net = net.init_parameters(model)
# net = net.init_parameters_vgg16(vgg16_model)
# print(net.conv1[0].bias.data)
load_model = 'model_epoch_loss/MODEL_EPOCH10_LOSS1.0056836081342109.pth'
# print(net.conv1[0].bias.data)
load = False
if load:
net.load_state_dict(torch.load(load_model))
print('Model loaded from {}'.format(load_model))
def main():
cudnn.benchmark = True
deeplab_caffe2pytorch = 'train_iter_20000.caffemodel.pth'
print('load model:', deeplab_caffe2pytorch)
pretrained_model = torch.load(deeplab_caffe2pytorch)
model = UNet_deeplab(in_channels=4, feature_length=512)
model = model.init_parameters(pretrained_model)
# seperate layers, to set different lr
param_exist = []
param_add = []
for k, (name, module) in enumerate(model.named_children()):
# existing layers including: conv1~conv5, fc6, fc7
if k < 7:
for param in module.parameters():
param_exist.append(param)
# adding layers including: fc7_1
else:
for param in module.parameters():
param_add.append(param)
model = model.cuda()
print(' + Number of params: {}'.format(
sum([p.data.nelement() for p in model.parameters()])))
optimizer = optim.RMSprop([{'params': param_exist, 'lr': LEARNING_RATE*0.1},
{'params': param_add}], lr=LEARNING_RATE,
weight_decay=WEIGHT_DECAY, eps=1e-12)
# use margin=2
criterion = nn.TripletMarginLoss(margin=2, p=2).cuda()
exp_dir = EXPERIMENT + 'ranking-test'
if os.path.exists(exp_dir):
shutil.rmtree(exp_dir)
exp = experiment.Experiment('ranking-test', EXPERIMENT)
exp.init()
START_EPOCH = exp.epoch
END_EPOCH = START_EPOCH + N_EPOCHS
for epoch in range(START_EPOCH, END_EPOCH):
since = time.time()
# # ### Collect data ###
# # delete existing folder and old data
if os.path.exists(res_root_path):
shutil.rmtree(res_root_path)
utils.collect_data(ori_train_base_rp, res_train_base_rp)
utils.collect_data(ori_val_base_rp, res_val_base_rp)
# data loader
train_loader, val_loader = utils.data_loader(res_root_path)
# # ### Train ###
trn_loss = utils.train(model, train_loader, optimizer, criterion, epoch)
trn_err = 0
print('Epoch {:d}: Train - Loss: {:.4f}\tErr: {:.4f}'.format(epoch, trn_loss, trn_err))
time_elapsed = time.time() - since
print('Train Time {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
### Test ###
val_loss = utils.test(model, val_loader, criterion, epoch)
val_err = 0
print('Val - Loss: {:.4f}, Error: {:.4f}'.format(val_loss, val_err))
time_elapsed = time.time() - since
print('Total Time {:.0f}m {:.0f}s\n'.format(
time_elapsed // 60, time_elapsed % 60))
### Save Metrics ###
exp.save_history('train', trn_loss, trn_err)
exp.save_history('val', val_loss, val_err)
### Checkpoint ###
exp.save_weights(model, trn_loss, val_loss, trn_err, val_err)
exp.save_optimizer(optimizer, val_loss)
## Early Stopping ##
if (epoch - exp.best_val_loss_epoch) > MAX_PATIENCE:
print(("Early stopping at epoch %d since no "
+"better loss found since epoch %.3").format(epoch, exp.best_val_loss))
break
# Adjust Lr ###--old method
utils.adjust_learning_rate(LEARNING_RATE, LR_DECAY, optimizer,
epoch, DECAY_LR_EVERY_N_EPOCHS)
exp.epoch += 1
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()