def get_dataset_loaders(model, dataset, workers): target_size = (model["common"]["image_size"],) * 2 batch_size = model["common"]["batch_size"] path = dataset["common"]["dataset"] mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225] transform = JointCompose( [ JointTransform(ConvertImageMode("RGB"), ConvertImageMode("P")), JointTransform(Resize(target_size, Image.BILINEAR), Resize(target_size, Image.NEAREST)), JointTransform(CenterCrop(target_size), CenterCrop(target_size)), JointRandomHorizontalFlip(0.5), JointRandomRotation(0.5, 90), JointRandomRotation(0.5, 90), JointRandomRotation(0.5, 90), JointTransform(ImageToTensor(), MaskToTensor()), JointTransform(Normalize(mean=mean, std=std), None), ] ) train_dataset = SlippyMapTilesConcatenation( [os.path.join(path, "training", "images")], os.path.join(path, "training", "labels"), transform ) val_dataset = SlippyMapTilesConcatenation( [os.path.join(path, "validation", "images")], os.path.join(path, "validation", "labels"), transform ) train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True, num_workers=workers) val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, drop_last=True, num_workers=workers) return train_loader, val_loader
def segment(self, image): # don't track tensors with autograd during prediction with torch.no_grad(): mean, std = self.dataset['stats']['mean'], self.dataset['stats'][ 'std'] transform = Compose([ ConvertImageMode(mode='RGB'), ImageToTensor(), Normalize(mean=mean, std=std) ]) image = transform(image) batch = image.unsqueeze(0).to(self.device) output = self.net(batch) output = output.cpu().data.numpy() output = output.squeeze(0) mask = output.argmax(axis=0).astype(np.uint8) mask = Image.fromarray(mask, mode='P') palette = make_palette(*self.dataset['common']['colors']) mask.putpalette(palette) return mask
def main(args): dataset = load_config(args.dataset) path = dataset["common"]["dataset"] num_classes = len(dataset["common"]["classes"]) train_transform = Compose([ConvertImageMode(mode="P"), MaskToTensor()]) train_dataset = SlippyMapTiles(os.path.join(path, "training", "labels"), transform=train_transform) n = 0 counts = np.zeros(num_classes, dtype=np.int64) loader = DataLoader(train_dataset, batch_size=1) for images, tile in tqdm(loader, desc="Loading", unit="image", ascii=True): image = torch.squeeze(images) image = np.array(image, dtype=np.uint8) n += image.shape[0] * image.shape[1] counts += np.bincount(image.ravel(), minlength=num_classes) # Class weighting scheme `w = 1 / ln(c + p)` see: # - https://arxiv.org/abs/1707.03718 # LinkNet: Exploiting Encoder Representations for Efficient Semantic Segmentation # - https://arxiv.org/abs/1606.02147 # ENet: A Deep Neural Network Architecture for Real-Time Semantic Segmentation probs = counts / n weights = 1 / np.log(1.02 + probs) weights.round(6, out=weights) print(weights.tolist())
def get_dataset_loaders(model, dataset): target_size = (model["common"]["image_size"], ) * 2 batch_size = model["common"]["batch_size"] path = dataset["common"]["dataset"] mean, std = dataset["stats"]["mean"], dataset["stats"]["std"] image_transform = Compose([ ConvertImageMode("RGB"), Resize(target_size, Image.BILINEAR), CenterCrop(target_size), ImageToTensor(), Normalize(mean=mean, std=std), ]) target_transform = Compose([ ConvertImageMode("P"), Resize(target_size, Image.NEAREST), CenterCrop(target_size), MaskToTensor() ]) train_dataset = SlippyMapTilesConcatenation( [os.path.join(path, "training", "images")], [image_transform], os.path.join(path, "training", "labels"), target_transform, ) val_dataset = SlippyMapTilesConcatenation( [os.path.join(path, "validation", "images")], [image_transform], os.path.join(path, "validation", "labels"), target_transform, ) train_loader = DataLoader(train_dataset, batch_size=batch_size, drop_last=True) val_loader = DataLoader(val_dataset, batch_size=batch_size, drop_last=True) return train_loader, val_loader
def get_dataset_loaders(model, dataset): target_size = (model['common']['image_size'], ) * 2 batch_size = model['common']['batch_size'] path = dataset['common']['dataset'] mean, std = dataset['stats']['mean'], dataset['stats']['std'] image_transform = Compose([ ConvertImageMode('RGB'), Resize(target_size, Image.BILINEAR), CenterCrop(target_size), ImageToTensor(), Normalize(mean=mean, std=std)]) target_transform = Compose([ ConvertImageMode('P'), Resize(target_size, Image.NEAREST), CenterCrop(target_size), MaskToTensor()]) train_dataset = SlippyMapTilesConcatenation( [os.path.join(path, 'training', 'images')], [image_transform], os.path.join(path, 'training', 'labels'), target_transform) val_dataset = SlippyMapTilesConcatenation( [os.path.join(path, 'validation', 'images')], [image_transform], os.path.join(path, 'validation', 'labels'), target_transform) train_sampler = RandomSubsetSampler(train_dataset, dataset['samples']['training']) val_sampler = RandomSubsetSampler(val_dataset, dataset['samples']['validation']) train_loader = DataLoader(train_dataset, batch_size=batch_size, sampler=train_sampler, drop_last=True) val_loader = DataLoader(val_dataset, batch_size=batch_size, sampler=val_sampler, drop_last=True) return train_loader, val_loader
def main(args): dataset = load_config(args.dataset) path = dataset["common"]["dataset"] train_transform = Compose([ConvertImageMode(mode="RGB"), ImageToTensor()]) train_dataset = SlippyMapTiles(os.path.join(path, "training", "images"), transform=train_transform) n = 0 mean = np.zeros(3, dtype=np.float64) loader = DataLoader(train_dataset, batch_size=1) for images, tile in tqdm(loader, desc="Loading", unit="image", ascii=True): image = torch.squeeze(images) assert image.size(0) == 3, "channel first" image = np.array(image, dtype=np.float64) n += image.shape[1] * image.shape[2] mean += np.sum(image, axis=(1, 2)) mean /= n mean.round(decimals=6, out=mean) print("mean: {}".format(mean.tolist())) std = np.zeros(3, dtype=np.float64) loader = DataLoader(train_dataset, batch_size=1) for images, tile in tqdm(loader, desc="Loading", unit="image", ascii=True): image = torch.squeeze(images) assert image.size(0) == 3, "channel first" image = np.array(image, dtype=np.float64) difference = np.transpose(image, (1, 2, 0)) - mean std += np.sum(np.square(difference), axis=(0, 1)) std = np.sqrt(std / (n - 1)) std.round(decimals=6, out=std) print("std: {}".format(std.tolist()))
def main(args): model = load_config(args.model) dataset = load_config(args.dataset) cuda = model["common"]["cuda"] device = torch.device("cuda" if cuda else "cpu") def map_location(storage, _): return storage.cuda() if cuda else storage.cpu() if cuda and not torch.cuda.is_available(): sys.exit("Error: CUDA requested but not available") num_classes = len(dataset["common"]["classes"]) # https://github.com/pytorch/pytorch/issues/7178 chkpt = torch.load(args.checkpoint, map_location=map_location) net = UNet(num_classes).to(device) net = nn.DataParallel(net) if cuda: torch.backends.cudnn.benchmark = True net.load_state_dict(chkpt["state_dict"]) net.eval() mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225] transform = Compose([ ConvertImageMode(mode="RGB"), ImageToTensor(), Normalize(mean=mean, std=std) ]) directory = BufferedSlippyMapDirectory(args.tiles, transform=transform, size=args.tile_size, overlap=args.overlap) loader = DataLoader(directory, batch_size=args.batch_size, num_workers=args.workers) # don't track tensors with autograd during prediction with torch.no_grad(): for images, tiles in tqdm(loader, desc="Eval", unit="batch", ascii=True): images = images.to(device) outputs = net(images) # manually compute segmentation mask class probabilities per pixel probs = nn.functional.softmax(outputs, dim=1).data.cpu().numpy() for tile, prob in zip(tiles, probs): x, y, z = list(map(int, tile)) # we predicted on buffered tiles; now get back probs for original image prob = directory.unbuffer(prob) # Quantize the floating point probabilities in [0,1] to [0,255] and store # a single-channel `.png` file with a continuous color palette attached. assert prob.shape[ 0] == 2, "single channel requires binary model" assert np.allclose( np.sum(prob, axis=0), 1. ), "single channel requires probabilities to sum up to one" foreground = prob[1:, :, :] anchors = np.linspace(0, 1, 256) quantized = np.digitize(foreground, anchors).astype(np.uint8) palette = continuous_palette_for_color("pink", 256) out = Image.fromarray(quantized.squeeze(), mode="P") out.putpalette(palette) os.makedirs(os.path.join(args.probs, str(z), str(x)), exist_ok=True) path = os.path.join(args.probs, str(z), str(x), str(y) + ".png") out.save(path, optimize=True)