/
eval_video_segmentation.py
243 lines (200 loc) · 9.15 KB
/
eval_video_segmentation.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
"""
###########################################################################
This video segmentation evaluation code is built on code from
https://github.com/facebookresearch/dino. Original copyright
notice shown below:
Adapted from https://github.com/Liusifei/UVC
Written by: Saksham Suri
###########################################################################
"""
# Copyright (c) Facebook, Inc. and its affiliates.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import copy
import glob
import queue
from urllib.request import urlopen
import argparse
import numpy as np
from tqdm import tqdm
import cv2
import torch
import torch.nn as nn
from torch.nn import functional as F
from PIL import Image
from torchvision import transforms
@torch.no_grad()
def eval_video_tracking_davis(args, model, frame_list, video_dir, first_seg, seg_ori, color_palette):
"""
Evaluate tracking on a video given first frame & segmentation
"""
os.makedirs(args.output_dir, exist_ok=True)
os.makedirs(os.path.join(args.output_dir, model.mod_id), exist_ok=True)
os.makedirs(os.path.join(args.output_dir, model.mod_id, f'{args.dataset}_vidseg_224_blk{args.blk}'), exist_ok=True)
output_dir = os.path.join(args.output_dir, model.mod_id, f'{args.dataset}_vidseg_224_blk{args.blk}')
video_folder = os.path.join(output_dir, video_dir.split('/')[-1])
os.makedirs(video_folder, exist_ok=True)
# The queue stores the n preceeding frames
que = queue.Queue(args.n_last_frames)
# first frame
frame1, ori_h, ori_w = read_frame(frame_list[0])
# extract first frame feature
frame1_feat = extract_feature(args, model, frame1, patch_size=args.patch, imsize=args.imsize).T # dim x h*w
# saving first segmentation
out_path = os.path.join(video_folder, "00000.png")
imwrite_indexed(out_path, seg_ori, color_palette)
mask_neighborhood = None
for cnt in tqdm(range(1, len(frame_list))):
frame_tar = read_frame(frame_list[cnt])[0]
# we use the first segmentation and the n previous ones
used_frame_feats = [frame1_feat] + [pair[0] for pair in list(que.queue)]
used_segs = [first_seg] + [pair[1] for pair in list(que.queue)]
frame_tar_avg, feat_tar, mask_neighborhood = label_propagation(args, model, frame_tar, used_frame_feats, used_segs, mask_neighborhood)
# pop out oldest frame if neccessary
if que.qsize() == args.n_last_frames:
que.get()
# push current results into queue
seg = copy.deepcopy(frame_tar_avg)
que.put([feat_tar, seg])
# upsampling & argmax
frame_tar_avg = F.interpolate(frame_tar_avg, scale_factor=args.patch, mode='bilinear', align_corners=False, recompute_scale_factor=False)[0]
frame_tar_avg = norm_mask(frame_tar_avg)
_, frame_tar_seg = torch.max(frame_tar_avg, dim=0)
# saving to disk
frame_tar_seg = np.array(frame_tar_seg.squeeze().cpu(), dtype=np.uint8)
frame_tar_seg = np.array(Image.fromarray(frame_tar_seg).resize((ori_w, ori_h), 0))
frame_nm = frame_list[cnt].split('/')[-1].replace(".jpg", ".png")
imwrite_indexed(os.path.join(video_folder, frame_nm), frame_tar_seg, color_palette)
def restrict_neighborhood(args, h, w):
# We restrict the set of source nodes considered to a spatial neighborhood of the query node (i.e. ``local attention'')
mask = torch.zeros(h, w, h, w)
for i in range(h):
for j in range(w):
for p in range(2 * args.size_mask_neighborhood + 1):
for q in range(2 * args.size_mask_neighborhood + 1):
if i - args.size_mask_neighborhood + p < 0 or i - args.size_mask_neighborhood + p >= h:
continue
if j - args.size_mask_neighborhood + q < 0 or j - args.size_mask_neighborhood + q >= w:
continue
mask[i, j, i - args.size_mask_neighborhood + p, j - args.size_mask_neighborhood + q] = 1
mask = mask.reshape(h * w, h * w)
return mask.cuda(non_blocking=True)
def norm_mask(mask):
c, h, w = mask.size()
for cnt in range(c):
mask_cnt = mask[cnt,:,:]
if(mask_cnt.max() > 0):
mask_cnt = (mask_cnt - mask_cnt.min())
mask_cnt = mask_cnt/mask_cnt.max()
mask[cnt,:,:] = mask_cnt
return mask
def label_propagation(args, model, frame_tar, list_frame_feats, list_segs, mask_neighborhood=None):
"""
propagate segs of frames in list_frames to frame_tar
"""
## we only need to extract feature of the target frame
feat_tar, h, w = extract_feature(args, model, frame_tar, return_h_w=True, patch_size = args.patch, imsize=args.imsize)
return_feat_tar = feat_tar.T # dim x h*w
ncontext = len(list_frame_feats)
feat_sources = torch.stack(list_frame_feats) # nmb_context x dim x h*w
feat_tar = F.normalize(feat_tar, dim=1, p=2)
feat_sources = F.normalize(feat_sources, dim=1, p=2)
feat_tar = feat_tar.unsqueeze(0).repeat(ncontext, 1, 1)
aff = torch.exp(torch.bmm(feat_tar, feat_sources) / 0.1) # nmb_context x h*w (tar: query) x h*w (source: keys)
if args.size_mask_neighborhood > 0:
if mask_neighborhood is None:
mask_neighborhood = restrict_neighborhood(args, h, w)
mask_neighborhood = mask_neighborhood.unsqueeze(0).repeat(ncontext, 1, 1)
aff *= mask_neighborhood
aff = aff.transpose(2, 1).reshape(-1, h * w) # nmb_context*h*w (source: keys) x h*w (tar: queries)
tk_val, _ = torch.topk(aff, dim=0, k=args.topk)
tk_val_min, _ = torch.min(tk_val, dim=0)
aff[aff < tk_val_min] = 0
aff = aff / torch.sum(aff, keepdim=True, axis=0)
list_segs = [s.cuda() for s in list_segs]
segs = torch.cat(list_segs)
nmb_context, C, h, w = segs.shape
segs = segs.reshape(nmb_context, C, -1).transpose(2, 1).reshape(-1, C).T # C x nmb_context*h*w
seg_tar = torch.mm(segs, aff)
seg_tar = seg_tar.reshape(1, C, h, w)
return seg_tar, return_feat_tar, mask_neighborhood
def extract_feature(args, model, frame, return_h_w=False, patch_size=16, imsize=224):
"""Extract one frame feature everytime."""
if(args.dense):
out = model.get_activations(frame)[0].float()
dim = out.shape[-1]
h, w = out.shape[0], out.shape[1]
out = out.reshape(-1, dim)
else:
frame = model.transform(frame)
out = model.get_activations(frame.unsqueeze(0).cuda())[0].float()
out = out[:, 1:, :] # we discard the [CLS] token
h, w = int(frame.shape[1] / patch_size), int(frame.shape[2] / patch_size)
dim = out.shape[-1]
out = out[0].reshape(h, w, dim)
out = out.reshape(-1, dim)
if return_h_w:
return out, h, w
return out
def imwrite_indexed(filename, array, color_palette):
""" Save indexed png for DAVIS."""
if np.atleast_3d(array).shape[2] != 1:
raise Exception("Saving indexed PNGs requires 2D array.")
im = Image.fromarray(array)
im.putpalette(color_palette.ravel())
im.save(filename, format='PNG')
def to_one_hot(y_tensor, n_dims=None):
"""
Take integer y (tensor or variable) with n dims &
convert it to 1-hot representation with n+1 dims.
"""
if(n_dims is None):
n_dims = int(y_tensor.max()+ 1)
_,h,w = y_tensor.size()
y_tensor = y_tensor.type(torch.LongTensor).view(-1, 1)
n_dims = n_dims if n_dims is not None else int(torch.max(y_tensor)) + 1
y_one_hot = torch.zeros(y_tensor.size()[0], n_dims).scatter_(1, y_tensor, 1)
y_one_hot = y_one_hot.view(h,w,n_dims)
return y_one_hot.permute(2, 0, 1).unsqueeze(0)
def read_frame_list(video_dir):
frame_list = [img for img in glob.glob(os.path.join(video_dir,"*.jpg"))]
frame_list = sorted(frame_list)
return frame_list
def read_frame(frame_dir):
img = Image.open(frame_dir)
ori_w, ori_h = img.size
return img, ori_h, ori_w
def read_seg(seg_dir, factor, scale_size=[224,224]):
seg = Image.open(seg_dir)
_w, _h = seg.size # note PIL.Image.Image's size is (w, h)
if len(scale_size) == 1:
if(_w > _h):
_th = scale_size[0]
_tw = (_th * _w) / _h
_tw = int((_tw // factor) * factor)
else:
_tw = scale_size[0]
_th = (_tw * _h) / _w
_th = int((_th // factor) * factor)
else:
_th = scale_size[1]
_tw = scale_size[0]
small_seg = np.array(seg.resize((_tw // factor, _th // factor), 0))
small_seg = torch.from_numpy(small_seg.copy()).contiguous().float().unsqueeze(0)
return to_one_hot(small_seg), np.asarray(seg)
def color_normalize(x, mean=[0.485, 0.456, 0.406], std=[0.228, 0.224, 0.225]):
for t, m, s in zip(x, mean, std):
t.sub_(m)
t.div_(s)
return x