def test(opt):
input_list_path = os.path.join(opt.data_path, "VOC{}".format(opt.year),
"ImageSets/Segmentation/{}.txt".format(opt.test_set))
image_ids = [id.strip() for id in open(input_list_path)]
output_folder = os.path.join(opt.output, "VOC{}_{}".format(opt.year, opt.test_set))
if os.path.isdir(output_folder):
shutil.rmtree(output_folder)
os.makedirs(output_folder)
if torch.cuda.is_available():
if opt.pre_trained_model_type == "model":
model = torch.load(opt.pre_trained_model_path)
else:
model = Deeplab(num_classes=21)
model.load_state_dict(torch.load(opt.pre_trained_model_path))
else:
if opt.pre_trained_model_type == "model":
model = torch.load(opt.pre_trained_model_path, map_location=lambda storage, loc: storage)
else:
model = Deeplab(num_classes=21)
model.load_state_dict(torch.load(opt.pre_trained_model_path, map_location=lambda storage, loc: storage))
model.eval()
for id in image_ids:
print (id)
img = np.zeros((513, 513, 3))
image_path = os.path.join(opt.data_path, "VOC{}".format(opt.year), "JPEGImages", "{}.jpg".format(id))
image = cv2.imread(image_path).astype(np.float32)
shutil.copy(os.path.join(opt.data_path, "VOC{}".format(opt.year), "JPEGImages", "{}.jpg".format(id)), output_folder + os.sep + id + ".jpg")
h,w,_ = image.shape
image[:, :, 0] -= 104.008
image[:, :, 1] -= 116.669
image[:, :, 2] -= 122.675
img[:h, :w, :] = image
img = np.transpose(np.array(img, dtype=np.float32), (2, 0, 1))
img = img[None, :, :, :]
img = torch.Tensor(img)
if torch.cuda.is_available():
img = img.cuda()
with torch.no_grad():
result = model(img)
interp = nn.UpsamplingBilinear2d(size=(513, 513))
output = interp(result[3]).cpu().data[0].numpy()
output = output[:, :h, :w]
result = np.argmax(output.transpose(1,2,0), axis=2).astype(np.uint8)
pred_colour = np.zeros((3, result.shape[0], result.shape[1]))
for k, v in full_to_colour.items():
pred_r = np.zeros((result.shape[0], result.shape[1]))
pred_r[(result == k)] = v[0]
pred_g = np.zeros((result.shape[0], result.shape[1]))
pred_g[(result == k)] = v[1]
pred_b = np.zeros((result.shape[0], result.shape[1]))
pred_b[(result == k)] = v[2]
uuu = np.stack((pred_r, pred_g, pred_b))
pred_colour += uuu
save_image(torch.from_numpy(pred_colour).float().div(255), os.path.join(output_folder + os.sep + id + "_prediction.jpg"))
def test(opt):
output_folder = os.path.join(opt.output, "VOC{}_{}".format(opt.year, opt.test_set))
if os.path.isdir(output_folder):
shutil.rmtree(output_folder)
os.makedirs(output_folder)
if torch.cuda.is_available():
if opt.pre_trained_model_type == "model":
model = torch.load(opt.pre_trained_model_path)
else:
model = Deeplab(num_classes=21)
model.load_state_dict(torch.load(opt.pre_trained_model_path))
else:
if opt.pre_trained_model_type == "model":
model = torch.load(opt.pre_trained_model_path, map_location=lambda storage, loc: storage)
else:
model = Deeplab(num_classes=21)
model.load_state_dict(torch.load(opt.pre_trained_model_path, map_location=lambda storage, loc: storage))
model.eval()
for image_path in glob.iglob("test_images/" + '*.jpg'):
if "prediction" in image_path:
continue
img = np.zeros((513, 513, 3))
image = cv2.imread(image_path).astype(np.float32)
h,w,_ = image.shape
image = cv2.resize(image, (int(w/4), int(h/4)))
h1, w1,_ = image.shape
image[:, :, 0] -= 104.008
image[:, :, 1] -= 116.669
image[:, :, 2] -= 122.675
img[:h1, :w1, :] = image
img = np.transpose(np.array(img, dtype=np.float32), (2, 0, 1))
img = img[None, :, :, :]
img = torch.Tensor(img)
if torch.cuda.is_available():
img = img.cuda()
with torch.no_grad():
result = model(img)
interp = nn.UpsamplingBilinear2d(size=(513, 513))
output = interp(result[3]).cpu().data[0].numpy()
output = output[:, :h1, :w1]
result = np.argmax(output.transpose(1,2,0), axis=2).astype(np.uint8)
pred_colour = np.zeros((3, result.shape[0], result.shape[1]))
for k, v in full_to_colour.items():
pred_r = np.zeros((result.shape[0], result.shape[1]))
pred_r[(result == k)] = v[0]
pred_g = np.zeros((result.shape[0], result.shape[1]))
pred_g[(result == k)] = v[1]
pred_b = np.zeros((result.shape[0], result.shape[1]))
pred_b[(result == k)] = v[2]
uuu = np.stack((pred_r, pred_g, pred_b))
pred_colour += uuu
save_image(torch.from_numpy(pred_colour).float().div(255), image_path[:-4] + "_prediction.jpg")
def __init__(self, num_classes=21,n=32,R=3):
super(RWN, self).__init__()
self.deeplab = Deeplab(num_classes=num_classes)
for param in self.deeplab.parameters():
param.requires_grad = False
self.n=n
self.R=R
self.k=k=3+64+64
self.upsample = nn.Upsample(size=(n,n), mode='bilinear', align_corners=True)
self.conv_parameter=np.random.rand(self.k)
def __init__(self, num_classes=21, n=32, R=3):
super(RWN, self).__init__()
self.deeplab = Deeplab(num_classes=num_classes)
for param in self.deeplab.parameters():
param.requires_grad = False
self.n = n
self.R = R
self.k = k = 3 + 64 + 64
self.upsample = nn.Upsample(size=(n, n),
mode='bilinear',
align_corners=True)
self.conv = nn.Conv2d(k, 1, (1, 1), bias=False)
self.pdist = nn.PairwiseDistance(p=1)
class RWN(nn.Module):
def __init__(self, num_classes=21,n=32,R=3):
super(RWN, self).__init__()
self.deeplab = Deeplab(num_classes=num_classes)
for param in self.deeplab.parameters():
param.requires_grad = False
self.n=n
self.R=R
self.k=k=3+64+64
self.upsample = nn.Upsample(size=(n,n), mode='bilinear', align_corners=True)
self.conv_parameter=np.random.rand(self.k)
def forward(self, x):
batch_size=x.size()[0]
input_size = x.size()[2:4]
n=self.n
n2=n**2
assert input_size[0]==input_size[1]==n
x2 = self.deeplab(x)[3]
conv1,conv2,_=self.deeplab.low_level_feature
x2=self.upsample(x2).view(batch_size,21,n2)
conv1=self.upsample(conv1)
conv2=self.upsample(conv2)
intergrated=torch.cat((x,conv1,conv2),1)
intergrated=intergrated.view(batch_size,self.k,self.n,self.n)
y=[]
for batch in range(batch_size):
W = None
to_select=itertools.product(range(self.n),range(self.n))
selected_indices = torch.LongTensor(random.sample(list(to_select), int(0.3 * n2)))
for ki in range(self.k):
intergrated_sampled = self._sampling(intergrated[batch][ki],selected_indices)
if W is None:
W=self._L1_distance_within_R(intergrated_sampled,ki)
else:
W+=self._L1_distance_within_R(intergrated_sampled,ki)
W=self._to_torch_sparse(W)
x_batch=x2[batch].transpose(0,1)
y.append(torch.matmul(W,x_batch).view(1,n2,21))
return torch.cat(y)
def _L1_distance_within_R(self,x,ki):
n2=self.n**2
indices_x=[]
data=[]
for i in range(x.shape[0]):
for j in range(x.shape[1]):
if x[i,j]==0:
break
new_x=utils.calc_D_between_scalar_matrix_within_R(x,i,j,self.R,self.n)
new_x.data=np.exp(new_x.data)
new_x.data=new_x.data/new_x.data.sum()
indices_x+=[(i * self.n + j)]*new_x.data.shape[0]
data.append(new_x)
indices_y,value=utils.get_flatten_indices(data,self.n)
value=np.array(value)*self.conv_parameter[ki].item()
return csr_matrix((value,(indices_x,indices_y)),shape=torch.Size([n2,n2]))
def _sampling(self,x,selected_indices):
""" converts dense tensor x to sparse format """
indices = selected_indices.t()
values = x[tuple(indices[i] for i in range(indices.shape[0]))]
return csc_matrix((values,(indices[0],indices[1])),shape=(self.n,self.n))
def _to_torch_sparse(self,w):
w=w.tocoo()
value=torch.FloatTensor(w.data)
indices=torch.LongTensor([w.col,w.row])
size=torch.Size(w.shape)
return torch.sparse.FloatTensor(indices,value,size)
def train(opt):
if torch.cuda.is_available():
torch.cuda.manual_seed(123)
else:
torch.manual_seed(123)
training_params = {
"batch_size": opt.batch_size,
"shuffle": True,
"drop_last": True,
"collate_fn": custom_collate_fn
}
test_params = {
"batch_size": opt.batch_size,
"shuffle": False,
"drop_last": False,
"collate_fn": custom_collate_fn
}
training_set = VOCDataset(opt.data_path, opt.dataset, opt.image_size)
training_generator = DataLoader(training_set, **training_params)
test_set = VOCDataset(opt.data_path,
opt.dataset,
opt.image_size,
is_training=False)
test_generator = DataLoader(test_set, **test_params)
model = Deeplab(num_classes=training_set.num_classes + 1)
#model.load_state_dict(torch.load(opt.pre_trained_model))
log_path = os.path.join(opt.log_path, "{}".format(opt.dataset))
if os.path.isdir(log_path):
shutil.rmtree(log_path)
#os.makedirs(log_path)
writer = SummaryWriter(log_path)
writer.add_graph(
model, torch.rand(opt.batch_size, 3, opt.image_size, opt.image_size))
if torch.cuda.is_available():
model.cuda()
best_loss = 1e10
best_epoch = 0
model.train()
num_iter_per_epoch = len(training_generator)
for epoch in range(opt.num_epoches):
for iter, batch in enumerate(training_generator):
current_step = epoch * num_iter_per_epoch + iter
current_lr = update_lr(opt.lr, current_step,
num_iter_per_epoch * opt.num_epoches)
optimizer = get_optimizer(model, current_lr, opt.momentum,
opt.decay)
if torch.cuda.is_available():
batch = [torch.Tensor(record).cuda() for record in batch]
else:
batch = [torch.Tensor(record) for record in batch]
image, gt1, gt2 = batch
gt1 = gt1.long()
gt2 = gt2.long()
optimizer.zero_grad()
results = model(image)
mul_losses = multiple_losses(results, [gt1, gt1, gt2, gt1])
mul_losses[4].backward()
optimizer.step()
print(
"Epoch: {}/{}, Iteration: {}/{}, Lr: {}, Loss: {:.2f} (1xloss: {:.2f} 0.75xloss: {:.2f} 0.5xloss: {:.2f} Max_merged_loss: {:.2f})"
.format(epoch + 1, opt.num_epoches, iter + 1,
num_iter_per_epoch, optimizer.param_groups[0]['lr'],
mul_losses[4], mul_losses[0], mul_losses[1],
mul_losses[2], mul_losses[3]))
writer.add_scalar('Train/Total_loss', mul_losses[4], current_step)
writer.add_scalar('Train/1x_scale_loss', mul_losses[0],
current_step)
writer.add_scalar('Train/0.75x_scale_loss', mul_losses[1],
current_step)
writer.add_scalar('Train/0.5x_scale_loss', mul_losses[2],
current_step)
writer.add_scalar('Train/Max_merged_loss', mul_losses[3],
current_step)
if epoch % opt.test_interval == 0:
model.eval()
loss_ls = []
loss_scale_1_ls = []
loss_scale_2_ls = []
loss_scale_3_ls = []
loss_max_merged_ls = []
for te_batch in test_generator:
if torch.cuda.is_available():
te_batch = [
torch.Tensor(record).cuda() for record in te_batch
]
else:
te_batch = [torch.Tensor(record) for record in te_batch]
te_image, te_gt1, te_gt2 = te_batch
te_gt1 = te_gt1.long()
te_gt2 = te_gt2.long()
num_sample = len(te_gt1)
with torch.no_grad():
te_results = model(te_image)
te_mul_losses = multiple_losses(
te_results, [te_gt1, te_gt1, te_gt2, te_gt1])
loss_ls.append(te_mul_losses[4] * num_sample)
loss_scale_1_ls.append(te_mul_losses[0] * num_sample)
loss_scale_2_ls.append(te_mul_losses[1] * num_sample)
loss_scale_3_ls.append(te_mul_losses[2] * num_sample)
loss_max_merged_ls.append(te_mul_losses[3] * num_sample)
te_loss = sum(loss_ls) / test_set.__len__()
te_scale_1_loss = sum(loss_scale_1_ls) / test_set.__len__()
te_scale_2_loss = sum(loss_scale_2_ls) / test_set.__len__()
te_scale_3_loss = sum(loss_scale_3_ls) / test_set.__len__()
te_max_merged_loss = sum(loss_max_merged_ls) / test_set.__len__()
print(
"Epoch: {}/{}, Lr: {}, Loss: {:.2f} (1xloss: {:.2f} 0.75xloss: {:.2f} 0.5xloss: {:.2f} Max_merged_loss: {:.2f})"
.format(epoch + 1, opt.num_epoches,
optimizer.param_groups[0]['lr'], te_loss,
te_scale_1_loss, te_scale_2_loss, te_scale_3_loss,
te_max_merged_loss))
writer.add_scalar('Test/Total_loss', te_loss, epoch)
writer.add_scalar('Test/1x_scale_loss', te_scale_1_loss, epoch)
writer.add_scalar('Test/0.75x_scale_loss', te_scale_2_loss, epoch)
writer.add_scalar('Test/0.5x_scale_loss', te_scale_3_loss, epoch)
writer.add_scalar('Test/Max_merged_loss', te_max_merged_loss,
epoch)
model.train()
if te_loss + opt.es_min_delta < best_loss:
best_loss = te_loss
best_epoch = epoch
torch.save(
model.state_dict(), opt.saved_path + os.sep +
"only_params_trained_deeplab_voc")
torch.save(
model, opt.saved_path + os.sep +
"whole_model_trained_deeplab_voc")
# Early stopping
if epoch - best_epoch > opt.es_patience > 0:
print(
"Stop training at epoch {}. The lowest loss achieved is {}"
.format(epoch, te_loss))
break
writer.close()
class RWN(nn.Module):
def __init__(self, num_classes=21, n=32, R=3):
super(RWN, self).__init__()
self.deeplab = Deeplab(num_classes=num_classes)
for param in self.deeplab.parameters():
param.requires_grad = False
self.n = n
self.R = R
self.k = k = 3 + 64 + 64
self.upsample = nn.Upsample(size=(n, n),
mode='bilinear',
align_corners=True)
self.conv = nn.Conv2d(k, 1, (1, 1), bias=False)
self.pdist = nn.PairwiseDistance(p=1)
def forward(self, x):
batch_size = x.size()[0]
input_size = x.size()[2:4]
n = self.n
n2 = n**2
assert input_size[0] == input_size[1] == n
x2 = self.deeplab(x)[3]
conv1, conv2, _ = self.deeplab.low_level_feature
x2 = self.upsample(x2).view(batch_size, 21, n2)
conv1 = self.upsample(conv1)
conv2 = self.upsample(conv2)
intergrated = torch.cat((x, conv1, conv2), 1)
intergrated = self.conv(intergrated).view(batch_size, self.n, self.n)
y = []
for batch in range(batch_size):
intergrated_sampled = self._sampling(intergrated[batch].detach())
W = self._L1_distance_within_R(intergrated_sampled)
x_batch = x2[batch].transpose(0, 1)
y.append(torch.matmul(W, x_batch).view(1, n2, 21))
return torch.cat(y)
def _L1_distance_within_R(self, x):
n2 = self.n**2
indices_x = []
data = []
for i in range(x.shape[0]):
for j in range(x.shape[1]):
if x[i, j] == 0:
break
new_x = utils.calc_D_between_scalar_matrix_within_R(
x, i, j, self.R, self.n)
new_x.data = np.exp(new_x.data)
new_x.data = new_x.data / new_x.data.sum()
indices_x += [(i * self.n + j)] * new_x.data.shape[0]
data.append(new_x)
indices_y, value = utils.get_flatten_indices(data, self.n)
return torch.sparse.FloatTensor(
torch.LongTensor([indices_x, indices_y]), torch.FloatTensor(value),
torch.Size([n2, n2]))
def _sampling(self, x, fractions=0.3):
""" converts dense tensor x to sparse format """
indices = torch.nonzero(x)
indices.requires_grad = False
if len(indices.shape) == 0: # if all elements are zeros
return None
indices_len = indices.shape[0]
selected_indices = random.sample(range(indices_len),
int(fractions * indices_len))
indices = indices[selected_indices]
indices = indices.t()
values = x[tuple(indices[i] for i in range(indices.shape[0]))]
return csc_matrix((values, (indices[0], indices[1])),
shape=(self.n, self.n))