Mem_every=mem_every,
Mem_number=None)
# Save results for quantitative eval ######################
for f in range(num_frames):
img_E = Image.fromarray(pred[f])
img_E.putpalette(palette)
img_E.save(os.path.join(test_path, '{:05d}.png'.format(f)))
torch.cuda.empty_cache()
if VIZ:
from helpers import overlay_davis
# visualize results #######################
viz_path = os.path.join('./viz/', code_name, seq_name)
if not os.path.exists(viz_path):
os.makedirs(viz_path)
for f in range(num_frames):
pF = (Fs[0, :, f].permute(1, 2, 0).numpy() * 255.).astype(np.uint8)
pE = pred[f]
canvas = overlay_davis(pF, pE, palette)
canvas = Image.fromarray(canvas)
canvas.save(os.path.join(viz_path, 'f{}.jpg'.format(f)))
vid_path = os.path.join('./viz/', code_name, '{}.mp4'.format(seq_name))
frame_path = os.path.join('./viz/', code_name, seq_name, 'f%d.jpg')
os.system(
'ffmpeg -framerate 10 -i {} {} -vcodec libx264 -crf 10 -pix_fmt yuv420p -nostats -loglevel 0 -y'
.format(frame_path, vid_path))
def run_test(cuda_device):
torch.set_grad_enabled(False) # Volatile
args = get_arguments()
GPU = args.g
YEAR = args.y
SET = args.s
VIZ = args.viz
DATA_ROOT = args.D
device = cuda_device
# Model and version
MODEL = 'STM'
print(MODEL, ': Testing on DAVIS', YEAR)
os.environ['CUDA_VISIBLE_DEVICES'] = GPU
print('--- CUDA:')
if VIZ:
print(
'--- Produce mask overaid video outputs. Evaluation will run slow.'
)
print('--- Require FFMPEG for encoding, Check folder ./viz')
palette = Image.open(DATA_ROOT +
'/Annotations/606332/00000.png').getpalette()
Testset = DAVIS_MO_Test(DATA_ROOT,
resolution='480p',
imset='test.txt',
single_object=(YEAR == 16))
Testloader = data.DataLoader(Testset,
batch_size=1,
shuffle=False,
num_workers=0,
pin_memory=False)
# model = nn.DataParallel(STM())
model = STM()
try:
model = model.to(device)
print('--- Model successfully sent to cuda')
except:
print('--- Cuda is not available')
model.eval() # turn-off BN
pth_path = '../user_data/STM.pth'
print('Loading weights:', pth_path)
state = model.state_dict()
# load entire saved model from checkpoint
checkpoint = torch.load(
pth_path,
map_location='cpu') # dict_keys(['epoch', 'model', 'optimizer'])
# checkpoint['model'] = {k: v for k, v in checkpoint['model'].items() if k in state}
# # overwrite entries in the existing state dict
# state.update(checkpoint['model'])
# load the new state dict
model = nn.DataParallel(model)
model.load_state_dict(checkpoint)
# params = []
# for key, value in dict(model.named_parameters()).items():
# if value.requires_grad:
# params += [{'params': [value]}]
# # load optimizere state dict
# optimizer = torch.optim.Adam(params, lr=1e-5)
# if 'optimizer' in checkpoint.keys():
# optimizer.load_state_dict(checkpoint['optimizer'])
# del checkpoint
# model.load_state_dict(torch.load(pth_path,map_location='cpu'))
code_name = '{}_DAVIS_{}_{}'.format(MODEL, YEAR, SET)
print('Start Testing:', code_name)
for seq, V in enumerate(Testloader):
# if seq < 21:
# continue
Fs, Ms, num_objects, info = V
#Fs: torch.Size([1, 3, 69, 480, 910])
#Ms: torch.Size([1, 11, 69, 480, 910])
#num_objects: tensor([[2]])
#info: {'name': ['bike-packing'], 'num_frames': tensor([69]),
# 'size_480p': [tensor([480]), tensor([910])]}
seq_name = info['name'][0] # = bike-packing
num_frames = info['num_frames'][0].item() # = 69
# if num_frames > 40:
# num_frames = 40
print('[{}]: num_frames: {}, num_objects: {}'.format(
seq_name, num_frames, num_objects[0][0]))
pred, Es = Run_video(model,
Fs,
Ms,
num_frames,
num_objects,
Mem_every=5,
Mem_number=None)
#pred: (69, 480, 910)
#Es: torch.Size([1, 11, 69, 480, 910])
# Save results for quantitative eval ######################
test_path = os.path.join('../user_data', code_name, seq_name)
print('test_path: ', test_path)
if not os.path.exists(test_path):
os.makedirs(test_path)
for f in range(num_frames):
img_E = Image.fromarray(pred[f])
#img_E: (910, 480)
#imagem com pixels rotulados de 0 a Nº de objetos
img_E.putpalette(palette)
#putpalette troca os rótulos pelos mesmos da anotação manual de
#onde se extraiu o palette
img_E.save(os.path.join(test_path, '{:05d}.png'.format(f)))
#Código para visualizar a máscara de segmentação sobre o frame:
if VIZ:
from helpers import overlay_davis
# visualize results #######################
viz_path = os.path.join('./viz/', code_name, seq_name)
if not os.path.exists(viz_path):
os.makedirs(viz_path)
for f in range(num_frames):
pF = (Fs[0, :, f].permute(1, 2, 0).numpy() * 255.).astype(
np.uint8)
pE = pred[f]
canvas = overlay_davis(pF, pE, palette)
canvas = Image.fromarray(canvas)
canvas.save(os.path.join(viz_path, 'f{}.jpg'.format(f)))