def video_test_without_train(self):
#base_path = '/home/zjh/code/vseg/dataset/tianchiyusai/'
#base_path = 'F:/vseg/dataset/media/tianchiyusai/tianchiyusai/'
base_path = self.base_path
set_type = 'test'
#set_type = 'val'
if set_type == 'test':
set_file_path = os.path.join(base_path, 'ImageSets', 'test.txt')
result_dir = os.path.join(base_path, 'result', 'test')
else:
set_file_path = os.path.join(base_path, 'ImageSets', 'val.txt')
result_dir = os.path.join('result', 'test')
if not os.path.exists(result_dir):
os.makedirs(result_dir)
with open(set_file_path) as f:
videos = [v.strip() for v in f.readlines()]
#bulit net and load weight
#model_path = 'pretrained/fparent_epoch-8.pth'
#model_path = 'offline_save_dir/lr_0.0005_wd_0.001/parent_epoch-1.pth'
model_path = self.model_path
state_dict = torch.load(model_path)
net = deeplab_resnet.Res_Deeplab_no_msc(2)
net.load_state_dict(state_dict)
if self.use_cuda:
torch.cuda.set_device(0)
net.cuda()
print('use cuda')
else:
print('use cpu')
net.eval()
start_time = datetime.datetime.now()
print('now we have {} videos to deal'.format(len(videos)))
with torch.no_grad():
for cnt, video in enumerate(videos):
print('Predict in {} {}/{}'.format(video, cnt + 1,
len(videos)))
first_mask_path = os.path.join(base_path, 'Annotations', video,
'00000.png')
img_dir = os.path.join(base_path, 'JPEGImages', video)
save_dir = os.path.join(result_dir, video)
if not os.path.exists(save_dir):
os.makedirs(save_dir)
shutil.copyfile(first_mask_path,
os.path.join(save_dir, '00000.png'))
self.test_without_train(net,
first_mask_path,
img_dir,
save_dir=save_dir,
use_cuda=self.use_cuda)
present_time = datetime.datetime.now()
remain = (present_time - start_time) * (len(videos) - cnt -
1) / (cnt + 1)
print('{} remain {}'.format(present_time, remain))
batch_size = int(3)
vbatch_size = 3
db_root_dir = Path.db_offline_train_root_dir()
nAveGrad = 4 # keep it even
save_dir = os.path.join(Path.save_offline_root_dir(),
'lr_' + str(base_lr) + '_wd_' + str(weight_decay))
if not os.path.exists(save_dir):
os.makedirs(os.path.join(save_dir))
learnRate = base_lr
######################################################################################################################
"""Initialise the network"""
net = deeplab_resnet.Res_Deeplab_no_msc(int(NoLabels))
net.float()
modelName = 'parent'
if use_cuda:
torch.cuda.set_device(0)
if resume_epoch == 0:
saved_state_dict = torch.load(
'pretrained/MS_DeepLab_resnet_pretrained_COCO_init.pth')
"""
Last layer structure is modified according to the number of segmentation classes.
For Mastrack, #classes=2
"""
# Training dataset and its iterator
db_train[i] = db.DAVIS17OnlineDataset(train=True,
inputRes=(480, 854),
db_root_dir=db_root_dir,
transform=composed_transforms,
seq_name=seq_name,
noIterations=noIterations,
object_id=i)
trainloader[i] = DataLoader(db_train[i],
batch_size=p['trainBatch'],
shuffle=False,
num_workers=2)
nets[i] = deeplab_resnet.Res_Deeplab_no_msc(int(args['--NoLabels']))
nets[i].float()
nets[i].train()
# Let us make it run on multiple GPUs!
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
nets[i] = nn.DataParallel(nets[i])
if torch.cuda.is_available():
nets[i].cuda()
print("Updating weights from: {}".format(
os.path.join(db_parent_root_dir,
'lr_' + str(parent_lr) + '_wd_' + str(parent_wd),
'parent_epoch-' + str(resume_epoch_parent) + '.pth')))
import torch
from torch.utils.data import DataLoader
import torch.backends.cudnn as cudnn
import torch.optim as optim
from utility_functions import *
import deeplab_resnet
from get import *
net = deeplab_resnet.Res_Deeplab_no_msc(2)
net.float()
saved_state_dict = torch.load(
'pretrained/MS_DeepLab_resnet_pretrained_COCO_init.pth')
for i in saved_state_dict:
i_parts = i.split('.')
if i_parts[1] == 'layer5':
saved_state_dict[i] = net.state_dict()[i]
if i_parts[1] == 'conv1':
saved_state_dict[i] = torch.cat(
(saved_state_dict[i], torch.FloatTensor(64, 1, 7, 7).normal_(
0, 0.0001)), 1)
net.load_state_dict(saved_state_dict)
# Let us make it run on multiple GPUs!
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
net = nn.DataParallel(net)
