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))