def test_http(self):
download_dir = '/download_dir'
uri = 'http://bucket/my/file.txt'
path = get_local_path(uri, download_dir)
self.assertEqual(path, '/download_dir/http/bucket/my/file.txt')
# simulate a zxy tile URI
uri = 'http://bucket/10/25/53?auth=426753'
path = get_local_path(uri, download_dir)
self.assertEqual(path, '/download_dir/http/bucket/10/25/53')
def unzip_data(self) -> List[str]:
"""Unzip dataset zip files.
Returns:
paths to directories that each contain contents of one zip file
"""
cfg = self.cfg
if cfg.data.uri.startswith('s3://') or cfg.data.uri.startswith('/'):
data_uri = cfg.data.uri
else:
data_uri = join(cfg.base_uri, cfg.data.uri)
data_dirs = []
zip_uris = [data_uri] if data_uri.endswith('.zip') else list_paths(
data_uri, 'zip')
for zip_ind, zip_uri in enumerate(zip_uris):
zip_path = get_local_path(zip_uri, self.data_cache_dir)
if not isfile(zip_path):
zip_path = download_if_needed(zip_uri, self.data_cache_dir)
with zipfile.ZipFile(zip_path, 'r') as zipf:
data_dir = join(self.tmp_dir, 'data', str(zip_ind))
data_dirs.append(data_dir)
zipf.extractall(data_dir)
return data_dirs
def unzip_data(self, uri: Union[str, List[str]]) -> List[str]:
"""Unzip dataset zip files.
Args:
uri: a list of URIs of zip files or the URI of a directory containing
zip files
Returns:
paths to directories that each contain contents of one zip file
"""
cfg = self.cfg
data_dirs = []
if isinstance(uri, list):
zip_uris = uri
else:
# TODO generalize this to work with any file system
if uri.startswith('s3://') or uri.startswith('/'):
data_uri = uri
else:
data_uri = join(cfg.base_uri, uri)
zip_uris = ([data_uri] if data_uri.endswith('.zip') else
list_paths(data_uri, 'zip'))
for zip_ind, zip_uri in enumerate(zip_uris):
zip_path = get_local_path(zip_uri, self.data_cache_dir)
if not isfile(zip_path):
zip_path = download_if_needed(zip_uri, self.data_cache_dir)
with zipfile.ZipFile(zip_path, 'r') as zipf:
data_dir = join(self.tmp_dir, 'data', str(uuid.uuid4()),
str(zip_ind))
data_dirs.append(data_dir)
zipf.extractall(data_dir)
return data_dirs
def crop_image(image_uri, window, crop_uri):
im_dataset = rasterio.open(image_uri)
rasterio_window = window.rasterio_format()
im = im_dataset.read(window=rasterio_window)
with tempfile.TemporaryDirectory() as tmp_dir:
crop_path = get_local_path(crop_uri, tmp_dir)
make_dir(crop_path, use_dirname=True)
meta = im_dataset.meta
meta['width'], meta['height'] = window.get_width(), window.get_height()
meta['transform'] = rasterio.windows.transform(rasterio_window,
im_dataset.transform)
with rasterio.open(crop_path, 'w', **meta) as dst:
dst.colorinterp = im_dataset.colorinterp
dst.write(im)
upload_or_copy(crop_path, crop_uri)
def __init__(self,
cfg: LearnerConfig,
tmp_dir: str,
model_path: Optional[str] = None):
"""Constructor.
Args:
cfg: configuration
tmp_dir: root of temp dirs
model_path: a local path to model weights. If provided, the model is loaded
and it is assumed that this Learner will be used for prediction only.
"""
self.cfg = cfg
self.tmp_dir = tmp_dir
# TODO make cache dirs configurable
torch_cache_dir = '/opt/data/torch-cache'
os.environ['TORCH_HOME'] = torch_cache_dir
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
self.data_cache_dir = '/opt/data/data-cache'
make_dir(self.data_cache_dir)
self.model = self.build_model()
self.model.to(self.device)
if model_path is not None:
if isfile(model_path):
self.model.load_state_dict(
torch.load(model_path, map_location=self.device))
else:
raise Exception(
'Model could not be found at {}'.format(model_path))
self.model.eval()
else:
log.info(self.cfg)
# ds = dataset, dl = dataloader
self.train_ds = None
self.train_dl = None
self.valid_ds = None
self.valid_dl = None
self.test_ds = None
self.test_dl = None
if cfg.output_uri.startswith('s3://'):
self.output_dir = get_local_path(cfg.output_uri, tmp_dir)
make_dir(self.output_dir, force_empty=True)
if not cfg.overfit_mode:
self.sync_from_cloud()
else:
self.output_dir = cfg.output_uri
make_dir(self.output_dir)
self.last_model_path = join(self.output_dir, 'last-model.pth')
self.config_path = join(self.output_dir, 'config.json')
self.train_state_path = join(self.output_dir, 'train-state.json')
self.log_path = join(self.output_dir, 'log.csv')
model_bundle_fn = basename(cfg.get_model_bundle_uri())
self.model_bundle_path = join(self.output_dir, model_bundle_fn)
self.metric_names = self.build_metric_names()
json_to_file(self.cfg.dict(), self.config_path)
self.load_init_weights()
self.load_checkpoint()
self.opt = self.build_optimizer()
self.setup_data()
self.start_epoch = self.get_start_epoch()
self.steps_per_epoch = len(
self.train_ds) // self.cfg.solver.batch_sz
self.step_scheduler = self.build_step_scheduler()
self.epoch_scheduler = self.build_epoch_scheduler()
self.setup_tensorboard()
def setup_data(self):
cfg = self.cfg
batch_sz = cfg.solver.batch_sz
num_workers = cfg.data.num_workers
# download and unzip data
if cfg.data.uri.startswith('s3://') or cfg.data.uri.startswith('/'):
data_uri = cfg.data.uri
else:
data_uri = join(cfg.base_uri, cfg.data.uri)
data_dirs = []
zip_uris = [data_uri] if data_uri.endswith('.zip') else list_paths(
data_uri, 'zip')
for zip_ind, zip_uri in enumerate(zip_uris):
zip_path = get_local_path(zip_uri, self.data_cache_dir)
if not isfile(zip_path):
zip_path = download_if_needed(zip_uri, self.data_cache_dir)
with zipfile.ZipFile(zip_path, 'r') as zipf:
data_dir = join(self.tmp_dir, 'data', str(zip_ind))
data_dirs.append(data_dir)
zipf.extractall(data_dir)
# build datasets -- one per zip file and then merge them into a single dataset
train_ds = []
valid_ds = []
test_ds = []
for data_dir in data_dirs:
train_dir = join(data_dir, 'train')
valid_dir = join(data_dir, 'valid')
transform = Compose(
[Resize((cfg.data.img_sz, cfg.data.img_sz)),
ToTensor()])
aug_transform = Compose([
RandomHorizontalFlip(),
RandomVerticalFlip(),
ColorJitter(0.1, 0.1, 0.1, 0.1),
Resize((cfg.data.img_sz, cfg.data.img_sz)),
ToTensor()
])
if isdir(train_dir):
if cfg.overfit_mode:
train_ds.append(
ImageRegressionDataset(train_dir,
cfg.data.class_names,
transform=transform))
else:
train_ds.append(
ImageRegressionDataset(train_dir,
cfg.data.class_names,
transform=aug_transform))
if isdir(valid_dir):
valid_ds.append(
ImageRegressionDataset(valid_dir,
cfg.data.class_names,
transform=transform))
test_ds.append(
ImageRegressionDataset(valid_dir,
cfg.data.class_names,
transform=transform))
train_ds, valid_ds, test_ds = \
ConcatDataset(train_ds), ConcatDataset(valid_ds), ConcatDataset(test_ds)
if cfg.overfit_mode:
train_ds = Subset(train_ds, range(batch_sz))
valid_ds = train_ds
test_ds = train_ds
elif cfg.test_mode:
train_ds = Subset(train_ds, range(batch_sz))
valid_ds = Subset(valid_ds, range(batch_sz))
test_ds = Subset(test_ds, range(batch_sz))
train_dl = DataLoader(train_ds,
shuffle=True,
batch_size=batch_sz,
num_workers=num_workers,
pin_memory=True)
valid_dl = DataLoader(valid_ds,
shuffle=True,
batch_size=batch_sz,
num_workers=num_workers,
pin_memory=True)
test_dl = DataLoader(test_ds,
shuffle=True,
batch_size=batch_sz,
num_workers=num_workers,
pin_memory=True)
self.train_ds, self.valid_ds, self.test_ds = (train_ds, valid_ds,
test_ds)
self.train_dl, self.valid_dl, self.test_dl = (train_dl, valid_dl,
test_dl)
def test_s3(self):
download_dir = '/download_dir'
uri = 's3://bucket/my/file.txt'
path = get_local_path(uri, download_dir)
self.assertEqual(path, '/download_dir/s3/bucket/my/file.txt')
def test_local(self):
download_dir = '/download_dir'
uri = '/my/file.txt'
path = get_local_path(uri, download_dir)
self.assertEqual(path, uri)
def save(self, labels):
"""Save.
Args:
labels - (SemanticSegmentationLabels) labels to be saved
"""
local_path = get_local_path(self.uri, self.tmp_dir)
make_dir(local_path, use_dirname=True)
transform = self.crs_transformer.get_affine_transform()
crs = self.crs_transformer.get_image_crs()
band_count = 1
dtype = np.uint8
if self.class_trans:
band_count = 3
mask = (np.zeros((self.extent.ymax, self.extent.xmax), dtype=np.uint8)
if self.vector_output else None)
# https://github.com/mapbox/rasterio/blob/master/docs/quickstart.rst
# https://rasterio.readthedocs.io/en/latest/topics/windowed-rw.html
with rasterio.open(local_path,
'w',
driver='GTiff',
height=self.extent.ymax,
width=self.extent.xmax,
count=band_count,
dtype=dtype,
transform=transform,
crs=crs) as dataset:
for window in labels.get_windows():
label_arr = labels.get_label_arr(window)
window = window.intersection(self.extent)
label_arr = label_arr[0:window.get_height(),
0:window.get_width()]
if mask is not None:
mask[window.ymin:window.ymax,
window.xmin:window.xmax] = label_arr
window = window.rasterio_format()
if self.class_trans:
rgb_labels = self.class_trans.class_to_rgb(label_arr)
for chan in range(3):
dataset.write_band(chan + 1,
rgb_labels[:, :, chan],
window=window)
else:
img = label_arr.astype(dtype)
dataset.write_band(1, img, window=window)
upload_or_copy(local_path, self.uri)
if self.vector_output:
import mask_to_polygons.vectorification as vectorification
import mask_to_polygons.processing.denoise as denoise
for vo in self.vector_output:
denoise_radius = vo.denoise
uri = vo.uri
mode = vo.get_mode()
class_id = vo.class_id
class_mask = np.array(mask == class_id, dtype=np.uint8)
local_geojson_path = get_local_path(uri, self.tmp_dir)
def transform(x, y):
return self.crs_transformer.pixel_to_map((x, y))
if denoise_radius > 0:
class_mask = denoise.denoise(class_mask, denoise_radius)
if uri and mode == 'buildings':
geojson = vectorification.geojson_from_mask(
mask=class_mask,
transform=transform,
mode=mode,
min_aspect_ratio=vo.min_aspect_ratio,
min_area=vo.min_area,
width_factor=vo.element_width_factor,
thickness=vo.element_thickness)
elif uri and mode == 'polygons':
geojson = vectorification.geojson_from_mask(
mask=class_mask, transform=transform, mode=mode)
if local_geojson_path:
with open(local_geojson_path, 'w') as file_out:
file_out.write(geojson)
upload_or_copy(local_geojson_path, uri)
def _zxy2geotiff(tile_schema, zoom, bounds, output_uri, make_cog=False):
"""Generates a GeoTIFF of a bounded region from a ZXY tile server.
Args:
tile_schema: (str) the URI schema for zxy tiles (ie. a slippy map tile server)
of the form /tileserver-uri/{z}/{x}/{y}.png. If {-y} is used, the tiles
are assumed to be indexed using TMS coordinates, where the y axis starts
at the southernmost point. The URI can be for http, S3, or the local
file system.
zoom: (int) the zoom level to use when retrieving tiles
bounds: (list) a list of length 4 containing min_lat, min_lng,
max_lat, max_lng
output_uri: (str) where to save the GeoTIFF. The URI can be for http, S3, or the
local file system
"""
min_lat, min_lng, max_lat, max_lng = bounds
if min_lat >= max_lat:
raise ValueError('min_lat must be < max_lat')
if min_lng >= max_lng:
raise ValueError('min_lng must be < max_lng')
is_tms = False
if '{-y}' in tile_schema:
tile_schema = tile_schema.replace('{-y}', '{y}')
is_tms = True
tmp_dir_obj = tempfile.TemporaryDirectory()
tmp_dir = tmp_dir_obj.name
# Get range of tiles that cover bounds.
output_path = get_local_path(output_uri, tmp_dir)
tile_sz = 256
t = mercantile.tile(min_lng, max_lat, zoom)
xmin, ymin = t.x, t.y
t = mercantile.tile(max_lng, min_lat, zoom)
xmax, ymax = t.x, t.y
# The supplied bounds are contained within the "tile bounds" -- ie. the
# bounds of the set of tiles that covers the supplied bounds. Therefore,
# we need to crop out the imagery that lies within the supplied bounds.
# We do this by computing a top, bottom, left, and right offset in pixel
# units of the supplied bounds against the tile bounds. Getting the offsets
# in pixel units involves converting lng/lat to web mercator units since we
# assume that is the CRS of the tiles. These offsets are then used to crop
# individual tiles and place them correctly into the output raster.
nw_merc_x, nw_merc_y = lnglat2merc(min_lng, max_lat)
left_pix_offset, top_pix_offset = merc2pixel(xmin, ymin, zoom, nw_merc_x,
nw_merc_y)
se_merc_x, se_merc_y = lnglat2merc(max_lng, min_lat)
se_left_pix_offset, se_top_pix_offset = merc2pixel(xmax, ymax, zoom,
se_merc_x, se_merc_y)
right_pix_offset = tile_sz - se_left_pix_offset
bottom_pix_offset = tile_sz - se_top_pix_offset
uncropped_height = tile_sz * (ymax - ymin + 1)
uncropped_width = tile_sz * (xmax - xmin + 1)
height = uncropped_height - top_pix_offset - bottom_pix_offset
width = uncropped_width - left_pix_offset - right_pix_offset
transform = rasterio.transform.from_bounds(nw_merc_x, se_merc_y, se_merc_x,
nw_merc_y, width, height)
with rasterio.open(
output_path,
'w',
driver='GTiff',
height=height,
width=width,
count=3,
crs='epsg:3857',
transform=transform,
dtype=rasterio.uint8) as dataset:
out_x = 0
for xi, x in enumerate(range(xmin, xmax + 1)):
tile_xmin, tile_xmax = 0, tile_sz - 1
if x == xmin:
tile_xmin += left_pix_offset
if x == xmax:
tile_xmax -= right_pix_offset
window_width = tile_xmax - tile_xmin + 1
out_y = 0
for yi, y in enumerate(range(ymin, ymax + 1)):
tile_ymin, tile_ymax = 0, tile_sz - 1
if y == ymin:
tile_ymin += top_pix_offset
if y == ymax:
tile_ymax -= bottom_pix_offset
window_height = tile_ymax - tile_ymin + 1
# Convert from xyz to tms if needed.
# https://gist.github.com/tmcw/4954720
if is_tms:
y = (2**zoom) - y - 1
tile_uri = tile_schema.format(x=x, y=y, z=zoom)
tile_path = download_if_needed(tile_uri, tmp_dir)
img = np.array(Image.open(tile_path))
img = img[tile_ymin:tile_ymax + 1, tile_xmin:tile_xmax + 1, :]
window = Window(out_x, out_y, window_width, window_height)
dataset.write(
np.transpose(img[:, :, 0:3], (2, 0, 1)), window=window)
out_y += window_height
out_x += window_width
if make_cog:
create_cog(output_path, output_uri, tmp_dir)
else:
upload_or_copy(output_path, output_uri)