def update_appear():
global pred_prob
global appear
global location
for th in range(appear.shape[0]):
bbox = gen_bbox(prob_to_label(combine_prob(pred_prob[th])),
range(instance_num))
for i in range(appear.shape[1]):
appear[th, i] = ((bbox[i, 2] - bbox[i, 0]) *
(bbox[i, 3] - bbox[i, 1]) > 1)
if appear[th, i] > 0:
location[th, i, 0] = float(bbox[i, 2] + bbox[i, 0]) / 2
location[th, i, 1] = float(bbox[i, 3] + bbox[i, 1]) / 2
else:
location[th, i, :] = location[th - 1, i, :]
def save_frame_mine(prob,
total_index,
dir_name='',
video_name='',
vis=True):
result = prob_to_label(combine_prob(prob))
result_show = np.dstack(
(colors[result, 0], colors[result, 1], colors[result,
2])).astype(np.uint8)
if args.output != '' and dir_name != '':
out_file = os.path.join(str(dataset_dir), str(dir_name),
str(video_name), '%05d.png' % total_index)
if not os.path.exists(os.path.split(out_file)[0]):
os.makedirs(os.path.split(out_file)[0])
if vis:
cv2.imwrite(out_file, result_show)
else:
cv2.imwrite(out_file, result)
return
def save_frame(th, do_pause, dir_name='', vis=True):
result = prob_to_label(combine_prob(pred_prob[th]))
result_show = np.dstack(
(colors[result, 0], colors[result, 1], colors[result,
2])).astype(np.uint8)
if d.output != '' and dir_name != '':
out_file = os.path.join(dataset_dir, 'Results', 'Segmentations',
resolution, d.output, dir_name, video_dir,
'%05d.png' % th)
if not os.path.exists(os.path.split(out_file)[0]):
os.makedirs(os.path.split(out_file)[0])
if vis:
cv2.imwrite(out_file, result_show)
else:
cv2.imwrite(out_file, result)
temp = cv2.resize(frames[th],
frame_0.shape[-2::-1]) * 0.3 + result_show * 0.7
cv2.imshow('Result', temp.astype(np.uint8))
if do_pause:
cv2.waitKey()
else:
cv2.waitKey(100)
def main():
def save_frame(th, do_pause, dir_name='', vis=True):
result = prob_to_label(combine_prob(pred_prob[th]))
result_show = np.dstack(
(colors[result, 0], colors[result, 1], colors[result,
2])).astype(np.uint8)
if args.output != '' and dir_name != '':
out_file = os.path.join(dataset_dir, 'Results', 'Segmentations',
resolution, args.output, dir_name,
video_dir, '%05d.png' % th)
if not os.path.exists(os.path.split(out_file)[0]):
os.makedirs(os.path.split(out_file)[0])
if vis:
cv2.imwrite(out_file, result_show)
else:
cv2.imwrite(out_file, result)
temp = cv2.resize(frames[th],
frame_0.shape[-2::-1]) * 0.3 + result_show * 0.7
return
cv2.imshow('Result', temp.astype(np.uint8))
if do_pause:
cv2.waitKey()
else:
cv2.waitKey(100)
colors = labelcolormap(256)
global pred_prob, frames, flow1, flow2, orig_mask, \
model, instance_num, fr_h_r, fr_w_r, appear, bbox_cnt, \
location, patch_shapes
args = parse_args()
cfg = Config.from_file(args.config)
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
model = []
model = getattr(models, cfg.model.name)(cfg.model)
if os.path.isfile(cfg.weight):
print(("=> loading checkpoint '{}'".format(cfg.weight)))
checkpoint = torch.load(cfg.weight)
model.load_state_dict(checkpoint['state_dict'], True)
print(("=> loaded checkpoint"))
else:
raise (("=> no checkpoint found at '{}'".format(cfg.weight)))
model = model.cuda()
model.eval()
cudnn.benchmark = True
# Setup dataset
dataset_dir = os.path.join('data/DAVIS')
resolution = '480p'
imageset_dir = os.path.join(dataset_dir, 'ImageSets', '2017',
args.testset + '.txt')
video_list = []
for line in open(imageset_dir).readlines():
if line.strip() != '':
video_list.append(line.strip())
person_all = pickle_load(os.path.join(dataset_dir, 'PersonSearch',
resolution, args.testset + '.pkl'),
encoding='latin')
object_all = pickle_load(os.path.join(dataset_dir, 'ObjectSearch',
resolution, args.testset + '.pkl'),
encoding='latin')
category_all = pickle_load(os.path.join(dataset_dir, 'Class', resolution,
args.testset + '.pkl'),
encoding='latin')
frame_cnt = 0
use_cache = False #(args.cache != '')
video_cnt = -1
for video_dir in video_list:
video_cnt += 1
frame_dir = os.path.join(dataset_dir, 'JPEGImages', resolution,
video_dir)
frame_fr_dir = os.path.join(dataset_dir, 'JPEGImages',
'Full-Resolution', video_dir)
label_dir = os.path.join(dataset_dir, 'Annotations', resolution,
video_dir)
flow_dir = os.path.join(dataset_dir, 'Flow', resolution, video_dir)
cache_dir = os.path.join(dataset_dir, 'Cache', resolution, args.cache,
video_dir)
frames_num = len(os.listdir(frame_dir))
if (video_cnt % args.gpu_num != args.gpu):
frame_cnt += frames_num
continue
frame_0 = cv2.imread(os.path.join(frame_dir, '%05d.jpg' % 0))
#label_0 = cv2.imread(os.path.join(label_dir, '%05d.png' % 0), cv2.IMREAD_UNCHANGED)
label_0 = cv2.imread(os.path.join(label_dir, '%05d.png' % 0),
cv2.IMREAD_GRAYSCALE)
label_0 = convertColorMask2Ordinary(label_0)
instance_num = label_0.max()
frames = [None for _ in range(frames_num)]
pred_prob = [None for _ in range(frames_num)]
flow1 = [None for _ in range(frames_num)]
flow2 = [None for _ in range(frames_num)]
person_reid = [[None for _ in range(instance_num)]
for _ in range(frames_num)]
object_reid = [[None for _ in range(instance_num)]
for _ in range(frames_num)]
category = category_all[video_dir]
orig_mask = [None for _ in range(instance_num)]
frames[0] = cv2.imread(os.path.join(frame_fr_dir, '%05d.jpg' % 0))
fr_h_r = float(frames[0].shape[0]) / float(frame_0.shape[0])
fr_w_r = float(frames[0].shape[1]) / float(frame_0.shape[1])
pred_prob[0] = label_to_prob(label_0, instance_num)
person_reid[0] = person_all[frame_cnt]
object_reid[0] = object_all[frame_cnt]
save_frame(0, False, 'result', False)
bbox = gen_bbox(label_0, range(instance_num), True)
for i in range(instance_num):
orig_mask[i] = pred_prob[0][bbox[i, 1]:bbox[i, 3],
bbox[i, 0]:bbox[i, 2], i * 2 + 1]
for th in range(1, frames_num):
frames[th] = cv2.imread(os.path.join(frame_fr_dir,
'%05d.jpg' % th))
pred_prob[th] = label_to_prob(np.zeros_like(label_0, np.uint8),
instance_num)
flow1[th - 1] = flo.readFlow(
os.path.join(flow_dir, '%05d.flo' % (th - 1)))
flow2[th] = flo.readFlow(os.path.join(flow_dir, '%05d.rflo' % th))
person_reid[th] = person_all[frame_cnt + th]
object_reid[th] = object_all[frame_cnt + th]
bbox_cnt = -1000 * np.ones((frames_num, instance_num))
bbox_cnt[0, :] = 0
for th in range(frames_num):
for i in range(instance_num):
person_reid[th][i] = person_reid[th][i][:, [0, 1, 2, 3, 5]]
object_reid[th][i] = object_reid[th][i][:, [0, 1, 2, 3, 5]]
frame_cnt += frames_num
cache_file = os.path.join(cache_dir, '%s.pkl' % video_dir)
if (use_cache and os.path.exists(cache_file)):
pred_prob, bbox_cnt = pickle_load(cache_file, encoding='latin')
else:
predict(1, frames_num, 1, range(instance_num))
if use_cache:
if not os.path.exists(os.path.split(cache_file)[0]):
os.makedirs(os.path.split(cache_file)[0])
pickle_dump((pred_prob, bbox_cnt), cache_file)
appear = np.zeros((frames_num, instance_num)).astype(int)
location = np.zeros((frames_num, instance_num, 2)).astype(int)
update_appear()
for th in range(frames_num):
save_frame(th, False, 'draft', True)
for reid_target in ['person', 'object']:
cache_file = os.path.join(cache_dir,
'%s_%s.pkl' % (reid_target, video_dir))
if (use_cache and os.path.exists(cache_file)):
pred_prob, bbox_cnt = pickle_load(cache_file, encoding='latin')
else:
target_instance = []
for i in range(instance_num):
if (reid_target == 'object'
or category[i][123] > 0.5): # person is 123
target_instance.append(i)
reid_score = person_reid if reid_target == 'person' else object_reid
draft_cnt = 0
while (True):
max_score = 0
for i in range(1, frames_num - 1):
temp_label = prob_to_label(combine_prob(pred_prob[i]))
bbox_i = gen_bbox(temp_label, range(instance_num),
False, 0.99)
for j in target_instance:
if bbox_cnt[
i,
j] != i and reid_score[i][j].shape[0] > 0:
bbox_id = np.argmax(reid_score[i][j][:, 4])
# retrieval
if (appear[i, j] == 0):
x1, y1, x2, y2 = reid_score[i][j][bbox_id,
0:4]
if (reid_score[i][j][bbox_id,
4] > max_score
and reid_score[i][j][bbox_id,
4] > reid_th):
bbox_now = reid_score[i][j][
bbox_id, 0:4].astype(int)
result = np.bincount(temp_label[
bbox_now[1]:bbox_now[3] + 1,
bbox_now[0]:bbox_now[2] +
1].flatten(),
minlength=j + 2)
flag = True
if flag:
for occ_instance in np.where(
result[1:] > 0)[0]:
if (IoU(
bbox_now,
bbox_i[occ_instance]) >
bbox_occ_th):
flag = False
if flag:
for k in target_instance:
if (k != j
and appear[i, k] == 0
and reid_score[i][k][
bbox_id, 4] >
reid_score[i][j][
bbox_id, 4]):
flag = False
if flag:
max_frame = i
max_instance = j
max_bbox = reid_score[i][j][
bbox_id, 0:4]
max_score = reid_score[i][j][
bbox_id, 4]
if (max_score == 0):
break
bbox_cnt[max_frame, max_instance] = max_frame
predict_single(max_frame, max_instance, max_bbox,
orig_mask[max_instance])
save_frame(max_frame, False,
'%s_%05d_checkpoint' % (reid_target, draft_cnt))
temp = 0
for i in range(max_frame - 1, -1, -1):
if appear[i, max_instance] != 0:
temp = i
break
predict(max_frame - 1, temp, -1, [max_instance])
temp = frames_num
for i in range(max_frame + 1, frames_num, 1):
if appear[i, max_instance] != 0:
temp = i
break
predict(max_frame + 1, temp, 1, [max_instance])
update_appear()
for th in range(frames_num):
save_frame(th, False,
'%s_%05d' % (reid_target, draft_cnt))
draft_cnt = draft_cnt + 1
for th in range(frames_num):
save_frame(th, False, '%s' % reid_target)
if use_cache:
pickle_dump((pred_prob, bbox_cnt), cache_file)
for th in range(frames_num):
save_frame(th, False, 'result', False)
def predict(st, en, step, instance_list):
global pred_prob
global bbox_cnt
for th in range(st, en, step):
if (step == 1):
warp_prob = flo.get_warp_label(flow1[th - 1], flow2[th],
combine_prob(pred_prob[th - 1]))
elif (step == -1):
warp_prob = flo.get_warp_label(flow2[th + 1], flow1[th],
combine_prob(pred_prob[th + 1]))
bbox = gen_bbox(prob_to_label(warp_prob), range(instance_num), True)
new_instance_list = []
abort = True
temp = prob_to_label(combine_prob(pred_prob[th - step]))
for i in range(instance_num):
if (th == st) or (np.count_nonzero(orig_mask[i]) <=
np.count_nonzero(temp == (i + 1)) * 10):
if np.abs(bbox_cnt[th][i] - th) <= np.abs((st - step) - th):
continue
if i in instance_list:
abort = False
bbox_cnt[th][i] = (st - step)
new_instance_list.append(i)
else:
for j in instance_list:
if IoU(bbox[i, :], bbox[j, :]) > 1e-6:
new_instance_list.append(i)
break
if abort:
break
new_instance_list = sorted(new_instance_list)
temp_image = frames[th].astype(float)
f_prob = [
np.zeros(
[bbox[idx, 3] - bbox[idx, 1], bbox[idx, 2] - bbox[idx, 0], 2])
for idx in new_instance_list
]
image_patch = np.zeros(
(len(new_instance_list), patch_shape[1], patch_shape[0], 3), float)
flow_patch = np.zeros(
(len(new_instance_list), patch_shape[1], patch_shape[0], 2), float)
warp_label_patch = np.zeros(
(len(new_instance_list), patch_shape[1], patch_shape[0], 1), float)
for i in range(len(new_instance_list)):
idx = new_instance_list[i]
warp_label_patch[i, ..., 0] = cv2.resize(
warp_prob[bbox[idx, 1]:bbox[idx, 3], bbox[idx, 0]:bbox[idx, 2],
idx + 1], patch_shape).astype(float)
image_patch[i, ...] = cv2.resize(
temp_image[
int(0.5 +
bbox[idx, 1] * fr_h_r):int(0.5 +
bbox[idx, 3] * fr_h_r),
int(0.5 +
bbox[idx, 0] * fr_w_r):int(0.5 +
bbox[idx, 2] * fr_w_r), :],
patch_shape).astype(float)
if (step == 1):
flow_patch[i, ...] = cv2.resize(
flow2[th][bbox[idx, 1]:bbox[idx, 3],
bbox[idx, 0]:bbox[idx, 2], :],
patch_shape).astype(float)
else:
flow_patch[i, ...] = cv2.resize(
flow1[th][bbox[idx, 1]:bbox[idx, 3],
bbox[idx, 0]:bbox[idx, 2], :],
patch_shape).astype(float)
image_patch = torch.from_numpy(image_patch.transpose(
0, 3, 1, 2)).contiguous().float().cuda()
warp_label_patch = torch.from_numpy(
warp_label_patch.transpose(0, 3, 1,
2)).contiguous().float().cuda()
flow_patch = torch.from_numpy(flow_patch.transpose(
0, 3, 1, 2)).contiguous().float().cuda()
with torch.no_grad():
prob = model(image_patch, flow_patch, warp_label_patch)
prob = torch.nn.functional.softmax(prob, dim=1)
if use_flip:
image_patch = flip(image_patch, 3)
warp_label_patch = flip(warp_label_patch, 3)
flow_patch = flip(flow_patch, 3)
flow_patch[:, 0, ...] = -flow_patch[:, 0, ...]
with torch.no_grad():
prob_f = model(image_patch, flow_patch, warp_label_patch)
prob_f = torch.nn.functional.softmax(prob_f, dim=1)
prob_f = flip(prob_f, 3)
prob = (prob + prob_f) / 2.0
prob = prob.data.cpu().numpy().transpose(0, 2, 3, 1)
for i in range(len(new_instance_list)):
idx = new_instance_list[i]
f_prob[i] += cv2.resize(
prob[i, ...],
(bbox[idx, 2] - bbox[idx, 0], bbox[idx, 3] - bbox[idx, 1]))
for i in range(len(new_instance_list)):
idx = new_instance_list[i]
pred_prob[th][..., idx * 2] = 1
pred_prob[th][..., idx * 2 + 1] = 0
pred_prob[th][bbox[idx, 1]:bbox[idx, 3], bbox[idx, 0]:bbox[idx, 2],
idx * 2] = f_prob[i][..., 0]
pred_prob[th][bbox[idx, 1]:bbox[idx, 3], bbox[idx, 0]:bbox[idx, 2],
idx * 2 + 1] = f_prob[i][..., 1]
def predict_single(th, i, bbox, in_warp_label_patch):
# first time
global pred_prob
new_instance_list = [i]
# source
temp = gen_bbox(prob_to_label(combine_prob(pred_prob[th])),
range(instance_num), True)
result = prob_to_label(
combine_prob(pred_prob[th][temp[i, 1]:temp[i, 3] + 1,
temp[i, 0]:temp[i, 2] + 1, :]))
result = np.unique(result)
for j in result:
if (j != 0) and ((j - 1) not in new_instance_list):
new_instance_list.append(j - 1)
# unknow anything about i
pred_prob[th][(pred_prob[th][..., i * 2 + 1] > 0.5), 0::2] = 1
pred_prob[th][(pred_prob[th][..., i * 2 + 1] > 0.5), 1::2] = 0
bbox_w = bbox[2] - bbox[0]
bbox_h = bbox[3] - bbox[1]
bbox_enlarge = 0.0
bbox[0] = np.clip((bbox[0] - bbox_enlarge * bbox_w), 0,
pred_prob[th].shape[1] - 1)
bbox[2] = np.clip((bbox[2] + bbox_enlarge * bbox_w), bbox[0] + 1,
pred_prob[th].shape[1])
bbox[1] = np.clip((bbox[1] - bbox_enlarge * bbox_h), 0,
pred_prob[th].shape[0] - 1)
bbox[3] = np.clip((bbox[3] + bbox_enlarge * bbox_h), bbox[1] + 1,
pred_prob[th].shape[0])
bbox = bbox.astype(int)
bbox_w = bbox[2] - bbox[0]
bbox_h = bbox[3] - bbox[1]
temp_image = frames[th].astype(float)
for T in range(5):
image_patch = np.zeros((1, patch_shape[1], patch_shape[0], 3), float)
flow_patch = np.zeros((1, patch_shape[1], patch_shape[0], 2), float)
image_patch[0, ...] = cv2.resize(
temp_image[int(0.5 + bbox[1] * fr_h_r):int(0.5 + bbox[3] * fr_h_r),
int(0.5 + bbox[0] * fr_w_r):int(0.5 +
bbox[2] * fr_w_r), :],
patch_shape).astype(float)
if T == 0:
warp_label_patch = cv2.resize(
in_warp_label_patch, patch_shape)[np.newaxis, ...,
np.newaxis].astype(float)
else:
warp_label_patch = cv2.resize(
pred_prob[th][bbox[1]:bbox[3], bbox[0]:bbox[2], i * 2 + 1],
patch_shape)[np.newaxis, ..., np.newaxis].astype(float)
flow_patch[0, ...] = cv2.resize(
flow2[th][bbox[1]:bbox[3], bbox[0]:bbox[2], :],
patch_shape).astype(float)
image_patch = torch.from_numpy(image_patch.transpose(
0, 3, 1, 2)).contiguous().float().cuda()
warp_label_patch = torch.from_numpy(
warp_label_patch.transpose(0, 3, 1,
2)).contiguous().float().cuda()
flow_patch = torch.from_numpy(flow_patch.transpose(
0, 3, 1, 2)).contiguous().float().cuda()
with torch.no_grad():
prob = model(image_patch, flow_patch, warp_label_patch)
prob = torch.nn.functional.softmax(prob, dim=1)
if use_flip:
image_patch = flip(image_patch, 3)
warp_label_patch = flip(warp_label_patch, 3)
flow_patch = flip(flow_patch, 3)
flow_patch[:, 0, ...] = -flow_patch[:, 0, ...]
with torch.no_grad():
prob2 = model(image_patch, flow_patch, warp_label_patch)
prob2 = torch.nn.functional.softmax(prob2, dim=1)
prob2 = flip(prob2, 3)
prob = (prob + prob2) / 2.0
prob = prob.data.cpu().numpy().transpose(0, 2, 3, 1)
prob = cv2.resize(prob[0, ...], (bbox[2] - bbox[0], bbox[3] - bbox[1]))
pred_prob[th][..., i * 2] = 1
pred_prob[th][..., i * 2 + 1] = 0
pred_prob[th][bbox[1]:bbox[3], bbox[0]:bbox[2], i * 2] = prob[..., 0]
pred_prob[th][bbox[1]:bbox[3], bbox[0]:bbox[2], i * 2 + 1] = prob[...,
1]
temp = pred_prob[th][..., (i * 2):(i * 2 + 2)].copy()
pred_prob[th][(pred_prob[th][..., i * 2 + 1] > 0.5), 0::2] = 1
pred_prob[th][(pred_prob[th][..., i * 2 + 1] > 0.5), 1::2] = 0
pred_prob[th][..., (i * 2):(i * 2 + 2)] = temp.copy()
# target
bbox = gen_bbox(prob_to_label(combine_prob(pred_prob[th])),
range(instance_num), True)
result = prob_to_label(
combine_prob(pred_prob[th][bbox[i, 1]:bbox[i, 3],
bbox[i, 0]:bbox[i, 2], :]))
result = np.unique(result)
for j in result:
if (j != 0) and ((j - 1) not in new_instance_list):
new_instance_list.append(j - 1)
new_instance_list = sorted(new_instance_list)
for T in range(5):
# second time
warp_prob = combine_prob(pred_prob[th])
bbox = gen_bbox(prob_to_label(warp_prob), range(instance_num), True)
image_patch = np.zeros(
(len(new_instance_list), patch_shape[1], patch_shape[0], 3), float)
flow_patch = np.zeros(
(len(new_instance_list), patch_shape[1], patch_shape[0], 2), float)
warp_label_patch = np.zeros(
(len(new_instance_list), patch_shape[1], patch_shape[0], 1), float)
for i in range(len(new_instance_list)):
idx = new_instance_list[i]
warp_label_patch[i, ..., 0] = cv2.resize(
warp_prob[bbox[idx, 1]:bbox[idx, 3], bbox[idx, 0]:bbox[idx, 2],
idx + 1], patch_shape).astype(float)
image_patch[i, ...] = cv2.resize(
temp_image[
int(0.5 +
bbox[idx, 1] * fr_h_r):int(0.5 +
bbox[idx, 3] * fr_h_r),
int(0.5 +
bbox[idx, 0] * fr_w_r):int(0.5 +
bbox[idx, 2] * fr_w_r), :],
patch_shape).astype(float)
flow_patch[i, ...] = cv2.resize(
flow2[th][bbox[idx, 1]:bbox[idx, 3], bbox[idx, 0]:bbox[idx,
2], :],
patch_shape).astype(float)
image_patch = torch.from_numpy(image_patch.transpose(
0, 3, 1, 2)).contiguous().float().cuda()
warp_label_patch = torch.from_numpy(
warp_label_patch.transpose(0, 3, 1,
2)).contiguous().float().cuda()
flow_patch = torch.from_numpy(flow_patch.transpose(
0, 3, 1, 2)).contiguous().float().cuda()
with torch.no_grad():
prob = model(image_patch, flow_patch, warp_label_patch)
prob = torch.nn.functional.softmax(prob, dim=1)
if use_flip:
image_patch = flip(image_patch, 3)
warp_label_patch = flip(warp_label_patch, 3)
flow_patch = flip(flow_patch, 3)
flow_patch[:, 0, ...] = -flow_patch[:, 0, ...]
with torch.no_grad():
prob2 = model(image_patch, flow_patch, warp_label_patch)
prob2 = torch.nn.functional.softmax(prob2, dim=1)
prob2 = flip(prob2, 3)
prob = (prob + prob2) / 2.0
prob = prob.data.cpu().numpy().transpose(0, 2, 3, 1)
for i in range(len(new_instance_list)):
idx = new_instance_list[i]
f_prob = cv2.resize(
prob[i, ...],
(bbox[idx, 2] - bbox[idx, 0], bbox[idx, 3] - bbox[idx, 1]))
pred_prob[th][..., idx * 2] = 1
pred_prob[th][..., idx * 2 + 1] = 0
pred_prob[th][bbox[idx, 1]:bbox[idx, 3], bbox[idx, 0]:bbox[idx, 2],
idx * 2] = f_prob[..., 0]
pred_prob[th][bbox[idx, 1]:bbox[idx, 3], bbox[idx, 0]:bbox[idx, 2],
idx * 2 + 1] = f_prob[..., 1]
def predict(st, en, step, instance_list):
# st = 1, step =1, instance_list: obj number id.
global pred_prob
global bbox_cnt
global count
for th in range(st, en, step):
if (step == 1):
temp_prob_ = combine_prob(pred_prob[th - 1])
# print (temp_prob_.shape) # (720, 1280, 3)
# temp_prob_[temp_prob_ > mask_th] = 1
warp_prob = flo.get_warp_label(flow1[th - 1], flow2[th],
temp_prob_)
elif (step == -1):
temp_prob_ = combine_prob(pred_prob[th + 1])
# temp_prob_[temp_prob_ > mask_th] = 1
warp_prob = flo.get_warp_label(flow2[th + 1], flow1[th],
temp_prob_)
bbox = gen_bbox(prob_to_label(warp_prob), range(instance_num), True)
new_instance_list = []
abort = True
temp = prob_to_label(combine_prob(pred_prob[th - step]))
for i in range(instance_num):
if (th == st) or (np.count_nonzero(orig_mask[i]) <=
np.count_nonzero(temp == (i + 1)) * 10):
if np.abs(bbox_cnt[th][i] - th) <= np.abs((st - step) - th):
continue
if i in instance_list:
abort = False
bbox_cnt[th][i] = (st - step)
new_instance_list.append(i)
else:
for j in instance_list:
if IoU(bbox[i, :], bbox[j, :]) > 1e-6:
new_instance_list.append(i)
break
if abort:
break
new_instance_list = sorted(new_instance_list)
temp_image = frames[th].astype(float)
f_prob = [
np.zeros(
[bbox[idx, 3] - bbox[idx, 1], bbox[idx, 2] - bbox[idx, 0], 2])
for idx in new_instance_list
]
image_patch = np.zeros(
(len(new_instance_list), patch_shape[1], patch_shape[0], 3), float)
flow_patch = np.zeros(
(len(new_instance_list), patch_shape[1], patch_shape[0], 2), float)
warp_label_patch = np.zeros(
(len(new_instance_list), patch_shape[1], patch_shape[0], 1), float)
for i in range(len(new_instance_list)):
idx = new_instance_list[i]
warp_label_patch_temp = cv2.resize(
warp_prob[bbox[idx, 1]:bbox[idx, 3], bbox[idx, 0]:bbox[idx, 2],
idx + 1], patch_shape).astype(float)
# print (warp_label_patch_temp.shape) # (433, 433)
warp_label_patch_temp[warp_label_patch_temp >= mask_th] = 1
warp_label_patch_temp[warp_label_patch_temp < mask_th] = 0
warp_label_patch[i, ..., 0] = warp_label_patch_temp
image_patch[i, ...] = cv2.resize(
temp_image[
int(0.5 +
bbox[idx, 1] * fr_h_r):int(0.5 +
bbox[idx, 3] * fr_h_r),
int(0.5 +
bbox[idx, 0] * fr_w_r):int(0.5 +
bbox[idx, 2] * fr_w_r), :],
patch_shape).astype(float)
if (step == 1):
flow_patch[i, ...] = cv2.resize(
flow2[th][bbox[idx, 1]:bbox[idx, 3],
bbox[idx, 0]:bbox[idx, 2], :],
patch_shape).astype(float)
else:
flow_patch[i, ...] = cv2.resize(
flow1[th][bbox[idx, 1]:bbox[idx, 3],
bbox[idx, 0]:bbox[idx, 2], :],
patch_shape).astype(float)
image_patch = torch.from_numpy(image_patch.transpose(0, 3, 1, 2))
warp_label_patch = torch.from_numpy(
warp_label_patch.transpose(0, 3, 1,
2)).contiguous().float().cuda()
flow_patch = torch.from_numpy(flow_patch.transpose(
0, 3, 1, 2)).contiguous().float().cuda()
print(image_patch.size(), 'image_patch')
print(warp_label_patch.size(), 'warp_label_patch')
print(flow_patch.size(), 'flow_patch')
#notice here!!!!!!! if annotate it, you will be mad.
print((image_patch[0]).cpu().numpy().transpose(1, 2, 0).shape)
cv2.imwrite('output/%3d_prob.png' % count,
(image_patch[0]).cpu().numpy().transpose(1, 2, 0))
count += 1
cv2.imwrite('output/%3d_prob.png' % count,
(warp_label_patch[0][0] * 255).cpu().numpy())
count += 1
value_scale = 255
mean = [0.485, 0.456, 0.406]
mean = [item * value_scale for item in mean]
std = [0.229, 0.224, 0.225]
std = [item * value_scale for item in std]
for i in range(len(new_instance_list)):
for t, m, s in zip(image_patch[i], mean, std):
t.sub_(m).div_(s)
# notice here!!!!!!! if annotate it, you will be mad.
image_patch = image_patch.contiguous().float().cuda()
with torch.no_grad():
prob = model(
image_patch, warp_label_patch,
flow_patch) #model = torch.nn.DataParallel(model).cuda()
prob = torch.nn.functional.softmax(prob, dim=1)
cv2.imwrite(
'output/%3d_prob.png' % count,
(torch.argmax(prob[0], dim=0, keepdim=False)).cpu().numpy() *
100)
count += 1
# warp_label_patch debug1
# probxxx = prob.data.cpu().numpy().transpose(0, 2, 3, 1)
# print (probxxx[0].shape)
# out = np.argmax(probxxx[0], axis=2)
# print (out.shape)
# cv2.imwrite("debugs/"+str(th)+"_out.png", (out * 255))
if use_flip:
image_patch = flip(image_patch, 3)
warp_label_patch = flip(warp_label_patch, 3)
flow_patch = flip(flow_patch, 3)
flow_patch[:, 0, ...] = -flow_patch[:, 0, ...]
with torch.no_grad():
prob_f = model(image_patch, warp_label_patch, flow_patch)
#cv2.imwrite('output/%d_probf.png' % count, (prob_f[0][0]*255).cpu().numpy())
#count += 1
prob_f = torch.nn.functional.softmax(prob_f, dim=1)
prob_f = flip(prob_f, 3)
prob = (prob + prob_f) / 2.0
prob = prob.data.cpu().numpy().transpose(0, 2, 3, 1)
for i in range(len(new_instance_list)):
idx = new_instance_list[i]
f_prob[i] += cv2.resize(
prob[i, ...],
(bbox[idx, 2] - bbox[idx, 0], bbox[idx, 3] - bbox[idx, 1]))
for i in range(len(new_instance_list)):
idx = new_instance_list[i]
pred_prob[th][..., idx * 2] = 1
pred_prob[th][..., idx * 2 + 1] = 0
pred_prob[th][bbox[idx, 1]:bbox[idx, 3], bbox[idx, 0]:bbox[idx, 2],
idx * 2] = f_prob[i][..., 0]
pred_prob[th][bbox[idx, 1]:bbox[idx, 3], bbox[idx, 0]:bbox[idx, 2],
idx * 2 + 1] = f_prob[i][..., 1]
