class RAN():
def __init__(self, weight, gpu_ids):
self.model = DeepLab(num_classes=2,
backbone='mobilenet',
output_stride=16)
torch.cuda.set_device(gpu_ids)
self.model = self.model.cuda()
assert weight is not None
if not os.path.isfile(weight):
raise RuntimeError("=> no checkpoint found at '{}'".format(weight))
checkpoint = torch.load(weight)
self.model.load_state_dict(checkpoint['state_dict'])
self.model.eval()
self.mean = (0.485, 0.456, 0.406)
self.std = (0.229, 0.224, 0.225)
def inference(self, img):
# normalize
img = cv2.resize(img, (480, 480))
img = img.astype(np.float32)
img /= 255.0
img -= self.mean
img /= self.std
img = img.transpose((2, 0, 1))
img = img[np.newaxis, :, :, :]
# to tensor
img = torch.from_numpy(img).float().cuda()
with torch.no_grad():
output = self.model(img)
return output
def main():
here = osp.dirname(osp.abspath(__file__))
trainOpts = TrainOptions()
args = trainOpts.get_arguments()
now = datetime.datetime.now()
args.out = osp.join(
here, 'results',
args.model + '_' + args.dataset + '_' + now.strftime('%Y%m%d__%H%M%S'))
if not osp.isdir(args.out):
os.makedirs(args.out)
log_file = osp.join(args.out, args.model + '_' + args.dataset + '.log')
mylog = open(log_file, 'w')
checkpoint_dir = osp.join(args.out, 'checkpoints')
os.makedirs(checkpoint_dir)
os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)
cuda = torch.cuda.is_available()
torch.manual_seed(1337)
if cuda:
torch.cuda.manual_seed(1337)
# 1. dataset
# MAIN_FOLDER = args.folder + 'Vaihingen/'
# DATA_FOLDER = MAIN_FOLDER + 'top/top_mosaic_09cm_area{}.tif'
# LABEL_FOLDER = MAIN_FOLDER + 'gts_for_participants/top_mosaic_09cm_area{}.tif'
# train_ids = ['1', '3', '5', '21','23', '26', '7', '13', '17', '32', '37']
# val_ids =['11','15', '28', '30', '34']
# train_set = ISPRS_dataset(train_ids, DATA_FOLDER, LABEL_FOLDER,cache=args.cache)
# train_loader = torch.utils.data.DataLoader(train_set,batch_size=args.batch_size)
# MAIN_FOLDER = args.folder + 'Potsdam_multiscale/'
# DATA_FOLDER1 = MAIN_FOLDER + '3_Ortho_IRRG/top_potsdam_{}_IRRG.tif'
# LABEL_FOLDER1 = MAIN_FOLDER + '5_Labels_for_participants/top_potsdam_{}_label.tif'
# DATA_FOLDER2 = MAIN_FOLDER + '3_Ortho_IRRG_2/top_potsdam_{}_IRRG.tif'
# LABEL_FOLDER2 = MAIN_FOLDER + '5_Labels_for_participants_2/top_potsdam_{}_label.tif'
# train_ids=['2_10','3_10','3_11','3_12','4_11','4_12','5_10','5_12',\
# '6_8','6_9','6_10','6_11','6_12','7_7','7_9','7_11','7_12']
# val_ids=[ '2_11', '2_12', '4_10', '5_11', '6_7', '7_8', '7_10']
# target_set = ISPRS_dataset_multi(2,train_ids, DATA_FOLDER1, LABEL_FOLDER1,DATA_FOLDER2, LABEL_FOLDER2,cache=args.cache)
# target_loader = torch.utils.data.DataLoader(target_set,batch_size=args.batch_size)
# MAIN_FOLDER = args.folder + 'Potsdam/'
# DATA_FOLDER = MAIN_FOLDER + '3_Ortho_IRRG/top_potsdam_{}_IRRG.tif'
# LABEL_FOLDER = MAIN_FOLDER + '5_Labels_for_participants/top_potsdam_{}_label.tif'
# ERODED_FOLDER = MAIN_FOLDER + '5_Labels_for_participants_no_Boundary/top_potsdam_{}_label_noBoundary.tif'
# train_ids=['2_10','3_10','3_11','3_12','4_11','4_12','5_10','5_12',\
# '6_8','6_9','6_10','6_11','6_12','7_7','7_9','7_11','7_12']
# val_ids=[ '2_11', '2_12', '4_10', '5_11', '6_7', '7_8', '7_10']
# train_set = ISPRS_dataset(train_ids, DATA_FOLDER, LABEL_FOLDER,cache=args.cache)
# train_loader = torch.utils.data.DataLoader(train_set,batch_size=args.batch_size)
# MAIN_FOLDER = args.folder + 'Vaihingen/'
# DATA_FOLDER = MAIN_FOLDER + 'top/top_mosaic_09cm_area{}.tif'
# LABEL_FOLDER = MAIN_FOLDER + 'gts_for_participants/top_mosaic_09cm_area{}.tif'
# train_ids = ['1', '3', '5', '21','23', '26', '7', '13', '17', '32', '37']
# val_ids =['11','15', '28', '30', '34']
# train_set = ISPRS_dataset(train_ids, DATA_FOLDER, LABEL_FOLDER,cache=args.cache)
# train_loader = torch.utils.data.DataLoader(train_set,batch_size=args.batch_size)
MAIN_FOLDER = args.folder + 'DeepGlobe/land-train_crop/'
DATA_FOLDER = MAIN_FOLDER + '{}_sat.jpg'
LABEL_FOLDER = MAIN_FOLDER + '{}_mask.png'
all_files = sorted(glob(DATA_FOLDER.replace('{}', '*')))
all_ids = [f.split('/')[-1].split('_')[0] for f in all_files]
train_ids = all_ids[:int(len(all_ids) / 3 * 2)]
val_ids = all_ids[int(len(all_ids) / 3 * 2):]
train_set = DeepGlobe_dataset(train_ids, DATA_FOLDER, LABEL_FOLDER)
train_loader = torch.utils.data.DataLoader(train_set,
batch_size=args.batch_size)
MAIN_FOLDER = args.folder + 'ISPRS_dataset/Vaihingen/'
DATA_FOLDER1 = MAIN_FOLDER + 'top/top_mosaic_09cm_area{}.tif'
LABEL_FOLDER1 = MAIN_FOLDER + 'gts_for_participants/top_mosaic_09cm_area{}.tif'
DATA_FOLDER2 = MAIN_FOLDER + 'resized_resolution5/top_mosaic_09cm_area{}.tif'
LABEL_FOLDER2 = MAIN_FOLDER + 'gts_for_participants5/top_mosaic_09cm_area{}.tif'
train_ids = ['1', '3', '5', '21', '23', '26', '7', '13', '17', '32', '37']
val_ids = ['11', '15', '28', '30', '34']
target_set = ISPRS_dataset_multi(5,
train_ids,
DATA_FOLDER1,
LABEL_FOLDER1,
DATA_FOLDER2,
LABEL_FOLDER2,
cache=args.cache)
target_loader = torch.utils.data.DataLoader(target_set,
batch_size=args.batch_size)
# val_set = ISPRS_dataset(val_ids, DATA_FOLDER1, LABEL_FOLDER1,cache=args.cache)
# val_loader = torch.utils.data.DataLoader(val_set,batch_size=args.batch_size)
LABELS = [
"roads", "buildings", "low veg.", "trees", "cars", "clutter", "unknown"
] # Label names
N_CLASS = len(LABELS) # Number of classes
# 2. model
if args.backbone == 'resnet':
model = DeepLab(num_classes=N_CLASS,
backbone=args.backbone,
output_stride=args.out_stride,
sync_bn=args.sync_bn,
freeze_bn=args.freeze_bn)
elif args.backbone == 'resnet_multiscale':
model = DeepLabCA(num_classes=N_CLASS,
backbone=args.backbone,
output_stride=args.out_stride,
sync_bn=args.sync_bn,
freeze_bn=args.freeze_bn)
else:
print('backbone not exists!')
train_params = [{
'params': model.get_1x_lr_params(),
'lr': args.lr
}, {
'params': model.get_10x_lr_params(),
'lr': args.lr * 10
}]
start_epoch = 0
start_iteration = 0
# 3. optimizer
lr = args.lr
# optimizer = torch.optim.SGD(model.parameters(), lr=args.lr,
# momentum=args.momentum, weight_decay=args.weight_decay)
netD_domain = FCDiscriminator(num_classes=N_CLASS)
netD_scale = FCDiscriminator(num_classes=N_CLASS)
optim_netG = torch.optim.SGD(train_params,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=args.nesterov)
optim_netD_domain = optim.Adam(netD_domain.parameters(),
lr=args.lr_D,
betas=(0.9, 0.99))
optim_netD_scale = optim.Adam(netD_scale.parameters(),
lr=args.lr_D,
betas=(0.9, 0.99))
if cuda:
model, netD_domain, netD_scale = model.cuda(), netD_domain.cuda(
), netD_scale.cuda()
if args.resume:
checkpoint = torch.load(args.resume)
model.load_state_dict(checkpoint)
bce_loss = torch.nn.BCEWithLogitsLoss()
# 4. training
iter_ = 0
no_optim = 0
val_best_loss = float('Inf')
factor = 10
max_iter = 50000
trainloader_iter = enumerate(train_loader)
targetloader_iter = enumerate(target_loader)
source_label = 0
target_label = 1
source_scale_label = 0
target_scale_label = 1
train_loss = []
train_acc = []
target_acc_s1 = []
target_acc_s2 = []
while iter_ < max_iter:
optim_netG.zero_grad()
adjust_learning_rate(optim_netG, iter_, args)
optim_netD_domain.zero_grad()
optim_netD_scale.zero_grad()
adjust_learning_rate_D(optim_netD_domain, iter_, args)
adjust_learning_rate_D(optim_netD_scale, iter_, args)
if iter_ % 1000 == 0:
train_loss = []
train_acc = []
target_acc_s1 = []
target_acc_s2 = []
for param in netD_domain.parameters():
param.requires_grad = False
for param in netD_scale.parameters():
param.requires_grad = False
_, batch = trainloader_iter.__next__()
im_s, label_s = batch
_, batch = targetloader_iter.__next__()
im_t_s1, label_t_s1, im_t_s2, label_t_s2 = batch
if cuda:
im_s, label_s = im_s.cuda(), label_s.cuda()
im_t_s1, label_t_s1, im_t_s2, label_t_s2 = im_t_s1.cuda(
), label_t_s1.cuda(), im_t_s2.cuda(), label_t_s2.cuda()
############
#TRAIN NETG#
############
#train with source
#optimize segmentation network with source data
pred_seg = model(im_s)
seg_loss = cross_entropy2d(pred_seg, label_s)
seg_loss /= len(im_s)
loss_data = seg_loss.data.item()
if np.isnan(loss_data):
# continue
raise ValueError('loss is nan while training')
seg_loss.backward()
# import pdb
# pdb.set_trace()
pred = np.argmax(pred_seg.data.cpu().numpy()[0], axis=0)
gt = label_s.data.cpu().numpy()[0]
train_acc.append(accuracy(pred, gt))
train_loss.append(loss_data)
#train with target
pred_s1 = model(im_t_s1)
pred = np.argmax(pred_s1.data.cpu().numpy()[0], axis=0)
gt = label_t_s1.data.cpu().numpy()[0]
target_acc_s1.append(accuracy(pred, gt))
pred_s2 = model(im_t_s2)
pred = np.argmax(pred_s2.data.cpu().numpy()[0], axis=0)
gt = label_t_s2.data.cpu().numpy()[0]
target_acc_s2.append(accuracy(pred, gt))
pred_d = netD_domain(F.softmax(pred_s1))
pred_s = netD_scale(F.softmax(pred_s2))
loss_adv_domain = bce_loss(
pred_d,
Variable(
torch.FloatTensor(
pred_d.data.size()).fill_(source_label)).cuda())
loss_adv_scale = bce_loss(
pred_s,
Variable(
torch.FloatTensor(
pred_s.data.size()).fill_(source_scale_label)).cuda())
loss = args.lambda_adv_domain * loss_adv_domain + args.lambda_adv_scale * loss_adv_scale
loss /= len(im_t_s1)
loss.backward()
############
#TRAIN NETD#
############
for param in netD_domain.parameters():
param.requires_grad = True
for param in netD_scale.parameters():
param.requires_grad = True
#train with source domain and source scale
pred_seg, pred_s1 = pred_seg.detach(), pred_s1.detach()
pred_d = netD_domain(F.softmax(pred_seg))
# pred_s=netD_scale(F.softmax(pred_seg))
pred_s = netD_scale(F.softmax(pred_s1))
loss_D_domain = bce_loss(
pred_d,
Variable(
torch.FloatTensor(
pred_d.data.size()).fill_(source_label)).cuda())
loss_D_scale = bce_loss(
pred_s,
Variable(
torch.FloatTensor(
pred_s.data.size()).fill_(source_scale_label)).cuda())
loss_D_domain = loss_D_domain / len(im_s) / 2
loss_D_scale = loss_D_scale / len(im_s) / 2
loss_D_domain.backward()
loss_D_scale.backward()
#train with target domain and target scale
pred_s1, pred_s2 = pred_s1.detach(), pred_s2.detach()
pred_d = netD_domain(F.softmax(pred_s1))
pred_s = netD_scale(F.softmax(pred_s2))
loss_D_domain = bce_loss(
pred_d,
Variable(
torch.FloatTensor(
pred_d.data.size()).fill_(target_label)).cuda())
loss_D_scale = bce_loss(
pred_s,
Variable(
torch.FloatTensor(
pred_s.data.size()).fill_(target_scale_label)).cuda())
loss_D_domain = loss_D_domain / len(im_s) / 2
loss_D_scale = loss_D_scale / len(im_s) / 2
loss_D_domain.backward()
loss_D_scale.backward()
optim_netG.step()
optim_netD_domain.step()
optim_netD_scale.step()
if iter_ % 100 == 0:
print(
'Train [{}/{} Source loss:{:.6f} acc:{:.4f} % Target s1 acc:{:4f}% Target s2 acc:{:4f}%]'
.format(iter_, max_iter,
sum(train_loss) / len(train_loss),
sum(train_acc) / len(train_acc),
sum(target_acc_s1) / len(target_acc_s1),
sum(target_acc_s2) / len(target_acc_s2)))
print(
'Train [{}/{} Source loss:{:.6f} acc:{:.4f} % Target s1 acc:{:4f}% Target s2 acc:{:4f}%]'
.format(iter_, max_iter,
sum(train_loss) / len(train_loss),
sum(train_acc) / len(train_acc),
sum(target_acc_s1) / len(target_acc_s1),
sum(target_acc_s2) / len(target_acc_s2)),
file=mylog)
if iter_ % 1000 == 0:
print('saving checkpoint.....')
torch.save(model.state_dict(),
osp.join(checkpoint_dir, 'iter{}.pth'.format(iter_)))
iter_ += 1
save_seg_name = '/model.pth'
save_act_name = '/action.pth'
# segmentation
import torch
from modeling.deeplab import DeepLab
n_class = 10
try:
model = DeepLab(num_classes=n_class,
backbone='xception',
output_stride=16,
sync_bn=bool(None),
freeze_bn=bool(False))
model = model.cuda()
checkpoint = torch.load(now_dir + cp_seg_name)
model.load_state_dict(checkpoint['state_dict'])
torch.save(model, now_dir + target_dir + save_seg_name)
print('segmentation model - OK!')
except:
print('segmentation model - Failed!')
# action
import torchvision.models as models
import torch.nn as nn
try:
model = models.resnext50_32x4d()
model.fc = nn.Sequential(nn.Linear(2048, 2048),