コード例 #1
0
    def test_transform_siamfc(self):
        base_dataset = VOT(self.vot_dir, return_rect=True, download=True)
        transform = TransformSiamFC(stats_path=self.stats_path)
        dataset = Pairwise(base_dataset, transform=transform, subset='train')
        self.assertGreater(len(dataset), 0)

        for crop_z, crop_x, labels, weights in dataset:
            self.assertAlmostEqual(weights[labels == 1].sum().item(),
                                   weights[labels == 0].sum().item())
            self.assertAlmostEqual(weights.sum().item(),
                                   labels[labels >= 0].numel())
            self.assertEqual(
                weights[labels == transform.ignore_label].sum().item(), 0)

        if self.visualize:
            crop_z, crop_x, labels, weights = random.choice(dataset)
            crop_z = F.to_pil_image(crop_z / 255.0)
            crop_x = F.to_pil_image(crop_x / 255.0)
            labels = self._rescale(labels.cpu().squeeze().numpy())
            weights = self._rescale(weights.cpu().squeeze().numpy())

            bndbox_z = np.array([31, 31, 64, 64])
            bndbox_x = np.array([95, 95, 64, 64])

            show_frame(crop_z, bndbox_z, fig_n=1, pause=1)
            show_frame(crop_x, bndbox_x, fig_n=2, pause=1)
            show_frame(labels, fig_n=3, pause=1, cmap='hot')
            show_frame(weights, fig_n=4, pause=5, cmap='hot')
コード例 #2
0
ファイル: test_siamfc.py プロジェクト: wpfhtl/open-vot
    def test_siamfc_train_v2(self):
        tracker = TrackerSiamFC(branch='alexv2')
        transform = TransformSiamFC(
            stats_path=self.stats_path, score_sz=33,
            r_pos=8, total_stride=4)

        base_dataset = VOT(self.vot_dir, return_rect=True, download=True)
        dataset = Pairwise(base_dataset, transform, pairs_per_video=1)
        dataloader = DataLoader(dataset, batch_size=2, shuffle=True)

        # training loop
        for it, batch in enumerate(dataloader):
            update_lr = it == 0
            loss = tracker.step(batch, backward=True, update_lr=update_lr)
            print('Iter: {} Loss: {:.6f}'.format(it + 1, loss))

        # val loop
        for it, batch in enumerate(dataloader):
            loss = tracker.step(batch, backward=False)
            print('Val. Iter: {} Loss: {:.6f}'.format(it + 1, loss))