def __init__(self,
checkpoint_path='checkpoint/checkpoint.pth',
imsize=512):
super().__init__()
self.conf_threshold = 0.6 # Any pixel below this confidence value will not be considered as a detection.
self.checkpoint_path = checkpoint_path
self.transform = transforms.Compose(
[transforms.Resize((imsize, imsize)),
transforms.ToTensor()])
# Create Model
self.numInputChannels = 3
self.num_of_classes = 2
self.output_stride = 8 # Chosen while training the model. Doesn't make diff in inference. Possible Values: 8, 16
self.net = deeplab_xception.DeepLabv3_plus(
nInputChannels=self.numInputChannels,
n_classes=self.num_of_classes,
os=self.output_stride,
pretrained=False)
self.net = nn.DataParallel(self.net)
self.net.load_state_dict(torch.load(self.checkpoint_path))
# Send model to Device (GPU or CPU)
self.device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
self.net.to(self.device)
self.net.eval()
save_dir_root = os.path.join(os.path.dirname(os.path.abspath(__file__)))
exp_name = os.path.dirname(os.path.abspath(__file__)).split('/')[-1]
if resume_epoch != 0:
runs = sorted(glob.glob(os.path.join(save_dir_root, 'run', 'run_*')))
run_id = int(runs[-1].split('_')[-1]) if runs else 0
else:
runs = sorted(glob.glob(os.path.join(save_dir_root, 'run', 'run_*')))
run_id = int(runs[-1].split('_')[-1]) + 1 if runs else 0
save_dir = os.path.join(save_dir_root, 'run', 'run_' + str(run_id))
# Network definition
net = deeplab_xception.DeepLabv3_plus(nInputChannels=3,
n_classes=21,
pretrained=True)
modelName = 'deeplabv3+'
criterion = utils.cross_entropy2d
if resume_epoch == 0:
print("Training deeplabv3+ from scratch...")
else:
print("Initializing weights from: {}...".format(
os.path.join(save_dir, 'models',
modelName + '_epoch-' + str(resume_epoch - 1) + '.pth')))
net.load_state_dict(
torch.load(os.path.join(
save_dir, 'models',
modelName + '_epoch-' + str(resume_epoch - 1) + '.pth'),
map_location=lambda storage, loc: storage)
p['lr'] = 1e-10 # Learning rate
p['wd'] = 5e-4 # Weight decay
p['momentum'] = 0.9 # Momentum
p['epoch_size'] = 2 # How many epochs to change learning rate
p['Model'] = 'deeplab' # Choose model: unet or deeplab
backbone = 'xception' # For deeplab only: Use xception or resnet as feature extractor,
num_of_classes = 2
imsize = 512 # 256 or 512
output_stride = 8 # 8 or 16, 8 is better. Controls output stride of the deeplab model, which increases resolution of convolutions.
numInputChannels = 3
# Network definition
if p['Model'] == 'deeplab':
if backbone == 'xception':
net = deeplab_xception.DeepLabv3_plus(nInputChannels=numInputChannels, n_classes=num_of_classes, os=output_stride, pretrained=True)
elif backbone == 'resnet':
net = deeplab_resnet.DeepLabv3_plus(nInputChannels=numInputChannels, n_classes=num_of_classes, os=output_stride, pretrained=True)
else:
raise NotImplementedError
modelName = 'deeplabv3plus-' + backbone
# Use the following optimizer
optimizer = optim.SGD(net.parameters(), lr=p['lr'], momentum=p['momentum'], weight_decay=p['wd'])
p['optimizer'] = str(optimizer)
# Use the following loss function
criterion = utils.cross_entropy2d
else:
raise NotImplementedError
save_dir_root = os.path.join(os.path.dirname(os.path.abspath(__file__)))
exp_name = os.path.dirname(os.path.abspath(__file__)).split('/')[-1]
if resume_epoch != 0:
runs = sorted(glob.glob(os.path.join(save_dir_root, 'run', 'run_*')))
run_id = int(runs[-1].split('_')[-1]) if runs else 0
else:
runs = sorted(glob.glob(os.path.join(save_dir_root, 'run', 'run_*')))
run_id = int(runs[-1].split('_')[-1]) + 1 if runs else 0
save_dir = os.path.join(save_dir_root, 'run')
# Network definition
if backbone == 'xception':
net = deeplab_xception.DeepLabv3_plus(nInputChannels=3,
n_classes=12,
os=16,
pretrained=False)
elif backbone == 'resnet':
net = deeplab_resnet.DeepLabv3_plus(nInputChannels=3,
n_classes=12,
os=16,
pretrained=False)
else:
raise NotImplementedError
modelName = 'deeplabv3plus-' + backbone + '-voc'
criterion = loss.dice_bce_loss()
# criterion = loss.FocalLoss(gamma=2)
print("Backbone network is {}".format(backbone))
p['momentum'] = 0.9 # Momentum
p['epoch_size'] = 40 # How many epochs to change learning rate
save_dir_root = os.path.join(os.path.dirname(os.path.abspath(__file__)))
exp_name = os.path.dirname(os.path.abspath(__file__)).split('/')[-1]
if resume_epoch == 0:
runs = sorted(glob.glob(os.path.join(save_dir_root, 'run', 'run_*')))
run_id = int(runs[-1].split('_')[-1]) + 1 if runs else 0
else:
run_id = 0
save_dir = os.path.join(save_dir_root, 'run', 'run_' + str(run_id))
# Network definition
net = deeplab_xception.DeepLabv3_plus(nInputChannels=3, n_classes=21)
modelName = 'deeplabv3+'
criterion = utils.cross_entropy2d
if resume_epoch == 0:
print("Training deeplabv3+ from scratch...")
else:
print("Initializing weights from: {}...".format(
os.path.join(save_dir, 'models', modelName + '_epoch-' + str(resume_epoch - 1) + '.pth')))
net.load_state_dict(
torch.load(os.path.join(save_dir, 'models', modelName + '_epoch-' + str(resume_epoch - 1) + '.pth'),
map_location=lambda storage, loc: storage)) # Load all tensors onto the CPU
if gpu_id >= 0:
torch.cuda.set_device(device=gpu_id)
net.cuda()