def test_fcn(self):
x = Variable(torch.randn(BATCH_SIZE, 3, 224, 224).fill_(1.0))
self.exportTest(toC(
fcn_resnet101(pretrained=False, pretrained_backbone=False)),
toC(x),
rtol=1e-3,
atol=1e-5)
def get_model(model_name: str, device, num_classes: int, pretrained: bool = False) -> torch.nn:
"""
Gets the model requested from the config file
:param model_name: name of model used in the program
:param device: name of device model is hosted on
:param num_classes: number of classes in the model (including background 0,0,0)
:param pretrained: boolean of starting on a pretrained model
:return: the nn.model to be trained
"""
print(pretrained)
model_dict = {'fcn_resnet50': models.fcn_resnet50(pretrained=pretrained, num_classes=num_classes).to(device).eval(),
'fcn_resnet101': models.fcn_resnet101(pretrained=pretrained,
num_classes=num_classes).to(device).eval(),
'deeplabv3_resnet50': models.deeplabv3_resnet50(pretrained=pretrained,
num_classes=num_classes).to(device).eval(),
'deeplabv3_resnet101': models.deeplabv3_resnet101(pretrained=pretrained,
num_classes=num_classes).to(device).eval(),
'deeplabv3_mobilenet_v3_large': models.deeplabv3_mobilenet_v3_large(pretrained=pretrained,
num_classes=num_classes).to(
device).eval(),
'lraspp_mobilenet_v3_large': models.lraspp_mobilenet_v3_large(pretrained=pretrained,
num_classes=num_classes).to(
device).eval()
}
try:
model = model_dict[model_name]
except KeyError:
print(f"KeyError, model_name is not valid allowable names are: {model_dict.keys()}")
model = None
model = model.eval()
return model
def __init__(self):
super(Encoder101, self).__init__()
self.resnet = deeplabv3_resnet101(pretrained=True)
self.conv1_p = nn.Conv2d(1,
64,
kernel_size=7,
stride=2,
padding=3,
bias=True)
resnet = fcn_resnet101(pretrained=True)
self.conv1 = resnet.backbone.conv1
self.bn1 = resnet.backbone.bn1
self.relu = resnet.backbone.relu # 1/2, 64
self.maxpool = resnet.backbone.maxpool
self.res2 = resnet.backbone.layer1 # 1/4, 256
self.res3 = resnet.backbone.layer2 # 1/8, 512
self.res4 = resnet.backbone.layer3 # 1/16, 1024
self.res5 = resnet.backbone.layer4 # 1/32, 2048
self.register_buffer(
'mean',
torch.FloatTensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
self.register_buffer(
'std',
torch.FloatTensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
def load_seg_model(cfg, load_model_only=False):
modeltype = cfg['training']['model']
if modeltype == 'fcn_resnet101':
n_channels = len(cfg['data']['input_classes'])
n_outputs = len(cfg['data']['output_classes'])
model = fcn_resnet101(pretrained=False, num_classes=n_outputs)
model.backbone.conv1 = nn.Conv2d(n_channels, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
elif modeltype == 'deeplabv3_resnet101':
n_channels = len(cfg['data']['input_classes'])
n_outputs = len(cfg['data']['output_classes'])
aux_loss = cfg['training']['aux_loss'] is not False
model = deeplabv3_resnet101(pretrained=False, num_classes=n_outputs, aux_loss=aux_loss)
model.backbone.conv1 = nn.Conv2d(n_channels, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
elif modeltype == 'higher_hrnet':
hrnet_cfg = get_hrnet_cfg(cfg)
model = get_hrnet_model(hrnet_cfg)
else:
raise ValueError()
if modelpath := cfg['resume'].get('path', None):
state = torch.load(modelpath)
model.load_state_dict(state['state_dict'])
if load_model_only:
return model
starting_epoch = state['epoch']
if conf_epoch := cfg['resume'].get('epoch', None):
print(
f"WARNING: Loaded model trained for {starting_epoch - 1} epochs but config explicitly overrides to {conf_epoch}")
starting_epoch = conf_epoch
def resume(self, file, test=False):
import torch
if test and not file:
self.fcn = fcn_resnet101(pretrained=True, num_classes=2)
return
if file:
print('Loading checkpoint from: ' + file)
checkpoint = torch.load(file)
checkpoint = checkpoint['model_state_dict']
self.load_state_dict(checkpoint)
def __init__(self,
n_channels=3,
n_classes=21,
softmax_out=True,
resnet_type=101,
pretrained=False):
super(FCN, self).__init__()
self.resnet_type = resnet_type
self.n_channels = n_channels
self.n_classes = n_classes
self.pretrained = pretrained
# Input conv is applied to convert the input to 3 ch depth
self.inconv = None
if n_channels != 3:
self.inconv = FwdConv(n_channels, 3, kernel_size=1, padding=0)
# Pre-trained model needs to be an identical network
if pretrained:
mid_classes = 21
else:
mid_classes = n_classes
# Maind body
if resnet_type == 50:
self.fcn_body = fcn_resnet50(pretrained=False,
num_classes=mid_classes)
self.pretrained = False
else:
self.fcn_body = fcn_resnet101(pretrained=pretrained,
num_classes=mid_classes)
if n_classes != 21:
self.fcn_body.classifier[-1] = nn.Conv2d(512,
n_classes,
kernel_size=(1, 1),
stride=(1, 1))
if self.fcn_body.aux_classifier != None:
self.fcn_body.aux_classifier[-1] = nn.Conv2d(512,
n_classes,
kernel_size=(1,
1),
stride=(1, 1))
# Softmax alternative
self.has_softmax = softmax_out
if softmax_out:
self.softmax = nn.Softmax2d()
else:
self.softmax = None
def __init__(self,
end_module,
keyword='out',
pretrain=False,
freeze=False):
super(ResNet101_FCN, self).__init__()
self.resnet = models.fcn_resnet101(pretrained=pretrain,
progress=True,
num_classes=21)
if freeze:
for module in self.resnet.parameters():
module.requires_grad = False
print("module frozen.")
self.end_module = end_module
self.output = None
self.key = keyword
def __init__(self):
self.model = seg.fcn_resnet101(pretrained=True, progress=True)
self.model.cuda()
self.model.eval()
def __init__(self, n_class=21):
super(FCN, self).__init__()
self.fcn = fcn_resnet101(pretrained=False, num_classes=n_class)
def fcn_resnet101(num_classes):
"""Return torchvision fcn_resnet101 and wrap it in SingleOutNet."""
return SingleOutNet(
seg_models.fcn_resnet101(num_classes=num_classes, aux_loss=True))
def train(opt):
# Specify the Model Architecture
if opt.architecture.lower() == "deeplabv3_resnet50":
model = models.deeplabv3_resnet50(pretrained=True, progress=True)
model.classifier = models.deeplabv3.DeepLabHead(2048, 4)
elif opt.architecture.lower() == "deeplabv3_resnet101":
model = models.deeplabv3_resnet101(pretrained=True, progress=True)
model.classifier = models.deeplabv3.DeepLabHead(2048, 4)
elif opt.architecture.lower() == "fcn_resnet50":
model = models.fcn_resnet50(pretrained=True, progress=True)
model.classifier = models.fcn.FCNHead(2048, 4)
elif opt.architecture.lower() == "fcn_resnet101":
model = models.fcn_resnet101(pretrained=True, progress=True)
model.classifier = models.fcn.FCNHead(2048, 4)
# Define Optimizer
if opt.optim.lower() == "adam":
modelOptim = torch.optim.Adam(model.parameters(), lr=opt.lr)
elif opt.optim.lower() == "sgd":
modelOptim = torch.optim.SGD(model.parameters(), lr=opt.lr)
# Define Loss Function
lossFnc = torch.nn.CrossEntropyLoss()
# Set Training and Validation Datasets
dataTransforms = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
segdata = {
x: Dataset(root=Path(opt.data_path) / x,
imageFolder="images",
maskFolder="masks",
transforms=dataTransforms)
for x in ["train", "valid"]
}
dataLoaders = {
x: DataLoader(segdata[x],
batch_size=opt.batch_size,
shuffle=True,
num_workers=opt.num_workers)
for x in ["train", "valid"]
}
# Set Training Device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)
# Create and Initialize Training Log File
#perfMetrics = {"f1-score": f1_score, "auroc": roc_auc_score}
perfMetrics = {"f1-score": f1_score}
fieldnames = ['epoch', 'train_loss', 'valid_loss'] + \
[f'train_{m}' for m in perfMetrics.keys()] + \
[f'valid_{m}' for m in perfMetrics.keys()]
with open(os.path.join('log.csv'), 'w', newline='') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writeheader()
#model = torch.load('trainedModel25.pth')
# Train
startTimer = time.time()
for epoch in range(1, opt.num_epochs + 1):
print('-' * 60)
print("Epoch: {}/{}".format(epoch, opt.num_epochs))
batchsummary = {a: [0] for a in fieldnames}
for phase in ["train", "valid"]:
if phase == 'train':
model.train()
else:
model.eval()
# Iterate over data.
for sample in tqdm(iter(dataLoaders[phase]), file=sys.stdout):
inputs = sample['image'].to(device)
masks = sample['mask'].to(device)
# zero the parameter gradients
modelOptim.zero_grad()
# track history if only in train
with torch.set_grad_enabled(phase == 'train'):
outputs = model(inputs)
loss = lossFnc(outputs['out'], masks)
y_pred = outputs['out'].data.cpu().numpy().ravel()
y_true = masks.data.cpu().numpy().ravel()
for name, metric in perfMetrics.items():
if name == 'f1-score':
# Use a classification threshold of 0.1
f1Classes = np.zeros(4)
nPixels = np.zeros(4)
for classID in range(4):
f1Classes[classID] = metric(
y_true == classID,
y_pred[classID * len(y_true):
(classID + 1) * len(y_true)] > 0.1)
nPixels[classID] = np.count_nonzero(
y_true == classID)
f1weights = nPixels / (np.sum(nPixels))
f1 = np.matmul(f1Classes, f1weights)
batchsummary[f'{phase}_{name}'].append(f1)
else:
batchsummary[f'{phase}_{name}'].append(
metric(y_true.astype('uint8'), y_pred))
# backward + optimize only if in training phase
if phase == 'train':
loss.backward()
modelOptim.step()
batchsummary['epoch'] = epoch
epoch_loss = loss
batchsummary[f'{phase}_loss'] = epoch_loss.item()
for field in fieldnames[3:]:
batchsummary[field] = np.mean(batchsummary[field])
print((f'train loss: {batchsummary["train_loss"]: .4f}, '
f'valid loss: {batchsummary["valid_loss"]: .4f}, '
f'train f1-score: {batchsummary["train_f1-score"]: .4f}, '
f'valid f1-score: {batchsummary["valid_f1-score"]: .4f}, '))
#f'train auroc: {batchsummary["train_auroc"]: .4f}, '
#f'valid auroc: {batchsummary["valid_auroc"]: .4f}, '))
with open(os.path.join('log.csv'), 'a', newline='') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writerow(batchsummary)
# deep copy the model
if phase == 'valid' and loss < 1e10:
best_loss = loss
best_model_wts = copy.deepcopy(model.state_dict())
time_elapsed = time.time() - startTimer
print('-' * 60)
print('Training completed in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
print(f'Lowest validation loss: {best_loss: .4f}')
# load best model weights
model.load_state_dict(best_model_wts)
def test_fcn(self):
x = Variable(torch.randn(BATCH_SIZE, 3, 224, 224).fill_(1.0))
self.exportTest(toC(fcn_resnet101()), toC(x), rtol=1e-3, atol=1e-5)
def predict(config_file, model_file, output_dir, n_classes, architecture,
multi_crop=False, aux_loss=None):
pred_dir = Path(output_dir) / 'predictions'
col_dir = Path(output_dir) / 'colorized'
pred_dir.mkdir(exist_ok=True)
col_dir.mkdir(exist_ok=True)
if architecture == 'deeplab':
model = deeplabv3_resnet101(pretrained=False, progress=True,
num_classes=n_classes, aux_loss=aux_loss)
elif architecture == 'fcn':
model = fcn_resnet101(pretrained=False, progress=True,
num_classes=n_classes, aux_loss=aux_loss)
else:
raise ValueError('Unknown architecture specified. '
'Please choose either \'fcn\' or \'deeplab\'.')
model.load_state_dict(torch.load(model_file))
model.eval()
model.cuda()
criterion = CrossEntropyLoss(ignore_index=255)
test_loader, data_list = get_val_loader(config_file)
data_loader = tqdm(test_loader, total=len(test_loader), ncols=80)
color_palette = np.loadtxt(
'semseg/data/cityscapes/cityscapes_colors.txt').astype('uint8')
if not multi_crop:
for i, (img, label) in enumerate(data_loader):
img_path, label_path = data_list[i]
img = torch.cuda.FloatTensor(img.numpy())
label = torch.cuda.FloatTensor(label.numpy()).long()
output = model(img)['out']
loss = criterion(output, label)
data_loader.set_postfix(loss=loss.item())
pred_path = pred_dir / Path(img_path).name
col_path = col_dir / Path(img_path).name
# Apply argmax to prediction image. The resulting indices
# represent an 8-bit image with each pixel value corresponding
# to the respective class label
values, indices = torch.max(torch.squeeze(output), 0)
pred = indices.cpu().numpy()
color_pred = colorize(pred, color_palette)
cv2.imwrite(str(pred_path), pred)
color_pred.save(str(col_path))
else:
test_args = config.load_cfg_from_cfg_file(config_file)
mean, std = get_mean_std(test_args.mean_std)
crop_h = test_args.train_h
crop_w = test_args.train_w
logger = get_logger()
predict_multi_crop(logger, test_loader, data_list, model, n_classes,
mean, test_args.base_size,
crop_h, crop_w, test_args.scales, pred_dir, col_dir,
color_palette)
type=str,
default='cityscapes_deeplab.yaml',
help='config file')
args0 = parser.parse_args()
cudnn.benchmark = True
args, data_loader = get_train_loader(args0.config)
if args.architecture == 'deeplab':
model = deeplabv3_resnet101(pretrained=False,
progress=True,
num_classes=args.classes,
aux_loss=aux_loss)
elif args.architecture == 'fcn':
model = fcn_resnet101(pretrained=False,
progress=True,
num_classes=args.classes,
aux_loss=aux_loss)
else:
raise ValueError('Unknown architecture specified. '
'Please choose either \'fcn\' or \'deeplab\'.')
model.cuda()
print(model)
criterion = CrossEntropyLoss(ignore_index=args.ignore_label)
optimizer = optim.SGD(model.parameters(),
lr=args.base_lr,
momentum=args.momentum)
loss_log = []
for epoch in range(args.epochs):
def __init__(self):
super(FCNResNet101Model, self).__init__()
self.model = fcn_resnet101(pretrained=True, progress=False)
self.MEAN = torch.tensor([0.485, 0.456, 0.406])
self.STD = torch.tensor([0.229, 0.224, 0.225])
#%%
import torch
from torchvision.models import segmentation
from PIL import Image
import numpy as np
from matplotlib import pyplot as plt
#%%
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
#%%
fcn_res101 = segmentation.fcn_resnet101(pretrained=True)
fcn_res101.to(device)
#%%
print(fcn_res101)
# %%
img = Image.open("1.jpg")
img = np.array(img)
img_tensor = torch.from_numpy(img.astype(np.float32))
img_tensor = torch.unsqueeze(img_tensor,0)
img_tensor = img_tensor.permute(0,3,1,2)
img_tensor = img_tensor.to(device)
#%%
predict = fcn_res101(img_tensor)
#%%
pre_numpy = predict['out'].detach().numpy()
pre_numpy_person = pre_numpy[0,15,:,:]
plt.imshow(pre_numpy_person)
plt.show()
