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,
|
||||
default=500,
|
||||
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',
|
||||
type=int,
|
||||
|
|
|
@ -41,6 +41,7 @@ settings_path = os.path.join(root_path, 'settings.yaml')
|
|||
|
||||
# Define supported image extensions
|
||||
supported_extensions = {'.jpg','.jpeg','.png', '.tif', '.tiff', '.bmp'}
|
||||
supported_video_extensions = {'.mp4', '.mov'}
|
||||
|
||||
# Define the number of cores
|
||||
num_cores = multiprocessing.cpu_count()
|
||||
|
|
|
@ -135,7 +135,7 @@ class Video2Dataset:
|
|||
if self.f is not None:
|
||||
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))
|
||||
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.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:
|
||||
|
@ -58,22 +59,25 @@ class ODMLoadDatasetStage(types.ODM_Stage):
|
|||
except Exception as e:
|
||||
log.ODM_WARNING("Cannot write benchmark file: %s" % str(e))
|
||||
|
||||
# check if the image filename is supported
|
||||
def valid_image_filename(filename):
|
||||
def valid_filename(filename, supported_extensions):
|
||||
(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
|
||||
def get_images(in_dir):
|
||||
entries = os.listdir(in_dir)
|
||||
valid, rejects = [], []
|
||||
for f in entries:
|
||||
if valid_image_filename(f):
|
||||
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)
|
||||
|
@ -86,6 +90,8 @@ class ODMLoadDatasetStage(types.ODM_Stage):
|
|||
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
|
||||
|
||||
|
@ -100,6 +106,51 @@ class ODMLoadDatasetStage(types.ODM_Stage):
|
|||
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": 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)
|
||||
if files:
|
||||
# create ODMPhoto list
|
||||
|
|
Ładowanie…
Reference in New Issue