#!/usr/bin/env python3 # Author: Piero Toffanin # License: AGPLv3 import os import sys sys.path.insert(0, os.path.join("..", "..", os.path.dirname(__file__))) import rasterio import numpy as np import numpy.ma as ma import multiprocessing import argparse import functools from opensfm import dataset default_dem_path = "odm_dem/dsm.tif" default_outdir = "orthorectified" default_image_list = "img_list.txt" parser = argparse.ArgumentParser(description='Orthorectification Tool') parser.add_argument('dataset', type=str, help='Path to ODM dataset') parser.add_argument('--dem', type=str, default=default_dem_path, help='Absolute path to DEM to use to orthorectify images. Default: %(default)s') parser.add_argument('--no-alpha', type=bool, help="Don't output an alpha channel") parser.add_argument('--interpolation', type=str, choices=('nearest', 'bilinear'), default='bilinear', help="Type of interpolation to use to sample pixel values.Default: %(default)s") parser.add_argument('--outdir', type=str, default=default_outdir, help="Output directory where to store results. Default: %(default)s") parser.add_argument('--image-list', type=str, default=default_image_list, help="Path to file that contains the list of image filenames to orthorectify. By default all images in a dataset are processed. Default: %(default)s") parser.add_argument('--images', type=str, default="", help="Comma-separated list of filenames to rectify. Use as an alternative to --image-list. Default: process all images.") parser.add_argument('--threads', type=int, default=multiprocessing.cpu_count(), help="Number of CPU processes to use. Default: %(default)s") args = parser.parse_args() dataset_path = args.dataset dem_path = os.path.join(dataset_path, default_dem_path) if args.dem == default_dem_path else args.dem interpolation = args.interpolation with_alpha = not args.no_alpha image_list = os.path.join(dataset_path, default_image_list) if args.image_list == default_image_list else args.image_list cwd_path = os.path.join(dataset_path, default_outdir) if args.outdir == default_outdir else args.outdir if not os.path.exists(cwd_path): os.makedirs(cwd_path) target_images = [] # all if args.images: target_images = list(map(str.strip, args.images.split(","))) print("Processing %s images" % len(target_images)) elif args.image_list: with open(image_list) as f: target_images = list(filter(lambda filename: filename != '', map(str.strip, f.read().split("\n")))) print("Processing %s images" % len(target_images)) if not os.path.exists(dem_path): print("Whoops! %s does not exist. Provide a path to a valid DEM" % dem_path) exit(1) def bilinear_interpolate(im, x, y): x = np.asarray(x) y = np.asarray(y) x0 = np.floor(x).astype(int) x1 = x0 + 1 y0 = np.floor(y).astype(int) y1 = y0 + 1 x0 = np.clip(x0, 0, im.shape[1]-1) x1 = np.clip(x1, 0, im.shape[1]-1) y0 = np.clip(y0, 0, im.shape[0]-1) y1 = np.clip(y1, 0, im.shape[0]-1) Ia = im[ y0, x0 ] Ib = im[ y1, x0 ] Ic = im[ y0, x1 ] Id = im[ y1, x1 ] wa = (x1-x) * (y1-y) wb = (x1-x) * (y-y0) wc = (x-x0) * (y1-y) wd = (x-x0) * (y-y0) return wa*Ia + wb*Ib + wc*Ic + wd*Id # Read DEM print("Reading DEM: %s" % dem_path) with rasterio.open(dem_path) as dem_raster: dem = dem_raster.read()[0] dem_has_nodata = dem_raster.profile.get('nodata') is not None if dem_has_nodata: dem_min_value = ma.array(dem, mask=dem==dem_raster.nodata).min() else: dem_min_value = dem.min() print("DEM Minimum: %s" % dem_min_value) h, w = dem.shape crs = dem_raster.profile.get('crs') dem_offset_x, dem_offset_y = (0, 0) if crs: print("DEM has a CRS: %s" % str(crs)) # Read coords.txt coords_file = os.path.join(dataset_path, "odm_georeferencing", "coords.txt") if not os.path.exists(coords_file): print("Whoops! Cannot find %s (we need that!)" % coords_file) exit(1) with open(coords_file) as f: line = f.readline() # discard # second line is a northing/easting offset line = f.readline().rstrip() dem_offset_x, dem_offset_y = map(float, line.split(" ")) print("DEM offset: (%s, %s)" % (dem_offset_x, dem_offset_y)) print("DEM dimensions: %sx%s pixels" % (w, h)) # Read reconstruction udata = dataset.UndistortedDataSet(dataset.DataSet(os.path.join(dataset_path, "opensfm"))) reconstructions = udata.load_undistorted_reconstruction() if len(reconstructions) == 0: raise Exception("No reconstructions available") max_workers = args.threads print("Using %s threads" % max_workers) reconstruction = reconstructions[0] for shot in reconstruction.shots.values(): if len(target_images) == 0 or shot.id in target_images: print("Processing %s..." % shot.id) shot_image = udata.load_undistorted_image(shot.id) r = shot.pose.get_rotation_matrix() Xs, Ys, Zs = shot.pose.get_origin() a1 = r[0][0] b1 = r[0][1] c1 = r[0][2] a2 = r[1][0] b2 = r[1][1] c2 = r[1][2] a3 = r[2][0] b3 = r[2][1] c3 = r[2][2] print("Camera pose: (%f, %f, %f)" % (Xs, Ys, Zs)) img_h, img_w, num_bands = shot_image.shape print("Image dimensions: %sx%s pixels" % (img_w, img_h)) f = shot.camera.focal * max(img_h, img_w) has_nodata = dem_raster.profile.get('nodata') is not None def process_pixels(step): imgout = np.full((num_bands, dem_bbox_h, dem_bbox_w), np.nan) minx = dem_bbox_w miny = dem_bbox_h maxx = 0 maxy = 0 for j in range(dem_bbox_miny, dem_bbox_maxy + 1): if j % max_workers == step: im_j = j - dem_bbox_miny for i in range(dem_bbox_minx, dem_bbox_maxx + 1): im_i = i - dem_bbox_minx # World coordinates Xa, Ya = dem_raster.xy(j, i) Za = dem[j][i] # Skip nodata if has_nodata and Za == dem_raster.nodata: continue # Remove offset (our cameras don't have the geographic offset) Xa -= dem_offset_x Ya -= dem_offset_y # Colinearity function http://web.pdx.edu/~jduh/courses/geog493f14/Week03.pdf dx = (Xa - Xs) dy = (Ya - Ys) dz = (Za - Zs) den = a3 * dx + b3 * dy + c3 * dz x = (img_w - 1) / 2.0 - (f * (a1 * dx + b1 * dy + c1 * dz) / den) y = (img_h - 1) / 2.0 - (f * (a2 * dx + b2 * dy + c2 * dz) / den) if x >= 0 and y >= 0 and x <= img_w - 1 and y <= img_h - 1: if interpolation == 'bilinear': xi = img_w - 1 - x yi = img_h - 1 - y values = bilinear_interpolate(shot_image, xi, yi) else: # nearest xi = img_w - 1 - int(round(x)) yi = img_h - 1 - int(round(y)) values = shot_image[yi][xi] # We don't consider all zero values (pure black) # to be valid sample values. This will sometimes miss # valid sample values. if not np.all(values == 0): minx = min(minx, im_i) miny = min(miny, im_j) maxx = max(maxx, im_i) maxy = max(maxy, im_j) for b in range(num_bands): imgout[b][im_j][im_i] = values[b] # for b in range(num_bands): # minx = min(minx, im_i) # miny = min(miny, im_j) # maxx = max(maxx, im_i) # maxy = max(maxy, im_j) # imgout[b][im_j][im_i] = 255 return (imgout, (minx, miny, maxx, maxy)) # Compute bounding box of image coverage # assuming a flat plane at Z = min Z # (Otherwise we have to scan the entire DEM) # The Xa,Ya equations are just derived from the colinearity equations # solving for Xa and Ya instead of x,y def dem_coordinates(cpx, cpy): """ :param cpx principal point X (image coordinates) :param cpy principal point Y (image coordinates) """ Za = dem_min_value m = (a3*b1*cpy - a1*b3*cpy - (a3*b2 - a2*b3)*cpx - (a2*b1 - a1*b2)*f) Xa = dem_offset_x + (m*Xs + (b3*c1*cpy - b1*c3*cpy - (b3*c2 - b2*c3)*cpx - (b2*c1 - b1*c2)*f)*Za - (b3*c1*cpy - b1*c3*cpy - (b3*c2 - b2*c3)*cpx - (b2*c1 - b1*c2)*f)*Zs)/m Ya = dem_offset_y + (m*Ys - (a3*c1*cpy - a1*c3*cpy - (a3*c2 - a2*c3)*cpx - (a2*c1 - a1*c2)*f)*Za + (a3*c1*cpy - a1*c3*cpy - (a3*c2 - a2*c3)*cpx - (a2*c1 - a1*c2)*f)*Zs)/m y, x = dem_raster.index(Xa, Ya) return (x, y) dem_ul = dem_coordinates(-(img_w - 1) / 2.0, -(img_h - 1) / 2.0) dem_ur = dem_coordinates((img_w - 1) / 2.0, -(img_h - 1) / 2.0) dem_lr = dem_coordinates((img_w - 1) / 2.0, (img_h - 1) / 2.0) dem_ll = dem_coordinates(-(img_w - 1) / 2.0, (img_h - 1) / 2.0) dem_bbox = [dem_ul, dem_ur, dem_lr, dem_ll] dem_bbox_x = np.array(list(map(lambda xy: xy[0], dem_bbox))) dem_bbox_y = np.array(list(map(lambda xy: xy[1], dem_bbox))) dem_bbox_minx = min(w - 1, max(0, dem_bbox_x.min())) dem_bbox_miny = min(h - 1, max(0, dem_bbox_y.min())) dem_bbox_maxx = min(w - 1, max(0, dem_bbox_x.max())) dem_bbox_maxy = min(h - 1, max(0, dem_bbox_y.max())) dem_bbox_w = 1 + dem_bbox_maxx - dem_bbox_minx dem_bbox_h = 1 + dem_bbox_maxy - dem_bbox_miny print("Iterating over DEM box: [(%s, %s), (%s, %s)] (%sx%s pixels)" % (dem_bbox_minx, dem_bbox_miny, dem_bbox_maxx, dem_bbox_maxy, dem_bbox_w, dem_bbox_h)) if max_workers > 1: with multiprocessing.Pool(max_workers) as p: results = p.map(process_pixels, range(max_workers)) else: results = [process_pixels(0)] results = list(filter(lambda r: r[1][0] <= r[1][2] and r[1][1] <= r[1][3], results)) # Merge image imgout, _ = results[0] for j in range(dem_bbox_miny, dem_bbox_maxy + 1): im_j = j - dem_bbox_miny resimg, _ = results[j % max_workers] for b in range(num_bands): imgout[b][im_j] = resimg[b][im_j] # Merge bounds minx = dem_bbox_w miny = dem_bbox_h maxx = 0 maxy = 0 for _, bounds in results: minx = min(bounds[0], minx) miny = min(bounds[1], miny) maxx = max(bounds[2], maxx) maxy = max(bounds[3], maxy) print("Output bounds: (%s, %s), (%s, %s) pixels" % (minx, miny, maxx, maxy)) if minx <= maxx and miny <= maxy: imgout = imgout[:,miny:maxy,minx:maxx] if with_alpha: alpha = np.zeros((imgout.shape[1], imgout.shape[2]), dtype=np.uint8) # Set all not-NaN indices to 255 alpha[~np.isnan(imgout[0])] = 255 # Cast imgout = imgout.astype(shot_image.dtype) dem_transform = dem_raster.profile['transform'] offset_x, offset_y = dem_raster.xy(dem_bbox_miny + miny, dem_bbox_minx + minx, offset='ul') profile = { 'driver': 'GTiff', 'width': imgout.shape[2], 'height': imgout.shape[1], 'count': num_bands + 1 if with_alpha else num_bands, 'dtype': imgout.dtype.name, 'transform': rasterio.transform.Affine(dem_transform[0], dem_transform[1], offset_x, dem_transform[3], dem_transform[4], offset_y), 'nodata': None, 'crs': crs } outfile = os.path.join(cwd_path, shot.id) if not outfile.endswith(".tif"): outfile = outfile + ".tif" with rasterio.open(outfile, 'w', **profile) as wout: for b in range(num_bands): wout.write(imgout[b], b + 1) if with_alpha: wout.write(alpha, num_bands + 1) print("Wrote %s" % outfile) else: print("Cannot orthorectify image (is the image inside the DEM bounds?)")