Spaces:
Runtime error
Runtime error
| import logging | |
| import math | |
| from collections import OrderedDict | |
| import mmcv | |
| import numpy as np | |
| import torch | |
| from torchvision.utils import save_image | |
| from models.archs.fcn_arch import FCNHead | |
| from models.archs.shape_attr_embedding_arch import ShapeAttrEmbedding | |
| from models.archs.unet_arch import ShapeUNet | |
| from models.losses.accuracy import accuracy | |
| from models.losses.cross_entropy_loss import CrossEntropyLoss | |
| logger = logging.getLogger('base') | |
| class ParsingGenModel(): | |
| """Paring Generation model. | |
| """ | |
| def __init__(self, opt): | |
| self.opt = opt | |
| self.device = torch.device('cuda') | |
| self.is_train = opt['is_train'] | |
| self.attr_embedder = ShapeAttrEmbedding( | |
| dim=opt['embedder_dim'], | |
| out_dim=opt['embedder_out_dim'], | |
| cls_num_list=opt['attr_class_num']).to(self.device) | |
| self.parsing_encoder = ShapeUNet( | |
| in_channels=opt['encoder_in_channels']).to(self.device) | |
| self.parsing_decoder = FCNHead( | |
| in_channels=opt['fc_in_channels'], | |
| in_index=opt['fc_in_index'], | |
| channels=opt['fc_channels'], | |
| num_convs=opt['fc_num_convs'], | |
| concat_input=opt['fc_concat_input'], | |
| dropout_ratio=opt['fc_dropout_ratio'], | |
| num_classes=opt['fc_num_classes'], | |
| align_corners=opt['fc_align_corners'], | |
| ).to(self.device) | |
| self.init_training_settings() | |
| self.palette = [[0, 0, 0], [255, 250, 250], [220, 220, 220], | |
| [250, 235, 215], [255, 250, 205], [211, 211, 211], | |
| [70, 130, 180], [127, 255, 212], [0, 100, 0], | |
| [50, 205, 50], [255, 255, 0], [245, 222, 179], | |
| [255, 140, 0], [255, 0, 0], [16, 78, 139], | |
| [144, 238, 144], [50, 205, 174], [50, 155, 250], | |
| [160, 140, 88], [213, 140, 88], [90, 140, 90], | |
| [185, 210, 205], [130, 165, 180], [225, 141, 151]] | |
| def init_training_settings(self): | |
| optim_params = [] | |
| for v in self.attr_embedder.parameters(): | |
| if v.requires_grad: | |
| optim_params.append(v) | |
| for v in self.parsing_encoder.parameters(): | |
| if v.requires_grad: | |
| optim_params.append(v) | |
| for v in self.parsing_decoder.parameters(): | |
| if v.requires_grad: | |
| optim_params.append(v) | |
| # set up optimizers | |
| self.optimizer = torch.optim.Adam( | |
| optim_params, | |
| self.opt['lr'], | |
| weight_decay=self.opt['weight_decay']) | |
| self.log_dict = OrderedDict() | |
| self.entropy_loss = CrossEntropyLoss().to(self.device) | |
| def feed_data(self, data): | |
| self.pose = data['densepose'].to(self.device) | |
| self.attr = data['attr'].to(self.device) | |
| self.segm = data['segm'].to(self.device) | |
| def optimize_parameters(self): | |
| self.attr_embedder.train() | |
| self.parsing_encoder.train() | |
| self.parsing_decoder.train() | |
| self.attr_embedding = self.attr_embedder(self.attr) | |
| self.pose_enc = self.parsing_encoder(self.pose, self.attr_embedding) | |
| self.seg_logits = self.parsing_decoder(self.pose_enc) | |
| loss = self.entropy_loss(self.seg_logits, self.segm) | |
| self.optimizer.zero_grad() | |
| loss.backward() | |
| self.optimizer.step() | |
| self.log_dict['loss_total'] = loss | |
| def get_vis(self, save_path): | |
| img_cat = torch.cat([ | |
| self.pose, | |
| self.segm, | |
| ], dim=3).detach() | |
| img_cat = ((img_cat + 1) / 2) | |
| img_cat = img_cat.clamp_(0, 1) | |
| save_image(img_cat, save_path, nrow=1, padding=4) | |
| def inference(self, data_loader, save_dir): | |
| self.attr_embedder.eval() | |
| self.parsing_encoder.eval() | |
| self.parsing_decoder.eval() | |
| acc = 0 | |
| num = 0 | |
| for _, data in enumerate(data_loader): | |
| pose = data['densepose'].to(self.device) | |
| attr = data['attr'].to(self.device) | |
| segm = data['segm'].to(self.device) | |
| img_name = data['img_name'] | |
| num += pose.size(0) | |
| with torch.no_grad(): | |
| attr_embedding = self.attr_embedder(attr) | |
| pose_enc = self.parsing_encoder(pose, attr_embedding) | |
| seg_logits = self.parsing_decoder(pose_enc) | |
| seg_pred = seg_logits.argmax(dim=1) | |
| acc += accuracy(seg_logits, segm) | |
| palette_label = self.palette_result(segm.cpu().numpy()) | |
| palette_pred = self.palette_result(seg_pred.cpu().numpy()) | |
| pose_numpy = ((pose[0] + 1) / 2. * 255.).expand( | |
| 3, | |
| pose[0].size(1), | |
| pose[0].size(2), | |
| ).cpu().numpy().clip(0, 255).astype(np.uint8).transpose(1, 2, 0) | |
| concat_result = np.concatenate( | |
| (pose_numpy, palette_pred, palette_label), axis=1) | |
| mmcv.imwrite(concat_result, f'{save_dir}/{img_name[0]}') | |
| self.attr_embedder.train() | |
| self.parsing_encoder.train() | |
| self.parsing_decoder.train() | |
| return (acc / num).item() | |
| def get_current_log(self): | |
| return self.log_dict | |
| def update_learning_rate(self, epoch): | |
| """Update learning rate. | |
| Args: | |
| current_iter (int): Current iteration. | |
| warmup_iter (int): Warmup iter numbers. -1 for no warmup. | |
| Default: -1. | |
| """ | |
| lr = self.optimizer.param_groups[0]['lr'] | |
| if self.opt['lr_decay'] == 'step': | |
| lr = self.opt['lr'] * ( | |
| self.opt['gamma']**(epoch // self.opt['step'])) | |
| elif self.opt['lr_decay'] == 'cos': | |
| lr = self.opt['lr'] * ( | |
| 1 + math.cos(math.pi * epoch / self.opt['num_epochs'])) / 2 | |
| elif self.opt['lr_decay'] == 'linear': | |
| lr = self.opt['lr'] * (1 - epoch / self.opt['num_epochs']) | |
| elif self.opt['lr_decay'] == 'linear2exp': | |
| if epoch < self.opt['turning_point'] + 1: | |
| # learning rate decay as 95% | |
| # at the turning point (1 / 95% = 1.0526) | |
| lr = self.opt['lr'] * ( | |
| 1 - epoch / int(self.opt['turning_point'] * 1.0526)) | |
| else: | |
| lr *= self.opt['gamma'] | |
| elif self.opt['lr_decay'] == 'schedule': | |
| if epoch in self.opt['schedule']: | |
| lr *= self.opt['gamma'] | |
| else: | |
| raise ValueError('Unknown lr mode {}'.format(self.opt['lr_decay'])) | |
| # set learning rate | |
| for param_group in self.optimizer.param_groups: | |
| param_group['lr'] = lr | |
| return lr | |
| def save_network(self, save_path): | |
| """Save networks. | |
| """ | |
| save_dict = {} | |
| save_dict['embedder'] = self.attr_embedder.state_dict() | |
| save_dict['encoder'] = self.parsing_encoder.state_dict() | |
| save_dict['decoder'] = self.parsing_decoder.state_dict() | |
| torch.save(save_dict, save_path) | |
| def load_network(self): | |
| checkpoint = torch.load(self.opt['pretrained_parsing_gen']) | |
| self.attr_embedder.load_state_dict(checkpoint['embedder'], strict=True) | |
| self.attr_embedder.eval() | |
| self.parsing_encoder.load_state_dict( | |
| checkpoint['encoder'], strict=True) | |
| self.parsing_encoder.eval() | |
| self.parsing_decoder.load_state_dict( | |
| checkpoint['decoder'], strict=True) | |
| self.parsing_decoder.eval() | |
| def palette_result(self, result): | |
| seg = result[0] | |
| palette = np.array(self.palette) | |
| assert palette.shape[1] == 3 | |
| assert len(palette.shape) == 2 | |
| color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8) | |
| for label, color in enumerate(palette): | |
| color_seg[seg == label, :] = color | |
| # convert to BGR | |
| color_seg = color_seg[..., ::-1] | |
| return color_seg | |