コード例 #1
0
            # Display the captured frame:
            testIm(frame)
    
            # Press q on keyboard to exit the program
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
        # Break the loop
        else:
            break
        
    


if __name__ == '__main__':
    net = FaceBox()
    net.load_state_dict(torch.load('weight/faceboxes.pt', map_location=lambda storage, loc:storage), strict=False) 
    net.eval()
    data_encoder = DataEncoder()

    # given image path, predict and show
    root_path = "picture/"
    #picture = 'img_36.jpg' timg.jpg
    picture = 'timg.jpg'
    #testIm(root_path + picture)
    #im = cv2.imread(root_path + picture)
    cap = cv2.VideoCapture(0)
    cvDemo()
    cap.release()
    cv2.destroyAllWindows()

コード例 #2
0
ファイル: trainvisdom.py プロジェクト: ysc703/faceboxes_plate
def train():
    use_gpu = torch.cuda.is_available()
    file_root = os.path.dirname(os.path.abspath(__file__))

    learning_rate = 0.00000001
    num_epochs = 300
    batch_size = 32

    net = FaceBox()
    if use_gpu:
        net.cuda()

    print('load model...')
    net.load_state_dict(torch.load('weight/faceboxes.pt'))

    criterion = MultiBoxLoss()
    optimizer = torch.optim.SGD(net.parameters(),
                                lr=learning_rate,
                                momentum=0.99,
                                weight_decay=0.00055)
    #optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate, weight_decay=1e-4)
    scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
                                                     milestones=[248, 300],
                                                     gamma=0.1)

    train_dataset = ListDataset(root=file_root,
                                list_file='./traintest.txt',
                                train=True,
                                transform=[transforms.ToTensor()])
    train_loader = DataLoader(train_dataset,
                              batch_size=batch_size,
                              shuffle=True,
                              num_workers=4,
                              drop_last=True)
    val_dataset = ListDataset(root=file_root,
                              list_file='./new_val.txt',
                              train=False,
                              transform=[transforms.ToTensor()])
    val_loader = DataLoader(val_dataset,
                            batch_size=batch_size,
                            shuffle=False,
                            num_workers=1)
    print('the dataset has %d images' % (len(train_dataset)))
    print('the batch_size is %d' % (batch_size))

    num_iter = 0
    vis = visdom.Visdom()
    win = vis.line(Y=np.array([0]), X=np.array([0]))

    net.train()
    for epoch in range(num_epochs):
        scheduler.step()

        print('\n\nStarting epoch %d / %d' % (epoch + 1, num_epochs))
        print('Learning Rate for this epoch: {}'.format(learning_rate))

        total_loss = 0.
        net.train()
        for i, (images, loc_targets, conf_targets) in enumerate(train_loader):

            if use_gpu:
                images = images.cuda()
                loc_targets = loc_targets.cuda()
                conf_targets = conf_targets.cuda()

            loc_preds, conf_preds = net(images)
            loss = criterion(loc_preds, loc_targets, conf_preds, conf_targets)
            total_loss += loss.item()

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            if (i + 1) % 10 == 0:
                print(
                    'Epoch [{}/{}], Iter [{}/{}] Loss: {:.4f}, average_loss: {:.4f}'
                    .format(epoch + 1, num_epochs, i + 1, len(train_loader),
                            loss.item(), total_loss / (i + 1)))

                vis.line(Y=np.array([total_loss / (i + 1)]),
                         X=np.array([num_iter]),
                         win=win,
                         name='train',
                         update='append')
                num_iter += 1
        val_loss = 0.0
        net.eval()
        for idx, (images, loc_targets, conf_targets) in enumerate(val_loader):
            with torch.no_grad():
                if use_gpu:
                    images = images.cuda()
                    loc_targets = loc_targets.cuda()
                    conf_targets = conf_targets.cuda()

                loc_preds, conf_preds = net(images)
                loss = criterion(loc_preds, loc_targets, conf_preds,
                                 conf_targets)
                val_loss += loss.item()
        val_loss /= len(val_dataset) / batch_size
        vis.line(Y=np.array([val_loss]),
                 X=np.array([epoch * 30 + 30]),
                 win=win,
                 name='val',
                 update='append')
        print('loss of val is {}'.format(val_loss))

        if not os.path.exists('weight/'):
            os.mkdir('weight')

        print('saving model ...')
        torch.save(net.state_dict(), 'weight/faceboxes.pt')