def predict(config, cover, args, palette, chkpt, nn, device, mode): assert mode in ["predict", "predict_translate"], "Predict unknown mode" loader_module = load_module("robosat_pink.loaders.{}".format( chkpt["loader"].lower())) loader_predict = getattr(loader_module, chkpt["loader"])(config, chkpt["shape_in"][1:3], args.dataset, cover, mode=mode) loader = DataLoader(loader_predict, batch_size=args.bs, num_workers=args.workers) assert len(loader), "Empty predict dataset directory. Check your path." tiled = [] for images, tiles in tqdm(loader, desc="Eval", unit="batch", ascii=True): images = images.to(device) for tile, prob in zip( tiles, torch.nn.functional.softmax(nn(images), dim=1).data.cpu().numpy()): x, y, z = list(map(int, tile)) mask = np.around(prob[1:, :, :]).astype(np.uint8).squeeze() if mode == "predict": tile_label_to_file(args.out, mercantile.Tile(x, y, z), palette, mask) if mode == "predict_translate": tile_translate_to_file(args.out, mercantile.Tile(x, y, z), palette, mask, config["model"]["ms"]) tiled.append(mercantile.Tile(x, y, z)) return tiled
def main(args): chkpt = torch.load(args.checkpoint, map_location=torch.device("cpu")) model_module = load_module("robosat_pink.models.{}".format(chkpt["nn"].lower())) nn = getattr(model_module, chkpt["nn"])(chkpt["shape_in"], chkpt["shape_out"]).to("cpu") print("RoboSat.pink - export model to {}".format(args.type)) print("Model: {} - UUID: {} - Torch {}".format(chkpt["nn"], chkpt["uuid"], torch.__version__)) print(chkpt["doc_string"]) try: # https://github.com/pytorch/pytorch/issues/9176 nn.module.state_dict(chkpt["state_dict"]) except AttributeError: nn.state_dict(chkpt["state_dict"]) nn.eval() batch = torch.rand(1, *chkpt["shape_in"]) if args.type == "onnx": torch.onnx.export( nn, torch.autograd.Variable(batch), args.out, input_names=["input", "shape_in", "shape_out"], output_names=["output"], dynamic_axes={"input": {0: "num_batch"}, "output": {0: "num_batch"}}, ) if args.type == "jit": torch.jit.trace(nn, batch).save(args.out)
def __init__(self, metrics, config=None): self.config = config self.metrics = {metric: 0.0 for metric in metrics} self.modules = { metric: load_module("robosat_pink.metrics." + metric) for metric in metrics } self.n = 0
def __init__(self, shape_in, shape_out, encoder="resnet50", train_config=None): self.doc = """ U-Net inspired encoder-decoder architecture with a ResNet encoder as proposed by Alexander Buslaev. - https://arxiv.org/abs/1505.04597 - U-Net: Convolutional Networks for Biomedical Image Segmentation - https://arxiv.org/pdf/1804.08024 - Angiodysplasia Detection and Localization Using DCNN - https://arxiv.org/abs/1806.00844 - TernausNetV2: Fully Convolutional Network for Instance Segmentation """ if encoder in ["resnet50", "resnet101", "resnet152"]: self.doc += "https://arxiv.org/abs/1512.03385 - Deep Residual Learning for Image Recognition" elif encoder in ["resnext50_32x4d", "resnext101_32x8d"]: self.doc += "https://arxiv.org/pdf/1611.05431 - Aggregated Residual Transformations for DNN" elif encoder in ["wide_resnet50_2", "wide_resnet101_2"]: self.doc += "https://arxiv.org/abs/1605.07146 - Wide Residual Networks" else: encoders = "Resnet50, Resnet101, Resnet152, Resnext50_32x4d, Resnext101_32x8d, Wide_resnet50_2, Wide_resnet101_2" assert False, "Albunet, expects as encoder: " + encoders self.version = 1 num_filters = 32 num_channels = shape_in[0] num_classes = shape_out[0] super().__init__() try: pretrained = train_config["model"]["pretrained"] except: pretrained = False models = load_module("torchvision.models") self.resnet = getattr(models, encoder)(pretrained=pretrained) assert num_channels if num_channels != 3: weights = nn.init.xavier_uniform_(torch.zeros((64, num_channels, 7, 7))) if pretrained: for c in range(min(num_channels, 3)): weights.data[:, c, :, :] = self.resnet.conv1.weight.data[:, c, :, :] self.resnet.conv1 = nn.Conv2d(num_channels, 64, kernel_size=7, stride=2, padding=3, bias=False) self.resnet.conv1.weight = nn.Parameter(weights) # No encoder reference, cf: https://github.com/pytorch/pytorch/issues/8392 self.center = DecoderBlock(2048, num_filters * 8) self.dec0 = DecoderBlock(2048 + num_filters * 8, num_filters * 8) self.dec1 = DecoderBlock(1024 + num_filters * 8, num_filters * 8) self.dec2 = DecoderBlock(512 + num_filters * 8, num_filters * 2) self.dec3 = DecoderBlock(256 + num_filters * 2, num_filters * 2 * 2) self.dec4 = DecoderBlock(num_filters * 2 * 2, num_filters) self.dec5 = ConvRelu(num_filters, num_filters) self.final = nn.Conv2d(num_filters, num_classes, kernel_size=1)
def main(args): print( "RoboSat.pink - extract {} from {}. Could take some time. Please wait." .format(args.type, args.pbf)) module = load_module("robosat_pink.osm.{}".format(args.type.lower())) osmium_handler = getattr(module, "{}Handler".format(args.type))() osmium_handler.apply_file(filename=os.path.expanduser(args.pbf), locations=True) osmium_handler.save(os.path.expanduser(args.out))
def main(args): config = load_config(args.config) args.out = os.path.expanduser(args.out) config["model"][ "loader"] = args.loader if args.loader else config["model"]["loader"] config["model"]["bs"] = args.bs if args.bs else config["model"]["bs"] config["model"]["lr"] = args.lr if args.lr else config["model"]["lr"] config["model"]["ts"] = tuple(map( int, args.ts.split(","))) if args.ts else config["model"]["ts"] config["model"]["nn"] = args.nn if args.nn else config["model"]["nn"] config["model"]["encoder"] = args.encoder if args.encoder else config[ "model"]["encoder"] config["model"][ "loss"] = args.loss if args.loss else config["model"]["loss"] config["model"]["da"] = args.da if args.da else config["model"]["da"] config["model"]["dap"] = args.dap if args.dap else config["model"]["dap"] args.workers = config["model"]["bs"] if not args.workers else args.workers check_classes(config) check_channels(config) check_model(config) assert os.path.isdir(os.path.expanduser( args.dataset)), "Dataset is not a directory" if args.no_training and args.no_validation: sys.exit() log = Logs(os.path.join(args.out, "log")) csv_train = None if args.no_training else open( os.path.join(args.out, "training.csv"), mode="a") csv_val = None if args.no_validation else open( os.path.join(args.out, "validation.csv"), mode="a") if torch.cuda.is_available(): log.log("RoboSat.pink - training on {} GPUs, with {} workers".format( torch.cuda.device_count(), args.workers)) log.log("(Torch:{} Cuda:{} CudNN:{})".format( torch.__version__, torch.version.cuda, torch.backends.cudnn.version())) device = torch.device("cuda") torch.backends.cudnn.benchmark = True else: log.log("RoboSat.pink - training on CPU, with {} workers - (Torch:{})". format(args.workers, torch.__version__)) log.log("") log.log("==========================================================") log.log("WARNING: Are you -really- sure about not training on GPU ?") log.log("==========================================================") log.log("") device = torch.device("cpu") log.log("--- Input tensor from Dataset: {} ---".format(args.dataset)) num_channel = 1 # 1-based numerotation for channel in config["channels"]: for band in channel["bands"]: log.log("Channel {}:\t\t {}[band: {}]".format( num_channel, channel["name"], band)) num_channel += 1 log.log("--- Output Classes ---") for c, classe in enumerate(config["classes"]): log.log("Class {}:\t\t {}".format(c, classe["title"])) log.log("--- Hyper Parameters ---") for hp in config["model"]: log.log("{}{}".format(hp.ljust(25, " "), config["model"][hp])) loader = load_module("robosat_pink.loaders.{}".format( config["model"]["loader"].lower())) loader_train = getattr(loader, config["model"]["loader"])( config, config["model"]["ts"], os.path.join(args.dataset, "training"), None, "train") loader_val = getattr(loader, config["model"]["loader"])( config, config["model"]["ts"], os.path.join(args.dataset, "validation"), None, "train") encoder = config["model"]["encoder"].lower() nn_module = load_module("robosat_pink.nn.{}".format( config["model"]["nn"].lower())) nn = getattr(nn_module, config["model"]["nn"])(loader_train.shape_in, loader_train.shape_out, encoder, config).to(device) nn = torch.nn.DataParallel(nn) optimizer = Adam(nn.parameters(), lr=config["model"]["lr"]) resume = 0 if args.checkpoint: chkpt = torch.load(os.path.expanduser(args.checkpoint), map_location=device) nn.load_state_dict(chkpt["state_dict"]) log.log("--- Using Checkpoint ---") log.log("Path:\t\t {}".format(args.checkpoint)) log.log("UUID:\t\t {}".format(chkpt["uuid"])) if args.resume: optimizer.load_state_dict(chkpt["optimizer"]) resume = chkpt["epoch"] assert resume < args.epochs, "Epoch asked, already reached by the given checkpoint" loss_module = load_module("robosat_pink.losses.{}".format( config["model"]["loss"].lower())) criterion = getattr(loss_module, config["model"]["loss"])().to(device) bs = config["model"]["bs"] train_loader = DataLoader(loader_train, batch_size=bs, shuffle=True, drop_last=True, num_workers=args.workers) val_loader = DataLoader(loader_val, batch_size=bs, shuffle=False, drop_last=True, num_workers=args.workers) if args.no_training: epoch = 0 process(val_loader, config, log, csv_val, epoch, device, nn, criterion, "eval") sys.exit() for epoch in range(resume + 1, args.epochs + 1): # 1-N based UUID = uuid.uuid1() log.log("---{}Epoch: {}/{} -- UUID: {}".format(os.linesep, epoch, args.epochs, UUID)) process(train_loader, config, log, csv_train, epoch, device, nn, criterion, "train", optimizer) try: # https://github.com/pytorch/pytorch/issues/9176 nn_doc = nn.module.doc nn_version = nn.module.version except AttributeError: nn_doc = nn.doc nn_version = nn.version states = { "uuid": UUID, "model_version": nn_version, "producer_name": "RoboSat.pink", "producer_version": rsp.__version__, "model_licence": "MIT", "domain": "pink.RoboSat", # reverse-DNS "doc_string": nn_doc, "shape_in": loader_train.shape_in, "shape_out": loader_train.shape_out, "state_dict": nn.state_dict(), "epoch": epoch, "nn": config["model"]["nn"], "encoder": config["model"]["encoder"], "optimizer": optimizer.state_dict(), "loader": config["model"]["loader"], } checkpoint_path = os.path.join(args.out, "checkpoint-{:05d}.pth".format(epoch)) if epoch == args.epochs or not (epoch % args.saving): log.log("[Saving checkpoint]") torch.save(states, checkpoint_path) if not args.no_validation: process(val_loader, config, log, csv_val, epoch, device, nn, criterion, "eval")
def main(args): config = load_config(args.config) check_channels(config) check_classes(config) palette = make_palette([classe["color"] for classe in config["classes"]]) args.workers = torch.cuda.device_count() * 2 if torch.device( "cuda") and not args.workers else args.workers cover = [tile for tile in tiles_from_csv(os.path.expanduser(args.cover)) ] if args.cover else None log = Logs(os.path.join(args.out, "log")) if torch.cuda.is_available(): log.log("RoboSat.pink - predict on {} GPUs, with {} workers".format( torch.cuda.device_count(), args.workers)) log.log("(Torch:{} Cuda:{} CudNN:{})".format( torch.__version__, torch.version.cuda, torch.backends.cudnn.version())) device = torch.device("cuda") torch.backends.cudnn.enabled = True torch.backends.cudnn.benchmark = True else: log.log("RoboSat.pink - predict on CPU, with {} workers".format( args.workers)) log.log("") log.log("============================================================") log.log("WARNING: Are you -really- sure about not predicting on GPU ?") log.log("============================================================") log.log("") device = torch.device("cpu") chkpt = torch.load(args.checkpoint, map_location=device) model_module = load_module("robosat_pink.models.{}".format( chkpt["nn"].lower())) nn = getattr(model_module, chkpt["nn"])(chkpt["shape_in"], chkpt["shape_out"]).to(device) nn = torch.nn.DataParallel(nn) nn.load_state_dict(chkpt["state_dict"]) nn.eval() log.log("Model {} - UUID: {}".format(chkpt["nn"], chkpt["uuid"])) mode = "predict" if not args.translate else "predict_translate" loader_module = load_module("robosat_pink.loaders.{}".format( chkpt["loader"].lower())) loader_predict = getattr(loader_module, chkpt["loader"])(config, chkpt["shape_in"][1:3], args.dataset, cover, mode=mode) loader = DataLoader(loader_predict, batch_size=args.bs, num_workers=args.workers) assert len(loader), "Empty predict dataset directory. Check your path." tiled = [] with torch.no_grad( ): # don't track tensors with autograd during prediction for images, tiles in tqdm(loader, desc="Eval", unit="batch", ascii=True): images = images.to(device) outputs = nn(images) probs = torch.nn.functional.softmax(outputs, dim=1).data.cpu().numpy() for tile, prob in zip(tiles, probs): x, y, z = list(map(int, tile)) mask = np.around(prob[1:, :, :]).astype(np.uint8).squeeze() if args.translate: tile_translate_to_file(args.out, mercantile.Tile(x, y, z), palette, mask) else: tile_label_to_file(args.out, mercantile.Tile(x, y, z), palette, mask) tiled.append(mercantile.Tile(x, y, z)) if not args.no_web_ui and not args.translate: template = "leaflet.html" if not args.web_ui_template else args.web_ui_template base_url = args.web_ui_base_url if args.web_ui_base_url else "." web_ui(args.out, base_url, tiled, tiled, "png", template)
def main(args): chkpt = torch.load(args.checkpoint, map_location=torch.device("cpu")) UUID = chkpt["uuid"] if "uuid" in chkpt else uuid.uuid1() try: nn_name = chkpt["nn"] except: assert args.nn, "--nn mandatory as not already in input .pth" nn_name = args.nn try: encoder = chkpt["encoder"] except: assert args.encoder, "--encoder mandatory as not already in input .pth" encoder = args.encoder try: loader = chkpt["loader"] except: assert args.loader, "--loader mandatory as not already in input .pth" doc_string = args.doc_string try: doc_string = chkpt["doc_string"] except: assert args.doc_string, "--doc_string mandatory as not already in input .pth" doc_string = args.doc_string try: shape_in = chkpt["shape_in"] except: assert args.shape_in, "--shape_in mandatory as not already in input .pth" shape_in = tuple(map(int, args.shape_in.split(","))) try: shape_out = chkpt["shape_out"] except: assert args.shape_out, "--shape_out mandatory as not already in input .pth" shape_out = tuple(map(int, args.shape_out.split(","))) nn_module = load_module("robosat_pink.nn.{}".format(nn_name.lower())) nn = getattr(nn_module, nn_name)(shape_in, shape_out, encoder.lower()).to("cpu") print("RoboSat.pink - export model to {}".format(args.type), file=sys.stderr) print("Model: {}".format(nn_name, file=sys.stderr)) print("UUID: {}".format(UUID, file=sys.stderr)) if args.type == "pth": states = { "uuid": UUID, "model_version": None, "producer_name": "RoboSat.pink", "producer_version": rsp.__version__, "model_licence": "MIT", "domain": "pink.RoboSat", # reverse-DNS "doc_string": doc_string, "shape_in": shape_in, "shape_out": shape_out, "state_dict": nn.state_dict(), "epoch": 0, "nn": nn_name, "encoder": encoder, "optimizer": None, "loader": loader, } torch.save(states, args.out) else: try: # https://github.com/pytorch/pytorch/issues/9176 nn.module.state_dict(chkpt["state_dict"]) except AttributeError: nn.state_dict(chkpt["state_dict"]) nn.eval() batch = torch.rand(1, *shape_in) if args.type == "onnx": torch.onnx.export( nn, torch.autograd.Variable(batch), args.out, input_names=["input", "shape_in", "shape_out"], output_names=["output"], dynamic_axes={ "input": { 0: "num_batch" }, "output": { 0: "num_batch" } }, ) if args.type == "jit": torch.jit.trace(nn, batch).save(args.out)
def main(args): config = load_config(args.config) args.out = os.path.expanduser(args.out) args.workers = torch.cuda.device_count() * 2 if torch.device("cuda") and not args.workers else args.workers config["model"]["loader"] = args.loader if args.loader else config["model"]["loader"] config["model"]["bs"] = args.bs if args.bs else config["model"]["bs"] config["model"]["lr"] = args.lr if args.lr else config["model"]["lr"] config["model"]["ts"] = args.ts if args.ts else config["model"]["ts"] config["model"]["nn"] = args.nn if args.nn else config["model"]["nn"] config["model"]["loss"] = args.loss if args.loss else config["model"]["loss"] config["model"]["da"] = args.da if args.da else config["model"]["da"] config["model"]["dap"] = args.dap if args.dap else config["model"]["dap"] check_classes(config) check_channels(config) check_model(config) if not os.path.isdir(os.path.expanduser(args.dataset)): sys.exit("ERROR: dataset {} is not a directory".format(args.dataset)) log = Logs(os.path.join(args.out, "log")) if torch.cuda.is_available(): log.log("RoboSat.pink - training on {} GPUs, with {} workers".format(torch.cuda.device_count(), args.workers)) log.log("(Torch:{} Cuda:{} CudNN:{})".format(torch.__version__, torch.version.cuda, torch.backends.cudnn.version())) device = torch.device("cuda") torch.backends.cudnn.benchmark = True else: log.log("RoboSat.pink - training on CPU, with {} workers - (Torch:{})".format(args.workers, torch.__version__)) log.log("WARNING: Are you really sure sure about not training on GPU ?") device = torch.device("cpu") loader = load_module("robosat_pink.loaders.{}".format(config["model"]["loader"].lower())) loader_train = getattr(loader, config["model"]["loader"])( config, config["model"]["ts"], os.path.join(args.dataset, "training"), "train" ) loader_val = getattr(loader, config["model"]["loader"])( config, config["model"]["ts"], os.path.join(args.dataset, "validation"), "train" ) model_module = load_module("robosat_pink.models.{}".format(config["model"]["nn"].lower())) nn = getattr(model_module, config["model"]["nn"])(loader_train.shape_in, loader_train.shape_out, config).to(device) nn = torch.nn.DataParallel(nn) optimizer = Adam(nn.parameters(), lr=config["model"]["lr"]) resume = 0 if args.checkpoint: chkpt = torch.load(os.path.expanduser(args.checkpoint), map_location=device) nn.load_state_dict(chkpt["state_dict"]) log.log("Using checkpoint: {}".format(args.checkpoint)) if args.resume: optimizer.load_state_dict(chkpt["optimizer"]) resume = chkpt["epoch"] if resume >= args.epochs: sys.exit("ERROR: Epoch {} already reached by the given checkpoint".format(config["model"]["epochs"])) loss_module = load_module("robosat_pink.losses.{}".format(config["model"]["loss"].lower())) criterion = getattr(loss_module, config["model"]["loss"])().to(device) bs = config["model"]["bs"] train_loader = DataLoader(loader_train, batch_size=bs, shuffle=True, drop_last=True, num_workers=args.workers) val_loader = DataLoader(loader_val, batch_size=bs, shuffle=False, drop_last=True, num_workers=args.workers) log.log("--- Input tensor from Dataset: {} ---".format(args.dataset)) num_channel = 1 # 1-based numerotation for channel in config["channels"]: for band in channel["bands"]: log.log("Channel {}:\t\t {}[band: {}]".format(num_channel, channel["name"], band)) num_channel += 1 log.log("--- Hyper Parameters ---") for hp in config["model"]: log.log("{}{}".format(hp.ljust(25, " "), config["model"][hp])) for epoch in range(resume, args.epochs): UUID = uuid.uuid1() log.log("---{}Epoch: {}/{} -- UUID: {}".format(os.linesep, epoch + 1, args.epochs, UUID)) process(train_loader, config, log, device, nn, criterion, "train", optimizer) if not args.no_validation: process(val_loader, config, log, device, nn, criterion, "eval") try: # https://github.com/pytorch/pytorch/issues/9176 nn_doc = nn.module.doc nn_version = nn.module.version except AttributeError: nn_version = nn.version nn_doc == nn.doc states = { "uuid": UUID, "model_version": nn_version, "producer_name": "RoboSat.pink", "producer_version": "0.4.0", "model_licence": "MIT", "domain": "pink.RoboSat", # reverse-DNS "doc_string": nn_doc, "shape_in": loader_train.shape_in, "shape_out": loader_train.shape_out, "state_dict": nn.state_dict(), "epoch": epoch + 1, "nn": config["model"]["nn"], "optimizer": optimizer.state_dict(), "loader": config["model"]["loader"], } checkpoint_path = os.path.join(args.out, "checkpoint-{:05d}.pth".format(epoch + 1)) torch.save(states, checkpoint_path)
def main(args): config = load_config(args.config) check_channels(config) check_classes(config) palette = make_palette([classe["color"] for classe in config["classes"]]) if not args.bs: try: args.bs = config["model"]["bs"] except: pass assert args.bs, "For rsp predict, model/bs must be set either in config file, or pass trought parameter --bs" args.workers = args.bs if not args.workers else args.workers cover = [tile for tile in tiles_from_csv(os.path.expanduser(args.cover)) ] if args.cover else None log = Logs(os.path.join(args.out, "log")) if torch.cuda.is_available(): log.log("RoboSat.pink - predict on {} GPUs, with {} workers".format( torch.cuda.device_count(), args.workers)) log.log("(Torch:{} Cuda:{} CudNN:{})".format( torch.__version__, torch.version.cuda, torch.backends.cudnn.version())) device = torch.device("cuda") torch.backends.cudnn.enabled = True torch.backends.cudnn.benchmark = True else: log.log("RoboSat.pink - predict on CPU, with {} workers".format( args.workers)) log.log("") log.log("============================================================") log.log("WARNING: Are you -really- sure about not predicting on GPU ?") log.log("============================================================") log.log("") device = torch.device("cpu") chkpt = torch.load(args.checkpoint, map_location=device) nn_module = load_module("robosat_pink.nn.{}".format(chkpt["nn"].lower())) nn = getattr(nn_module, chkpt["nn"])(chkpt["shape_in"], chkpt["shape_out"], chkpt["encoder"].lower()).to(device) nn = torch.nn.DataParallel(nn) nn.load_state_dict(chkpt["state_dict"]) nn.eval() log.log("Model {} - UUID: {}".format(chkpt["nn"], chkpt["uuid"])) with torch.no_grad( ): # don't track tensors with autograd during prediction tiled = [] if args.passes in ["first", "both"]: log.log("== Predict First Pass ==") tiled = predict(config, cover, args, palette, chkpt, nn, device, "predict") if args.passes in ["second", "both"]: log.log("== Predict Second Pass ==") predict(config, cover, args, palette, chkpt, nn, device, "predict_translate") if not args.no_web_ui and tiled: template = "leaflet.html" if not args.web_ui_template else args.web_ui_template base_url = args.web_ui_base_url if args.web_ui_base_url else "." web_ui(args.out, base_url, tiled, tiled, "png", template)