def for_loop(net, data_loader, train_optimizer): is_train = train_optimizer is not None net.train() if is_train else net.eval() total_loss, total_time, total_num, preds, targets = 0.0, 0.0, 0, [], [] data_bar = tqdm(data_loader, dynamic_ncols=True) with (torch.enable_grad() if is_train else torch.no_grad()): for data, target, grad, boundary, name in data_bar: data, target, grad, boundary = data.cuda(), target.cuda(), grad.cuda(), boundary.cuda() torch.cuda.synchronize() start_time = time.time() seg, edge = net(data, grad) prediction = torch.argmax(seg.detach(), dim=1) torch.cuda.synchronize() end_time = time.time() semantic_loss = semantic_criterion(seg, target) edge_loss = edge_criterion(edge, target, boundary) task_loss = task_criterion(seg, edge, target) loss = semantic_loss + 20 * edge_loss + task_loss if is_train: train_optimizer.zero_grad() loss.backward() train_optimizer.step() total_num += data.size(0) total_time += end_time - start_time total_loss += loss.item() * data.size(0) preds.append(prediction.cpu()) targets.append(target.cpu()) if not is_train: if data_loader.dataset.split == 'test': # revert train id to regular id for key in sorted(trainId2label.keys(), reverse=True): prediction[prediction == key] = trainId2label[key].id # save pred images save_root = '{}/{}_{}_{}/{}'.format(save_path, backbone_type, crop_h, crop_w, data_loader.dataset.split) if not os.path.exists(save_root): os.makedirs(save_root) for pred_tensor, pred_name in zip(prediction, name): pred_img = ToPILImage()(pred_tensor.unsqueeze(dim=0).byte().cpu()) if data_loader.dataset.split == 'val': pred_img.putpalette(get_palette()) pred_name = pred_name.replace('leftImg8bit', 'color') path = '{}/{}'.format(save_root, pred_name) pred_img.save(path) data_bar.set_description('{} Epoch: [{}/{}] Loss: {:.4f} FPS: {:.0f}' .format(data_loader.dataset.split.capitalize(), epoch, epochs, total_loss / total_num, total_num / total_time)) # compute metrics preds = torch.cat(preds, dim=0) targets = torch.cat(targets, dim=0) pa, mpa, class_iou, category_iou = compute_metrics(preds, targets) print('{} Epoch: [{}/{}] PA: {:.2f}% mPA: {:.2f}% Class_mIOU: {:.2f}% Category_mIOU: {:.2f}%' .format(data_loader.dataset.split.capitalize(), epoch, epochs, pa * 100, mpa * 100, class_iou * 100, category_iou * 100)) return total_loss / total_num, pa * 100, mpa * 100, class_iou * 100, category_iou * 100
def train_val(net, data_loader, train_optimizer): is_train = train_optimizer is not None net.train() if is_train else net.eval() total_loss, total_correct, total_time, total_num, data_bar = 0.0, 0.0, 0.0, 0, tqdm( data_loader) with (torch.enable_grad() if is_train else torch.no_grad()): for data, target, name in data_bar: data, target = data.cuda(), target.cuda() torch.cuda.synchronize() start_time = time.time() out = net(data) prediction = torch.argmax(out, dim=1) torch.cuda.synchronize() end_time = time.time() loss = loss_criterion(out, target) if is_train: train_optimizer.zero_grad() loss.backward() train_optimizer.step() total_num += data.size(0) total_time += end_time - start_time total_loss += loss.item() * data.size(0) total_correct += torch.sum( prediction == target).item() / target.numel() * data.size(0) if not is_train and epoch % save_step == 0: # revert train id to regular id for key in trainId2label.keys(): prediction[prediction == key] = trainId2label[key].id # save pred images for pred_tensor, pred_name in zip(prediction, name): pred_img = ToPILImage()( pred_tensor.unsqueeze(dim=0).byte().cpu()) pred_img.putpalette(palette) pred_img.save('results/{}'.format( pred_name.replace('leftImg8bit', 'color'))) data_bar.set_description( '{} Epoch: [{}/{}] Loss: {:.4f} mPA: {:.2f}% FPS: {:.0f}'. format('Train' if is_train else 'Val', epoch, epochs, total_loss / total_num, total_correct / total_num * 100, total_num / total_time)) return total_loss / total_num, total_correct / total_num * 100
def savelabel(label, name): l = ToPILImage()(label[0].unsqueeze(0).byte()).convert('P') l.putpalette(lab_palette) l.save(name)
target = Image.new('RGB', (image_width * num_width, image_height)) images = [image] image = transform(image).unsqueeze(dim=0).cuda() # model load model = FastSCNN(in_channels=3, num_classes=19) model.load_state_dict( torch.load(model_weight, map_location=torch.device('cpu'))) model = model.cuda() model.eval() # predict and save image with torch.no_grad(): output = model(image) pred = torch.argmax(output, dim=1) pred_image = ToPILImage()(pred.byte().cpu()) pred_image.putpalette(palette) if 'test' not in input_pic: gt_image = Image.open('{}/gtFine/{}'.format( data_path, input_pic.replace('leftImg8bit', 'gtFine_color'))) images.append(gt_image) images.append(pred_image) # concat images for i in range(len(images)): left, top, right, bottom = image_width * i, 0, image_width * ( i + 1), image_height target.paste(images[i], (left, top, right, bottom)) target.save( os.path.split(input_pic)[-1].replace('leftImg8bit', 'result'))
grad = torch.from_numpy( np.expand_dims(np.asarray(grad, np.float32) / 255.0, axis=0).copy()).unsqueeze(dim=0).cuda() # model load model = GatedSCNN(model_weight.split('/')[-1].split('_')[0], num_classes=19) model.load_state_dict( torch.load(model_weight, map_location=torch.device('cpu'))) model = model.cuda() model.eval() # predict and save image with torch.no_grad(): output, _ = model(image, grad) pred = torch.argmax(output, dim=1) pred_image = ToPILImage()(pred.byte().cpu()) pred_image.putpalette(get_palette()) if 'test' not in input_pic: gt_image = Image.open('{}/gtFine/{}'.format( data_path, input_pic.replace('leftImg8bit', 'gtFine_color'))) images.append(gt_image) images.append(pred_image) # concat images for i in range(len(images)): left, top, right, bottom = image_width * i, 0, image_width * ( i + 1), image_height target.paste(images[i], (left, top, right, bottom)) target.save( os.path.split(input_pic)[-1].replace('leftImg8bit', 'result'))