kopia lustrzana https://github.com/OpenDroneMap/ODM
Extract video frames in load dataset
rodzic
64b687a3a6
commit
67787f0059
|
@ -658,7 +658,7 @@ def config(argv=None, parser=None):
|
||||||
action=StoreValue,
|
action=StoreValue,
|
||||||
default=500,
|
default=500,
|
||||||
metavar='<positive integer>',
|
metavar='<positive integer>',
|
||||||
help='Maximum number of frames to extract from video files for processing. Default: %(default)s')
|
help='Maximum number of frames to extract from video files for processing. Set to 0 for no limit. Default: %(default)s')
|
||||||
|
|
||||||
parser.add_argument('--video-resolution',
|
parser.add_argument('--video-resolution',
|
||||||
type=int,
|
type=int,
|
||||||
|
|
|
@ -41,6 +41,7 @@ settings_path = os.path.join(root_path, 'settings.yaml')
|
||||||
|
|
||||||
# Define supported image extensions
|
# Define supported image extensions
|
||||||
supported_extensions = {'.jpg','.jpeg','.png', '.tif', '.tiff', '.bmp'}
|
supported_extensions = {'.jpg','.jpeg','.png', '.tif', '.tiff', '.bmp'}
|
||||||
|
supported_video_extensions = {'.mp4', '.mov'}
|
||||||
|
|
||||||
# Define the number of cores
|
# Define the number of cores
|
||||||
num_cores = multiprocessing.cpu_count()
|
num_cores = multiprocessing.cpu_count()
|
||||||
|
|
|
@ -135,7 +135,7 @@ class Video2Dataset:
|
||||||
if self.f is not None:
|
if self.f is not None:
|
||||||
self.f.close()
|
self.f.close()
|
||||||
|
|
||||||
if self.parameters.limit is not None and self.global_idx >= self.parameters.limit:
|
if self.parameters.limit is not None and self.parameters.limit > 0 and self.global_idx >= self.parameters.limit:
|
||||||
log.ODM_INFO("Limit of {} frames reached, trimming dataset".format(self.parameters.limit))
|
log.ODM_INFO("Limit of {} frames reached, trimming dataset".format(self.parameters.limit))
|
||||||
output_file_paths = limit_files(output_file_paths, self.parameters.limit)
|
output_file_paths = limit_files(output_file_paths, self.parameters.limit)
|
||||||
|
|
||||||
|
|
|
@ -15,6 +15,7 @@ from opendm import ai
|
||||||
from opendm.skyremoval.skyfilter import SkyFilter
|
from opendm.skyremoval.skyfilter import SkyFilter
|
||||||
from opendm.bgfilter import BgFilter
|
from opendm.bgfilter import BgFilter
|
||||||
from opendm.concurrency import parallel_map
|
from opendm.concurrency import parallel_map
|
||||||
|
from opendm.video.video2dataset import Parameters, Video2Dataset
|
||||||
|
|
||||||
def save_images_database(photos, database_file):
|
def save_images_database(photos, database_file):
|
||||||
with open(database_file, 'w') as f:
|
with open(database_file, 'w') as f:
|
||||||
|
@ -58,22 +59,25 @@ class ODMLoadDatasetStage(types.ODM_Stage):
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
log.ODM_WARNING("Cannot write benchmark file: %s" % str(e))
|
log.ODM_WARNING("Cannot write benchmark file: %s" % str(e))
|
||||||
|
|
||||||
# check if the image filename is supported
|
def valid_filename(filename, supported_extensions):
|
||||||
def valid_image_filename(filename):
|
|
||||||
(pathfn, ext) = os.path.splitext(filename)
|
(pathfn, ext) = os.path.splitext(filename)
|
||||||
return ext.lower() in context.supported_extensions and pathfn[-5:] != "_mask"
|
return ext.lower() in supported_extensions and pathfn[-5:] != "_mask"
|
||||||
|
|
||||||
# Get supported images from dir
|
# Get supported images from dir
|
||||||
def get_images(in_dir):
|
def get_images(in_dir):
|
||||||
entries = os.listdir(in_dir)
|
entries = os.listdir(in_dir)
|
||||||
valid, rejects = [], []
|
valid, rejects = [], []
|
||||||
for f in entries:
|
for f in entries:
|
||||||
if valid_image_filename(f):
|
if valid_filename(f, context.supported_extensions):
|
||||||
valid.append(f)
|
valid.append(f)
|
||||||
else:
|
else:
|
||||||
rejects.append(f)
|
rejects.append(f)
|
||||||
return valid, rejects
|
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):
|
def find_mask(photo_path, masks):
|
||||||
(pathfn, ext) = os.path.splitext(os.path.basename(photo_path))
|
(pathfn, ext) = os.path.splitext(os.path.basename(photo_path))
|
||||||
k = "{}_mask".format(pathfn)
|
k = "{}_mask".format(pathfn)
|
||||||
|
@ -86,6 +90,8 @@ class ODMLoadDatasetStage(types.ODM_Stage):
|
||||||
else:
|
else:
|
||||||
log.ODM_WARNING("Image mask {} has a space. Spaces are currently not supported for image masks.".format(mask))
|
log.ODM_WARNING("Image mask {} has a space. Spaces are currently not supported for image masks.".format(mask))
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
# get images directory
|
# get images directory
|
||||||
images_dir = tree.dataset_raw
|
images_dir = tree.dataset_raw
|
||||||
|
|
||||||
|
@ -100,6 +106,51 @@ class ODMLoadDatasetStage(types.ODM_Stage):
|
||||||
if not os.path.exists(images_dir):
|
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))
|
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": 0.2,
|
||||||
|
"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)
|
files, rejects = get_images(images_dir)
|
||||||
if files:
|
if files:
|
||||||
# create ODMPhoto list
|
# create ODMPhoto list
|
||||||
|
|
Ładowanie…
Reference in New Issue