kopia lustrzana https://github.com/OpenDroneMap/ODM
338 wiersze
16 KiB
Python
338 wiersze
16 KiB
Python
import os
|
|
import json
|
|
|
|
from opendm import context
|
|
from opendm import io
|
|
from opendm import types
|
|
from opendm.photo import PhotoCorruptedException
|
|
from opendm import log
|
|
from opendm import system
|
|
from opendm.geo import GeoFile
|
|
from shutil import copyfile
|
|
from opendm import progress
|
|
from opendm import boundary
|
|
from opendm import ai
|
|
from opendm.skyremoval.skyfilter import SkyFilter
|
|
from opendm.bgfilter import BgFilter
|
|
from opendm.concurrency import parallel_map
|
|
from opendm.video.video2dataset import Parameters, Video2Dataset
|
|
|
|
def save_images_database(photos, database_file):
|
|
with open(database_file, 'w') as f:
|
|
f.write(json.dumps([p.__dict__ for p in photos]))
|
|
|
|
log.ODM_INFO("Wrote images database: %s" % database_file)
|
|
|
|
def load_images_database(database_file):
|
|
# Empty is used to create types.ODM_Photo class
|
|
# instances without calling __init__
|
|
class Empty:
|
|
pass
|
|
|
|
result = []
|
|
|
|
log.ODM_INFO("Loading images database: %s" % database_file)
|
|
|
|
with open(database_file, 'r') as f:
|
|
photos_json = json.load(f)
|
|
for photo_json in photos_json:
|
|
p = Empty()
|
|
for k in photo_json:
|
|
setattr(p, k, photo_json[k])
|
|
p.__class__ = types.ODM_Photo
|
|
result.append(p)
|
|
|
|
return result
|
|
|
|
class ODMLoadDatasetStage(types.ODM_Stage):
|
|
def process(self, args, outputs):
|
|
outputs['start_time'] = system.now_raw()
|
|
tree = types.ODM_Tree(args.project_path, args.gcp, args.geo, args.align)
|
|
outputs['tree'] = tree
|
|
|
|
if io.file_exists(tree.benchmarking):
|
|
# Delete the previously made file
|
|
try:
|
|
os.remove(tree.benchmarking)
|
|
with open(tree.benchmarking, 'a') as b:
|
|
b.write('ODM Benchmarking file created %s\nNumber of Cores: %s\n\n' % (system.now(), context.num_cores))
|
|
except Exception as e:
|
|
log.ODM_WARNING("Cannot write benchmark file: %s" % str(e))
|
|
|
|
def valid_filename(filename, supported_extensions):
|
|
(pathfn, ext) = os.path.splitext(filename)
|
|
return ext.lower() in supported_extensions and pathfn[-5:] != "_mask"
|
|
|
|
# Get supported images from dir
|
|
def get_images(in_dir):
|
|
entries = os.listdir(in_dir)
|
|
valid, rejects = [], []
|
|
for f in entries:
|
|
if valid_filename(f, context.supported_extensions):
|
|
valid.append(f)
|
|
else:
|
|
rejects.append(f)
|
|
return valid, rejects
|
|
|
|
def search_video_files(in_dir):
|
|
entries = os.listdir(in_dir)
|
|
return [os.path.join(in_dir, f) for f in entries if valid_filename(f, context.supported_video_extensions)]
|
|
|
|
def find_mask(photo_path, masks):
|
|
(pathfn, ext) = os.path.splitext(os.path.basename(photo_path))
|
|
k = "{}_mask".format(pathfn)
|
|
|
|
mask = masks.get(k)
|
|
if mask:
|
|
# Spaces are not supported due to OpenSfM's mask_list.txt format reqs
|
|
if not " " in mask:
|
|
return mask
|
|
else:
|
|
log.ODM_WARNING("Image mask {} has a space. Spaces are currently not supported for image masks.".format(mask))
|
|
|
|
|
|
|
|
# get images directory
|
|
images_dir = tree.dataset_raw
|
|
|
|
# define paths and create working directories
|
|
system.mkdir_p(tree.odm_georeferencing)
|
|
|
|
log.ODM_INFO('Loading dataset from: %s' % images_dir)
|
|
|
|
# check if we rerun cell or not
|
|
images_database_file = os.path.join(tree.root_path, 'images.json')
|
|
if not io.file_exists(images_database_file) or self.rerun():
|
|
if not os.path.exists(images_dir):
|
|
raise system.ExitException("There are no images in %s! Make sure that your project path and dataset name is correct. The current is set to: %s" % (images_dir, args.project_path))
|
|
|
|
# Check if we need to extract video frames
|
|
frames_db_file = os.path.join(images_dir, 'frames.json')
|
|
if not os.path.exists(frames_db_file) or self.rerun():
|
|
video_files = search_video_files(images_dir)
|
|
|
|
# If we're re-running the pipeline, and frames have been extracted during a previous run
|
|
# we need to remove those before re-extracting them
|
|
if len(video_files) > 0 and os.path.exists(frames_db_file) and self.rerun():
|
|
log.ODM_INFO("Re-run, removing previously extracted video frames")
|
|
frames = []
|
|
try:
|
|
with open(frames_db_file, 'r') as f:
|
|
frames = json.loads(f.read())
|
|
except Exception as e:
|
|
log.ODM_WARNING("Cannot check previous video extraction: %s" % str(e))
|
|
|
|
for f in frames:
|
|
fp = os.path.join(images_dir, f)
|
|
if os.path.isfile(fp):
|
|
os.remove(fp)
|
|
|
|
if len(video_files) > 0:
|
|
log.ODM_INFO("Found video files (%s), extracting frames" % len(video_files))
|
|
|
|
try:
|
|
params = Parameters({
|
|
"input": video_files,
|
|
"output": images_dir,
|
|
|
|
"blur_threshold": 200,
|
|
"distance_threshold": 10,
|
|
"black_ratio_threshold": 0.98,
|
|
"pixel_black_threshold": 0.30,
|
|
"use_srt": True,
|
|
"max_dimension": args.video_resolution,
|
|
"limit": args.video_limit,
|
|
})
|
|
v2d = Video2Dataset(params)
|
|
frames = v2d.ProcessVideo()
|
|
|
|
with open(frames_db_file, 'w') as f:
|
|
f.write(json.dumps([os.path.basename(f) for f in frames]))
|
|
except Exception as e:
|
|
log.ODM_WARNING("Could not extract video frames: %s" % str(e))
|
|
|
|
files, rejects = get_images(images_dir)
|
|
if files:
|
|
# create ODMPhoto list
|
|
path_files = [os.path.join(images_dir, f) for f in files]
|
|
|
|
# Lookup table for masks
|
|
masks = {}
|
|
for r in rejects:
|
|
(p, ext) = os.path.splitext(r)
|
|
if p[-5:] == "_mask" and ext.lower() in context.supported_extensions:
|
|
masks[p] = r
|
|
|
|
photos = []
|
|
with open(tree.dataset_list, 'w') as dataset_list:
|
|
log.ODM_INFO("Loading %s images" % len(path_files))
|
|
for f in path_files:
|
|
try:
|
|
p = types.ODM_Photo(f)
|
|
p.set_mask(find_mask(f, masks))
|
|
photos.append(p)
|
|
dataset_list.write(photos[-1].filename + '\n')
|
|
except PhotoCorruptedException:
|
|
log.ODM_WARNING("%s seems corrupted and will not be used" % os.path.basename(f))
|
|
|
|
# Check if a geo file is available
|
|
if tree.odm_geo_file is not None and os.path.isfile(tree.odm_geo_file):
|
|
log.ODM_INFO("Found image geolocation file")
|
|
gf = GeoFile(tree.odm_geo_file)
|
|
updated = 0
|
|
for p in photos:
|
|
entry = gf.get_entry(p.filename)
|
|
if entry:
|
|
p.update_with_geo_entry(entry)
|
|
p.compute_opk()
|
|
updated += 1
|
|
log.ODM_INFO("Updated %s image positions" % updated)
|
|
|
|
# GPSDOP override if we have GPS accuracy information (such as RTK)
|
|
if 'gps_accuracy_is_set' in args:
|
|
log.ODM_INFO("Forcing GPS DOP to %s for all images" % args.gps_accuracy)
|
|
|
|
for p in photos:
|
|
p.override_gps_dop(args.gps_accuracy)
|
|
|
|
# Override projection type
|
|
if args.camera_lens != "auto":
|
|
log.ODM_INFO("Setting camera lens to %s for all images" % args.camera_lens)
|
|
|
|
for p in photos:
|
|
p.override_camera_projection(args.camera_lens)
|
|
|
|
# Automatic sky removal
|
|
if args.sky_removal:
|
|
# For each image that :
|
|
# - Doesn't already have a mask, AND
|
|
# - Is not nadir (or if orientation info is missing, or if camera lens is fisheye), AND
|
|
# - There are no spaces in the image filename (OpenSfM requirement)
|
|
# Automatically generate a sky mask
|
|
|
|
# Generate list of sky images
|
|
sky_images = []
|
|
for p in photos:
|
|
if p.mask is None and (args.camera_lens in ['fisheye', 'spherical'] or p.pitch is None or (abs(p.pitch) > 20)) and (not " " in p.filename):
|
|
sky_images.append({'file': os.path.join(images_dir, p.filename), 'p': p})
|
|
|
|
if len(sky_images) > 0:
|
|
log.ODM_INFO("Automatically generating sky masks for %s images" % len(sky_images))
|
|
model = ai.get_model("skyremoval", "https://github.com/OpenDroneMap/SkyRemoval/releases/download/v1.0.5/model.zip", "v1.0.5")
|
|
if model is not None:
|
|
sf = SkyFilter(model=model)
|
|
|
|
def parallel_sky_filter(item):
|
|
try:
|
|
mask_file = sf.run_img(item['file'], images_dir)
|
|
|
|
# Check and set
|
|
if mask_file is not None and os.path.isfile(mask_file):
|
|
item['p'].set_mask(os.path.basename(mask_file))
|
|
log.ODM_INFO("Wrote %s" % os.path.basename(mask_file))
|
|
else:
|
|
log.ODM_WARNING("Cannot generate mask for %s" % img)
|
|
except Exception as e:
|
|
log.ODM_WARNING("Cannot generate mask for %s: %s" % (img, str(e)))
|
|
|
|
parallel_map(parallel_sky_filter, sky_images, max_workers=args.max_concurrency)
|
|
|
|
log.ODM_INFO("Sky masks generation completed!")
|
|
else:
|
|
log.ODM_WARNING("Cannot load AI model (you might need to be connected to the internet?)")
|
|
else:
|
|
log.ODM_INFO("No sky masks will be generated (masks already provided, or images are nadir)")
|
|
|
|
# End sky removal
|
|
|
|
# Automatic background removal
|
|
if args.bg_removal:
|
|
# For each image that :
|
|
# - Doesn't already have a mask, AND
|
|
# - There are no spaces in the image filename (OpenSfM requirement)
|
|
|
|
# Generate list of sky images
|
|
bg_images = []
|
|
for p in photos:
|
|
if p.mask is None and (not " " in p.filename):
|
|
bg_images.append({'file': os.path.join(images_dir, p.filename), 'p': p})
|
|
|
|
if len(bg_images) > 0:
|
|
log.ODM_INFO("Automatically generating background masks for %s images" % len(bg_images))
|
|
model = ai.get_model("bgremoval", "https://github.com/OpenDroneMap/ODM/releases/download/v2.9.0/u2net.zip", "v2.9.0")
|
|
if model is not None:
|
|
bg = BgFilter(model=model)
|
|
|
|
def parallel_bg_filter(item):
|
|
try:
|
|
mask_file = bg.run_img(item['file'], images_dir)
|
|
|
|
# Check and set
|
|
if mask_file is not None and os.path.isfile(mask_file):
|
|
item['p'].set_mask(os.path.basename(mask_file))
|
|
log.ODM_INFO("Wrote %s" % os.path.basename(mask_file))
|
|
else:
|
|
log.ODM_WARNING("Cannot generate mask for %s" % img)
|
|
except Exception as e:
|
|
log.ODM_WARNING("Cannot generate mask for %s: %s" % (img, str(e)))
|
|
|
|
parallel_map(parallel_bg_filter, bg_images, max_workers=args.max_concurrency)
|
|
|
|
log.ODM_INFO("Background masks generation completed!")
|
|
else:
|
|
log.ODM_WARNING("Cannot load AI model (you might need to be connected to the internet?)")
|
|
else:
|
|
log.ODM_INFO("No background masks will be generated (masks already provided)")
|
|
|
|
# End bg removal
|
|
|
|
# Save image database for faster restart
|
|
save_images_database(photos, images_database_file)
|
|
else:
|
|
raise system.ExitException('Not enough supported images in %s' % images_dir)
|
|
else:
|
|
# We have an images database, just load it
|
|
photos = load_images_database(images_database_file)
|
|
|
|
log.ODM_INFO('Found %s usable images' % len(photos))
|
|
log.logger.log_json_images(len(photos))
|
|
|
|
# Create reconstruction object
|
|
reconstruction = types.ODM_Reconstruction(photos)
|
|
|
|
if tree.odm_georeferencing_gcp and not args.use_exif:
|
|
reconstruction.georeference_with_gcp(tree.odm_georeferencing_gcp,
|
|
tree.odm_georeferencing_coords,
|
|
tree.odm_georeferencing_gcp_utm,
|
|
tree.odm_georeferencing_model_txt_geo,
|
|
rerun=self.rerun())
|
|
else:
|
|
reconstruction.georeference_with_gps(tree.dataset_raw,
|
|
tree.odm_georeferencing_coords,
|
|
tree.odm_georeferencing_model_txt_geo,
|
|
rerun=self.rerun())
|
|
|
|
reconstruction.save_proj_srs(os.path.join(tree.odm_georeferencing, tree.odm_georeferencing_proj))
|
|
outputs['reconstruction'] = reconstruction
|
|
|
|
# Try to load boundaries
|
|
if args.boundary:
|
|
if reconstruction.is_georeferenced():
|
|
outputs['boundary'] = boundary.load_boundary(args.boundary, reconstruction.get_proj_srs())
|
|
else:
|
|
args.boundary = None
|
|
log.ODM_WARNING("Reconstruction is not georeferenced, but boundary file provided (will ignore boundary file)")
|
|
|
|
# If sfm-algorithm is triangulation, check if photos have OPK
|
|
if args.sfm_algorithm == 'triangulation':
|
|
for p in photos:
|
|
if not p.has_opk():
|
|
log.ODM_WARNING("No omega/phi/kappa angles found in input photos (%s), switching sfm-algorithm to incremental" % p.filename)
|
|
args.sfm_algorithm = 'incremental'
|
|
break
|
|
|
|
# Rolling shutter cannot be done in non-georeferenced datasets
|
|
if args.rolling_shutter and not reconstruction.is_georeferenced():
|
|
log.ODM_WARNING("Reconstruction is not georeferenced, disabling rolling shutter correction")
|
|
args.rolling_shutter = False
|