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)
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
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)
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')
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)
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,