コード例 #1
0
                         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))
コード例 #2
0
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)))