-
Notifications
You must be signed in to change notification settings - Fork 4
/
Copy pathsplit_scenes.py
188 lines (168 loc) · 8.45 KB
/
split_scenes.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
"""Split Scenes Feature Core Code"""
import os
import glob
import argparse
import shutil
from typing import Callable
from webui_utils.simple_log import SimpleLog
from webui_utils.file_utils import create_directory, is_safe_path, split_filepath
from webui_utils.video_utils import get_detected_scenes, get_detected_breaks, scene_list_to_ranges
from webui_utils.mtqdm import Mtqdm
from resequence_files import ResequenceFiles
def main():
"""Use the Split Scenes feature from the command line"""
parser = argparse.ArgumentParser(description='Split a directory of PNG frame files')
parser.add_argument("--input_path", default=None, type=str,
help="Input path to PNG frame files to split")
parser.add_argument("--output_path", default=None, type=str,
help="Base path for frame group directories")
parser.add_argument("--file_ext", default="png", type=str,
help="File extension, default: 'png'; any extension or or '*'")
parser.add_argument("--type", default="scene", type=str,
help="Scene detect type 'scene' (default), 'break'")
parser.add_argument("--scene_threshold", default=0.6, type=float,
help="Threshold between 0.0 and 1.0 for scene detection (default 0.6)")
parser.add_argument("--break_duration", default=2.0, type=float,
help="Duration in seconds for break to be detectable (default 2.0)")
parser.add_argument("--break_ratio", default=0.98, type=float,
help="Percent 0.0 to 1.0 of frame that must be black to be detectable (default 0.98)")
parser.add_argument("--verbose", dest="verbose", default=False, action="store_true",
help="Show extra details")
args = parser.parse_args()
log = SimpleLog(args.verbose)
SplitScenes(args.input_path,
args.output_path,
args.file_ext,
args.type,
args.scene_threshold,
args.break_duration,
args.break_ratio,
log.log).split()
class SplitScenes:
"""Encapsulate logic for Split Scenes feature"""
def __init__(self,
input_path : str,
output_path : str,
file_ext : str,
type : str,
scene_threshold : float,
break_duration : float,
break_ratio : float,
log_fn : Callable | None):
self.input_path = input_path
self.output_path = output_path
self.file_ext = file_ext
self.type = type
self.scene_threshold = scene_threshold
self.break_duration = break_duration
self.break_ratio = break_ratio
self.dry_run = False
self.log_fn = log_fn
valid_types = ["scene", "break"]
if not is_safe_path(self.input_path):
raise ValueError("'input_path' must be a legal path")
if not is_safe_path(self.output_path):
raise ValueError("'output_path' must be a legal path")
if not self.type in valid_types:
raise ValueError(f"'type' must be one of {', '.join([t for t in valid_types])}")
def split_scenes(self, format : str="png", move_files : bool=False):
files = sorted(glob.glob(os.path.join(self.input_path, f"*.{self.file_ext}")))
num_files = len(files)
num_width = len(str(num_files))
self.log(f"calling `get_detected_scenes` with input path '{self.input_path}'" +\
f" threshold '{self.scene_threshold}'")
scenes = get_detected_scenes(self.input_path, float(self.scene_threshold), type=format)
# add one more final fake detection past the end to include frames past the last detection
scenes.append(num_files+1)
ranges = scene_list_to_ranges(scenes, num_files)
group_paths = []
with Mtqdm().open_bar(total=len(ranges), desc="Scenes") as scene_bar:
for _range in ranges:
first_index = _range["first_frame"]
last_index = _range["last_frame"]
if last_index >= num_files:
last_index = num_files
group_size = _range["scene_size"]
group_name = f"{str(first_index).zfill(num_width)}" +\
f"-{str(last_index).zfill(num_width)}"
group_path = os.path.join(self.output_path, group_name)
group_paths.append(group_path)
if self.dry_run:
self.log(f"[Dry Run] Creating directory {group_path}")
else:
create_directory(group_path)
group_files = [files[index] for index in range(first_index, last_index+1)]
base_filename = f"{group_name}-{self.type}-split-frame"
ResequenceFiles(self.input_path,
self.file_ext,
base_filename,
0, 1, 1, 0,
num_width,
False,
self.log,
output_path=group_path).resequence(move_files=move_files,
file_list=group_files)
Mtqdm().update_bar(scene_bar)
return group_paths
def split_breaks(self, format : str="jpg", move_files : bool=False):
files = sorted(glob.glob(os.path.join(self.input_path, f"*.{self.file_ext}")))
num_files = len(files)
num_width = len(str(num_files))
self.log(f"calling `get_detected_breaks` with input path '{self.input_path}'" +\
f" duration '{self.break_duration}' ratio '{self.break_ratio}'")
scenes = get_detected_breaks(self.input_path, float(self.break_duration),
float(self.break_ratio), type=format)
# add one more final fake detection past the end to include frames past the last detection
scenes.append(num_files+1)
ranges = scene_list_to_ranges(scenes, num_files)
group_paths = []
with Mtqdm().open_bar(total=len(ranges), desc="Scenes") as scene_bar:
for _range in ranges:
first_index = _range["first_frame"]
last_index = _range["last_frame"]
if last_index >= num_files:
last_index = num_files
group_size = _range["scene_size"]
group_name = f"{str(first_index).zfill(num_width)}" +\
f"-{str(last_index).zfill(num_width)}"
group_path = os.path.join(self.output_path, group_name)
group_paths.append(group_path)
if self.dry_run:
self.log(f"[Dry Run] Creating directory {group_path}")
else:
create_directory(group_path)
group_files = [files[index] for index in range(first_index, last_index+1)]
base_filename = f"{group_name}-{self.type}-split-frame"
ResequenceFiles(self.input_path,
self.file_ext,
base_filename,
0, 1, 1, 0,
num_width,
False,
self.log,
output_path=group_path).resequence(move_files=move_files,
file_list=group_files)
Mtqdm().update_bar(scene_bar)
return group_paths
def split(self, type : str="png", move_files : bool=False) -> list:
"""Invoke the Split Scenes feature"""
if self.type == "scene":
if self.scene_threshold < 0.0 or self.scene_threshold > 1.0:
raise ValueError("'scene_threshold' must be between 0.0 and 1.0")
self.split_scenes(type, move_files)
else:
if self.break_duration < 0.0:
raise ValueError("'break_duration' >= 0.0")
if self.break_ratio < 0.0 or self.break_ratio > 1.0:
raise ValueError("'break_ratio' must be between 0.0 and 1.0")
self.split_breaks(type, move_files)
if self.dry_run:
print(f"[Dry Run] Creating base output path {self.output_path}")
else:
create_directory(self.output_path)
def log(self, message : str) -> None:
"""Logging"""
if self.log_fn:
self.log_fn(message)
if __name__ == '__main__':
main()