Exemple #1
0
def train(cfg):
    model = ResNetUNet(pretrain=cfg.imagenet_pretrain, backbone=cfg.backbone)
    if cfg.pretrained_pth:
        model.load_state_dict(
            torch.load(cfg.pretrained_pth, map_location='cpu'))
    if torch.cuda.device_count() > 1:
        model = DataParallel(model)
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    model.to(device)

    trainset = SynthTextDataset(
        cfg) if cfg.mode == 'pretrain' else FinetuneDataset(cfg)
    train_loader = DataLoader(trainset, batch_size=cfg.batch_size, shuffle=True, \
                                      num_workers=cfg.num_workers, pin_memory=False, drop_last=cfg.drop_last)

    file_num = len(trainset)
    batch_num = int(file_num / cfg.batch_size)
    criterion = Loss_OHEM()

    optimizer = optim.Adam(model.parameters(),
                           lr=cfg.high_lr,
                           weight_decay=cfg.weight_decay)
    scheduler = lr_scheduler.CosineAnnealingLR(optimizer,
                                               cfg.epoch_iter * batch_num,
                                               cfg.low_lr)

    model.train()
    for epoch in range(cfg.epoch_iter):
        epoch_loss = 0
        epoch_time = time.time()
        for i, (img, text, ignore, rho, theta) in enumerate(train_loader):
            img, text, ignore, rho, theta = list(
                map(lambda x: x.to(device), [img, text, ignore, rho, theta]))
            pred_cls, pred_rho, pred_theta = model(img)
            loss = criterion(text, ignore, rho, theta, pred_cls, pred_rho,
                             pred_theta)

            epoch_loss += loss.item()
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            scheduler.step()

            print('Epoch is [{}/{}], mini-batch is [{}/{}], batch_loss is {:.8f}'.format(\
                         epoch+1, cfg.epoch_iter, i+1, batch_num, loss.item()))
            sys.stdout.flush()

        if (epoch + 1) % cfg.save_interval == 0:
            torch.save(
                model.module.state_dict(),
                os.path.join(cfg.pths_path,
                             'model_epoch_{}.pth'.format(epoch + 1)))
            print(time.asctime(time.localtime(time.time())))

        print('epoch_loss is {:.8f}, epoch_time is {:.8f}'.format(
            epoch_loss / batch_num,
            time.time() - epoch_time))
        print(time.asctime(time.localtime(time.time())))
        print('=' * 50)
Exemple #2
0
def setup_model_dataloader(args, batch_size):
    num_class = 1
    model = ResNetUNet(num_class)
    device = torch.device("cuda")
    model_file = os.path.join(args.model_dir, 'best_model.pth')
    model.load_state_dict(torch.load(model_file, map_location="cuda:0"))
    model.to(device)

    # Set model to the evaluation mode
    model.eval()

    # Setup dataset
    # Need to be careful here. This isn't perfect.
    # I'm assuming that the dataset isn't changing between training and inference time
    bee_ds = BeePointDataset(root_dir=args.data_dir)
    if 1:
        dbfile = open(os.path.join(args.model_dir, 'test_ds.pkl'), 'rb')
        test_ds = pickle.load(dbfile)
        #test_loader = DataLoader(test_ds, batch_size=1, shuffle=False, num_workers=1)
        test_loader = DataLoader(test_ds,
                                 batch_size=batch_size,
                                 shuffle=True,
                                 num_workers=1,
                                 collate_fn=helper.bee_collate_fn)
    else:
        # Just use the defaults in the bee_ds
        test_loader = DataLoader(bee_ds,
                                 batch_size=batch_size,
                                 shuffle=True,
                                 num_workers=1,
                                 collate_fn=helper.bee_collate_fn)

    return model, test_loader, device
Exemple #3
0
def setup_model_dataloader(args, batch_size):
    num_class = 1
    model = ResNetUNet(num_class)
    device = torch.device("cuda")
    model_file = os.path.join(args.model_dir, 'best_model.pth')
    model.load_state_dict(torch.load(model_file, map_location="cuda:0"))
    model.to(device)
    # Set model to the evaluation mode
    model.eval()
    return model, device
def main():
    device = utils.get_device()
    utils.set_seed(args.seed, device)  # set random seed

    dataset = LungSegDataSet(args.datapath)

    net = ResNetUNet(n_class=1).to(device)

    optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad,
                                        net.parameters()),
                                 lr=1e-4)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
                                                step_size=25,
                                                gamma=0.1)
    trainloader = torch.utils.data.DataLoader(dataset,
                                              batch_size=args.bstrain,
                                              shuffle=True,
                                              num_workers=args.nworkers)
    for epoch in range(args.maxepoch):
        scheduler.step()
        net = train(epoch, net, trainloader, optimizer, device)
    net = net.to('cpu')
    state = net.state_dict()
    torch.save(state, 'lungseg_net2.pth')
