diff --git a/contrib/orthorectify/orthorectify.py b/contrib/orthorectify/orthorectify.py index 36dda50d..9ce9ec82 100644 --- a/contrib/orthorectify/orthorectify.py +++ b/contrib/orthorectify/orthorectify.py @@ -1,8 +1,12 @@ #!/usr/bin/env python3 -import rasterio import os +import sys +sys.path.insert(0, os.path.join("..", "..", os.path.dirname(__file__))) + +import rasterio import numpy as np from opensfm import dataset +import multiprocessing # TODO: command argument parser @@ -12,6 +16,7 @@ dem_path = "/datasets/brighton2/odm_meshing/tmp/mesh_dsm.tif" target_images = [] # all target_images.append("DJI_0018.JPG") + # Read DSM print("Reading DSM: %s" % dem_path) with rasterio.open(dem_path) as dem_raster: @@ -26,52 +31,59 @@ with rasterio.open(dem_path) as dem_raster: if len(reconstructions) == 0: raise Exception("No reconstructions available") + max_workers = multiprocessing.cpu_count() + 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("Loading %s..." % shot.id) + 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() cam_w = shot.camera.width cam_h = shot.camera.height + img_w, img_h, num_bands = shot_image.shape f = shot.camera.focal - num_bands = shot_image.shape[2] - imgout = np.full((num_bands, w, h), np.nan, dtype=shot_image.dtype) + def process_pixels(step): + imgout = np.full((num_bands, w, h), np.nan, dtype=shot_image.dtype) + for i in range(w): + if i % max_workers == step: + for j in range(h): + # i, j = (168, 442) # TODO REMOVE + # World coordinates + Xa, Ya = dem_raster.xy(i, j) + Za = dem[j][i] + + # Colinearity function http://web.pdx.edu/~jduh/courses/geog493f14/Week03.pdf + dx = (Xa - Xs) + dy = (Ya - Ys) + dz = (Za - Zs) + + den = r[0][2] * dx + r[1][2] * dy + r[2][2] * dz + x = 0.5 + (-f * (r[0][0] * dx + r[1][0] * dy + r[2][0] * dz) / den) + y = 0.5 + (-f * (r[0][1] * dx + r[1][1] * dy + r[2][1] * dz) / den) + + # TEST + if x >= 0 and y >= 0 and x <= 1.0 and y <= 1.0: + for b in range(3): + imgout[b][i][j] = 255 + return imgout + + with multiprocessing.Pool(max_workers) as p: + results = p.map(process_pixels, range(max_workers)) + + # Merge + imgout = results[0] for i in range(w): - for j in range(h): + for b in range(num_bands): + imgout[b][i] = results[i % max_workers][b][i] - # i, j = (168, 442) # TODO REMOVE - - # World coordinates - Xa, Ya = dem_raster.xy(j, i) - Za = dem[j][i] - - # Colinearity function http://web.pdx.edu/~jduh/courses/geog493f14/Week03.pdf - dx = (Xa - Xs) - dy = (Ya - Ys) - dz = (Za - Zs) - - den = r[0][2] * dx + r[1][2] * dy + r[2][2] * dz - x = 0.5 + (-f * (r[0][0] * dx + r[1][0] * dy + r[2][0] * dz) / den) - y = 0.5 + (-f * (r[0][1] * dx + r[1][1] * dy + r[2][1] * dz) / den) - - # TODO: interpolate? - # xi = round(x * cam_w) - # yi = round(y * cam_h) - - for b in range(3): - - # TEST - if x > 0 and y > 0: - imgout[b][i][j] = 255 - - imgout[imgout == np.nan] = 0 print(w) print(h) @@ -85,5 +97,4 @@ with rasterio.open(dem_path) as dem_raster: } with rasterio.open("/datasets/brighton2/odm_meshing/tmp/out.tif", 'w', **profile) as wout: for b in range(num_bands): - wout.write(imgout[b], b + 1) - \ No newline at end of file + wout.write(imgout[b], b + 1) \ No newline at end of file