def predict(weights_path, img):
    """
    Args:
    weights_path (str): path to where the weights are stored, which will be loaded unto the model
    img (PIL): Nuclei PIL image. Image size = 1000x1000

    Returns:
    Mask (PIL): PIL image detailing the predicted segmentation of the nuclei
    """

    # Prepare image
    transform = transforms.Compose([
        transforms.Pad(
            12),  # given image is 1000x1000, pad it to make it 1024x1024
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],
                             [0.229, 0.224, 0.225])  # imagenet normalization
    ])

    img = transform(img)
    img = torch.unsqueeze(img, dim=0)

    # Prepare model
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    num_class = 3
    model = ResNetUNet(num_class).to(device)
    model.load_state_dict(torch.load(weights_path))
    model.eval()

    # Predict
    with torch.no_grad():
        outputs = torch.sigmoid(model(img.to(device)))

    pred = outputs.to('cpu').detach().numpy()[0].transpose((1, 2, 0))
    mask = Image.fromarray(np.uint8(pred * 255)).convert('RGB')

    return mask
    print("train val dataset size {},{}".format(len(train_set), len(val_set)))

    # global image_datasets

    image_datasets = {'train': train_set, 'val': val_set}

    dataloaders = {
        'train':
        DataLoader(train_set,
                   batch_size=args['batch_size'],
                   shuffle=True,
                   num_workers=args['n_work'],
                   pin_memory=args['pin']),
        'val':
        DataLoader(val_set,
                   batch_size=args['batch_size'],
                   shuffle=False,
                   num_workers=args['n_work'],
                   pin_memory=args['pin'])
    }

    model = ResNetUNet(args['num_class']).to(device)

    model.load_state_dict(torch.load(MODEL_PATH))

    model = model.to(device)
    inp = torch.randn((2, 3, 320, 320)).to(device)
    model(inp)

    evalModel(model)
Exemple #7
0
    bboxes = restore_bboxes(cls, rho, theta)
    bboxes[:, :8] = bboxes[:, :8] * 4 * (
        w / img.width + h / img.height) / 2  # restore scale and resize
    return bboxes


if __name__ == '__main__':
    img_files = [
        os.path.join(cfg.test.dataset_test, img_file)
        for img_file in sorted(os.listdir(cfg.test.dataset_test))
    ]
    img_path = np.random.choice(img_files)
    print(img_path)

    model_path = cfg.test.model_pth
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model = ResNetUNet(backbone='50')
    model.load_state_dict(torch.load(model_path))
    model.to(device)

    model.eval()
    img = Image.open(img_path).convert('RGB')
    img_array = np.array(img)
    bboxes = detect_single_image(img, model, device)
    print(bboxes.shape)

    for bbox in bboxes:
        pts = bbox[:8].astype(np.int32).reshape((-1, 1, 2))
        cv2.polylines(img_array, [pts], True, (0, 215, 255), 3)
    Image.fromarray(img_array).save('res.png')
Exemple #8
0
    dataloaders = {
        'train':
        DataLoader(train_set,
                   batch_size=args['batch_size'],
                   shuffle=True,
                   num_workers=args['n_work'],
                   pin_memory=args['pin']),
        'val':
        DataLoader(val_set,
                   batch_size=args['batch_size'],
                   shuffle=False,
                   num_workers=args['n_work'],
                   pin_memory=args['pin'])
    }

    model = ResNetUNet(args['num_class']).to(device)

    # if args['initialization'] == 'xavier':
    # 	model._initialize_()
    # elif args['initialization'] == 'kaiming':
    # 	model._kaiming_initialize_()

    if args['optimizer_choice'] == "SGD":
        optimizer_ft = torch.optim.SGD(model.parameters(),
                                       lr=0.001,
                                       momentum=0.9)

    elif args['optimizer_choice'] == "Adam":
        optimizer_ft = torch.optim.Adam(filter(lambda p: p.requires_grad,
                                               model.parameters()),
                                        lr=1e-2)
Exemple #9
0
from tools.config import Cfg
from train import Trainer
from model import ResNetUNet
from torch.optim import Adam
import torch
config = Cfg.load_config_from_file('config/config.yml')
params = {
    'print_every': 10,
    'valid_every': 5 * 10,
    'iters': 100000,
    'n_classses': 2
}

config['trainer'].update(params)
model = ResNetUNet(config['trainer']['n_classes'])
optimizer = Adam(model.parameters(), lr=1e-6)
train = Trainer(config, model, optimizer, pretrained=False)
x, y = next(iter(train.train_data_loader))
train.train()
            avg_jaccard_score = jaccard_score / len(thresholds)
        else:
            avg_dice_score += dice_score / len(thresholds)
            avg_jaccard_score += jaccard_score / len(thresholds)

    avg_dice_score /= len(test_loader)
    avg_jaccard_score /= len(test_loader)

    return avg_dice_score, avg_jaccard_score


if __name__ == "__main__":

    # Set up model
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model = ResNetUNet(3)
    model.load_state_dict(torch.load("./weights_3channel_dropout_1"))
    model.to(device)

    # Test Set Loader
    label_path = "./Test/Labels/"
    img_path = "./Test/Images/"
    trans = transforms.Compose([
        transforms.Pad(
            12),  # given image is 1000x1000, pad it to make it 1024x1024
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],
                             [0.229, 0.224, 0.225])  # imagenet normalization
    ])
    test_set = NucleiDataset(img_path, label_path, transform=trans)
    test_loader = DataLoader(test_set,