Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
iperov
GitHub Repository: iperov/deepfacelab
Path: blob/master/main.py
623 views
1
if __name__ == "__main__":
2
# Fix for linux
3
import multiprocessing
4
multiprocessing.set_start_method("spawn")
5
6
from core.leras import nn
7
nn.initialize_main_env()
8
import os
9
import sys
10
import time
11
import argparse
12
13
from core import pathex
14
from core import osex
15
from pathlib import Path
16
from core.interact import interact as io
17
18
if sys.version_info[0] < 3 or (sys.version_info[0] == 3 and sys.version_info[1] < 6):
19
raise Exception("This program requires at least Python 3.6")
20
21
class fixPathAction(argparse.Action):
22
def __call__(self, parser, namespace, values, option_string=None):
23
setattr(namespace, self.dest, os.path.abspath(os.path.expanduser(values)))
24
25
exit_code = 0
26
27
parser = argparse.ArgumentParser()
28
subparsers = parser.add_subparsers()
29
30
def process_extract(arguments):
31
osex.set_process_lowest_prio()
32
from mainscripts import Extractor
33
Extractor.main( detector = arguments.detector,
34
input_path = Path(arguments.input_dir),
35
output_path = Path(arguments.output_dir),
36
output_debug = arguments.output_debug,
37
manual_fix = arguments.manual_fix,
38
manual_output_debug_fix = arguments.manual_output_debug_fix,
39
manual_window_size = arguments.manual_window_size,
40
face_type = arguments.face_type,
41
max_faces_from_image = arguments.max_faces_from_image,
42
image_size = arguments.image_size,
43
jpeg_quality = arguments.jpeg_quality,
44
cpu_only = arguments.cpu_only,
45
force_gpu_idxs = [ int(x) for x in arguments.force_gpu_idxs.split(',') ] if arguments.force_gpu_idxs is not None else None,
46
)
47
48
p = subparsers.add_parser( "extract", help="Extract the faces from a pictures.")
49
p.add_argument('--detector', dest="detector", choices=['s3fd','manual'], default=None, help="Type of detector.")
50
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.")
51
p.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the extracted files will be stored.")
52
p.add_argument('--output-debug', action="store_true", dest="output_debug", default=None, help="Writes debug images to <output-dir>_debug\ directory.")
53
p.add_argument('--no-output-debug', action="store_false", dest="output_debug", default=None, help="Don't writes debug images to <output-dir>_debug\ directory.")
54
p.add_argument('--face-type', dest="face_type", choices=['half_face', 'full_face', 'whole_face', 'head', 'mark_only'], default=None)
55
p.add_argument('--max-faces-from-image', type=int, dest="max_faces_from_image", default=None, help="Max faces from image.")
56
p.add_argument('--image-size', type=int, dest="image_size", default=None, help="Output image size.")
57
p.add_argument('--jpeg-quality', type=int, dest="jpeg_quality", default=None, help="Jpeg quality.")
58
p.add_argument('--manual-fix', action="store_true", dest="manual_fix", default=False, help="Enables manual extract only frames where faces were not recognized.")
59
p.add_argument('--manual-output-debug-fix', action="store_true", dest="manual_output_debug_fix", default=False, help="Performs manual reextract input-dir frames which were deleted from [output_dir]_debug\ dir.")
60
p.add_argument('--manual-window-size', type=int, dest="manual_window_size", default=1368, help="Manual fix window size. Default: 1368.")
61
p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Extract on CPU..")
62
p.add_argument('--force-gpu-idxs', dest="force_gpu_idxs", default=None, help="Force to choose GPU indexes separated by comma.")
63
64
p.set_defaults (func=process_extract)
65
66
def process_sort(arguments):
67
osex.set_process_lowest_prio()
68
from mainscripts import Sorter
69
Sorter.main (input_path=Path(arguments.input_dir), sort_by_method=arguments.sort_by_method)
70
71
p = subparsers.add_parser( "sort", help="Sort faces in a directory.")
72
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.")
73
p.add_argument('--by', dest="sort_by_method", default=None, choices=("blur", "motion-blur", "face-yaw", "face-pitch", "face-source-rect-size", "hist", "hist-dissim", "brightness", "hue", "black", "origname", "oneface", "final-by-blur", "final-by-size", "absdiff"), help="Method of sorting. 'origname' sort by original filename to recover original sequence." )
74
p.set_defaults (func=process_sort)
75
76
def process_util(arguments):
77
osex.set_process_lowest_prio()
78
from mainscripts import Util
79
80
if arguments.add_landmarks_debug_images:
81
Util.add_landmarks_debug_images (input_path=arguments.input_dir)
82
83
if arguments.recover_original_aligned_filename:
84
Util.recover_original_aligned_filename (input_path=arguments.input_dir)
85
86
if arguments.save_faceset_metadata:
87
Util.save_faceset_metadata_folder (input_path=arguments.input_dir)
88
89
if arguments.restore_faceset_metadata:
90
Util.restore_faceset_metadata_folder (input_path=arguments.input_dir)
91
92
if arguments.pack_faceset:
93
io.log_info ("Performing faceset packing...\r\n")
94
from samplelib import PackedFaceset
95
PackedFaceset.pack( Path(arguments.input_dir) )
96
97
if arguments.unpack_faceset:
98
io.log_info ("Performing faceset unpacking...\r\n")
99
from samplelib import PackedFaceset
100
PackedFaceset.unpack( Path(arguments.input_dir) )
101
102
if arguments.export_faceset_mask:
103
io.log_info ("Exporting faceset mask..\r\n")
104
Util.export_faceset_mask( Path(arguments.input_dir) )
105
106
p = subparsers.add_parser( "util", help="Utilities.")
107
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.")
108
p.add_argument('--add-landmarks-debug-images', action="store_true", dest="add_landmarks_debug_images", default=False, help="Add landmarks debug image for aligned faces.")
109
p.add_argument('--recover-original-aligned-filename', action="store_true", dest="recover_original_aligned_filename", default=False, help="Recover original aligned filename.")
110
p.add_argument('--save-faceset-metadata', action="store_true", dest="save_faceset_metadata", default=False, help="Save faceset metadata to file.")
111
p.add_argument('--restore-faceset-metadata', action="store_true", dest="restore_faceset_metadata", default=False, help="Restore faceset metadata to file. Image filenames must be the same as used with save.")
112
p.add_argument('--pack-faceset', action="store_true", dest="pack_faceset", default=False, help="")
113
p.add_argument('--unpack-faceset', action="store_true", dest="unpack_faceset", default=False, help="")
114
p.add_argument('--export-faceset-mask', action="store_true", dest="export_faceset_mask", default=False, help="")
115
116
p.set_defaults (func=process_util)
117
118
def process_train(arguments):
119
osex.set_process_lowest_prio()
120
121
122
kwargs = {'model_class_name' : arguments.model_name,
123
'saved_models_path' : Path(arguments.model_dir),
124
'training_data_src_path' : Path(arguments.training_data_src_dir),
125
'training_data_dst_path' : Path(arguments.training_data_dst_dir),
126
'pretraining_data_path' : Path(arguments.pretraining_data_dir) if arguments.pretraining_data_dir is not None else None,
127
'pretrained_model_path' : Path(arguments.pretrained_model_dir) if arguments.pretrained_model_dir is not None else None,
128
'no_preview' : arguments.no_preview,
129
'force_model_name' : arguments.force_model_name,
130
'force_gpu_idxs' : [ int(x) for x in arguments.force_gpu_idxs.split(',') ] if arguments.force_gpu_idxs is not None else None,
131
'cpu_only' : arguments.cpu_only,
132
'silent_start' : arguments.silent_start,
133
'execute_programs' : [ [int(x[0]), x[1] ] for x in arguments.execute_program ],
134
'debug' : arguments.debug,
135
}
136
from mainscripts import Trainer
137
Trainer.main(**kwargs)
138
139
p = subparsers.add_parser( "train", help="Trainer")
140
p.add_argument('--training-data-src-dir', required=True, action=fixPathAction, dest="training_data_src_dir", help="Dir of extracted SRC faceset.")
141
p.add_argument('--training-data-dst-dir', required=True, action=fixPathAction, dest="training_data_dst_dir", help="Dir of extracted DST faceset.")
142
p.add_argument('--pretraining-data-dir', action=fixPathAction, dest="pretraining_data_dir", default=None, help="Optional dir of extracted faceset that will be used in pretraining mode.")
143
p.add_argument('--pretrained-model-dir', action=fixPathAction, dest="pretrained_model_dir", default=None, help="Optional dir of pretrain model files. (Currently only for Quick96).")
144
p.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Saved models dir.")
145
p.add_argument('--model', required=True, dest="model_name", choices=pathex.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), help="Model class name.")
146
p.add_argument('--debug', action="store_true", dest="debug", default=False, help="Debug samples.")
147
p.add_argument('--no-preview', action="store_true", dest="no_preview", default=False, help="Disable preview window.")
148
p.add_argument('--force-model-name', dest="force_model_name", default=None, help="Forcing to choose model name from model/ folder.")
149
p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Train on CPU.")
150
p.add_argument('--force-gpu-idxs', dest="force_gpu_idxs", default=None, help="Force to choose GPU indexes separated by comma.")
151
p.add_argument('--silent-start', action="store_true", dest="silent_start", default=False, help="Silent start. Automatically chooses Best GPU and last used model.")
152
153
p.add_argument('--execute-program', dest="execute_program", default=[], action='append', nargs='+')
154
p.set_defaults (func=process_train)
155
156
def process_exportdfm(arguments):
157
osex.set_process_lowest_prio()
158
from mainscripts import ExportDFM
159
ExportDFM.main(model_class_name = arguments.model_name, saved_models_path = Path(arguments.model_dir))
160
161
p = subparsers.add_parser( "exportdfm", help="Export model to use in DeepFaceLive.")
162
p.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Saved models dir.")
163
p.add_argument('--model', required=True, dest="model_name", choices=pathex.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), help="Model class name.")
164
p.set_defaults (func=process_exportdfm)
165
166
def process_merge(arguments):
167
osex.set_process_lowest_prio()
168
from mainscripts import Merger
169
Merger.main ( model_class_name = arguments.model_name,
170
saved_models_path = Path(arguments.model_dir),
171
force_model_name = arguments.force_model_name,
172
input_path = Path(arguments.input_dir),
173
output_path = Path(arguments.output_dir),
174
output_mask_path = Path(arguments.output_mask_dir),
175
aligned_path = Path(arguments.aligned_dir) if arguments.aligned_dir is not None else None,
176
force_gpu_idxs = arguments.force_gpu_idxs,
177
cpu_only = arguments.cpu_only)
178
179
p = subparsers.add_parser( "merge", help="Merger")
180
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.")
181
p.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the merged files will be stored.")
182
p.add_argument('--output-mask-dir', required=True, action=fixPathAction, dest="output_mask_dir", help="Output mask directory. This is where the mask files will be stored.")
183
p.add_argument('--aligned-dir', action=fixPathAction, dest="aligned_dir", default=None, help="Aligned directory. This is where the extracted of dst faces stored.")
184
p.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Model dir.")
185
p.add_argument('--model', required=True, dest="model_name", choices=pathex.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), help="Model class name.")
186
p.add_argument('--force-model-name', dest="force_model_name", default=None, help="Forcing to choose model name from model/ folder.")
187
p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Merge on CPU.")
188
p.add_argument('--force-gpu-idxs', dest="force_gpu_idxs", default=None, help="Force to choose GPU indexes separated by comma.")
189
p.set_defaults(func=process_merge)
190
191
videoed_parser = subparsers.add_parser( "videoed", help="Video processing.").add_subparsers()
192
193
def process_videoed_extract_video(arguments):
194
osex.set_process_lowest_prio()
195
from mainscripts import VideoEd
196
VideoEd.extract_video (arguments.input_file, arguments.output_dir, arguments.output_ext, arguments.fps)
197
p = videoed_parser.add_parser( "extract-video", help="Extract images from video file.")
198
p.add_argument('--input-file', required=True, action=fixPathAction, dest="input_file", help="Input file to be processed. Specify .*-extension to find first file.")
199
p.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the extracted images will be stored.")
200
p.add_argument('--output-ext', dest="output_ext", default=None, help="Image format (extension) of output files.")
201
p.add_argument('--fps', type=int, dest="fps", default=None, help="How many frames of every second of the video will be extracted. 0 - full fps.")
202
p.set_defaults(func=process_videoed_extract_video)
203
204
def process_videoed_cut_video(arguments):
205
osex.set_process_lowest_prio()
206
from mainscripts import VideoEd
207
VideoEd.cut_video (arguments.input_file,
208
arguments.from_time,
209
arguments.to_time,
210
arguments.audio_track_id,
211
arguments.bitrate)
212
p = videoed_parser.add_parser( "cut-video", help="Cut video file.")
213
p.add_argument('--input-file', required=True, action=fixPathAction, dest="input_file", help="Input file to be processed. Specify .*-extension to find first file.")
214
p.add_argument('--from-time', dest="from_time", default=None, help="From time, for example 00:00:00.000")
215
p.add_argument('--to-time', dest="to_time", default=None, help="To time, for example 00:00:00.000")
216
p.add_argument('--audio-track-id', type=int, dest="audio_track_id", default=None, help="Specify audio track id.")
217
p.add_argument('--bitrate', type=int, dest="bitrate", default=None, help="Bitrate of output file in Megabits.")
218
p.set_defaults(func=process_videoed_cut_video)
219
220
def process_videoed_denoise_image_sequence(arguments):
221
osex.set_process_lowest_prio()
222
from mainscripts import VideoEd
223
VideoEd.denoise_image_sequence (arguments.input_dir, arguments.factor)
224
p = videoed_parser.add_parser( "denoise-image-sequence", help="Denoise sequence of images, keeping sharp edges. Helps to remove pixel shake from the predicted face.")
225
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory to be processed.")
226
p.add_argument('--factor', type=int, dest="factor", default=None, help="Denoise factor (1-20).")
227
p.set_defaults(func=process_videoed_denoise_image_sequence)
228
229
def process_videoed_video_from_sequence(arguments):
230
osex.set_process_lowest_prio()
231
from mainscripts import VideoEd
232
VideoEd.video_from_sequence (input_dir = arguments.input_dir,
233
output_file = arguments.output_file,
234
reference_file = arguments.reference_file,
235
ext = arguments.ext,
236
fps = arguments.fps,
237
bitrate = arguments.bitrate,
238
include_audio = arguments.include_audio,
239
lossless = arguments.lossless)
240
241
p = videoed_parser.add_parser( "video-from-sequence", help="Make video from image sequence.")
242
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input file to be processed. Specify .*-extension to find first file.")
243
p.add_argument('--output-file', required=True, action=fixPathAction, dest="output_file", help="Input file to be processed. Specify .*-extension to find first file.")
244
p.add_argument('--reference-file', action=fixPathAction, dest="reference_file", help="Reference file used to determine proper FPS and transfer audio from it. Specify .*-extension to find first file.")
245
p.add_argument('--ext', dest="ext", default='png', help="Image format (extension) of input files.")
246
p.add_argument('--fps', type=int, dest="fps", default=None, help="FPS of output file. Overwritten by reference-file.")
247
p.add_argument('--bitrate', type=int, dest="bitrate", default=None, help="Bitrate of output file in Megabits.")
248
p.add_argument('--include-audio', action="store_true", dest="include_audio", default=False, help="Include audio from reference file.")
249
p.add_argument('--lossless', action="store_true", dest="lossless", default=False, help="PNG codec.")
250
251
p.set_defaults(func=process_videoed_video_from_sequence)
252
253
facesettool_parser = subparsers.add_parser( "facesettool", help="Faceset tools.").add_subparsers()
254
255
def process_faceset_enhancer(arguments):
256
osex.set_process_lowest_prio()
257
from mainscripts import FacesetEnhancer
258
FacesetEnhancer.process_folder ( Path(arguments.input_dir),
259
cpu_only=arguments.cpu_only,
260
force_gpu_idxs=arguments.force_gpu_idxs
261
)
262
263
p = facesettool_parser.add_parser ("enhance", help="Enhance details in DFL faceset.")
264
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory of aligned faces.")
265
p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Process on CPU.")
266
p.add_argument('--force-gpu-idxs', dest="force_gpu_idxs", default=None, help="Force to choose GPU indexes separated by comma.")
267
268
p.set_defaults(func=process_faceset_enhancer)
269
270
271
p = facesettool_parser.add_parser ("resize", help="Resize DFL faceset.")
272
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory of aligned faces.")
273
274
def process_faceset_resizer(arguments):
275
osex.set_process_lowest_prio()
276
from mainscripts import FacesetResizer
277
FacesetResizer.process_folder ( Path(arguments.input_dir) )
278
p.set_defaults(func=process_faceset_resizer)
279
280
def process_dev_test(arguments):
281
osex.set_process_lowest_prio()
282
from mainscripts import dev_misc
283
dev_misc.dev_gen_mask_files( arguments.input_dir )
284
285
p = subparsers.add_parser( "dev_test", help="")
286
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
287
p.set_defaults (func=process_dev_test)
288
289
# ========== XSeg
290
xseg_parser = subparsers.add_parser( "xseg", help="XSeg tools.").add_subparsers()
291
292
p = xseg_parser.add_parser( "editor", help="XSeg editor.")
293
294
def process_xsegeditor(arguments):
295
osex.set_process_lowest_prio()
296
from XSegEditor import XSegEditor
297
global exit_code
298
exit_code = XSegEditor.start (Path(arguments.input_dir))
299
300
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
301
302
p.set_defaults (func=process_xsegeditor)
303
304
p = xseg_parser.add_parser( "apply", help="Apply trained XSeg model to the extracted faces.")
305
306
def process_xsegapply(arguments):
307
osex.set_process_lowest_prio()
308
from mainscripts import XSegUtil
309
XSegUtil.apply_xseg (Path(arguments.input_dir), Path(arguments.model_dir))
310
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
311
p.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir")
312
p.set_defaults (func=process_xsegapply)
313
314
315
p = xseg_parser.add_parser( "remove", help="Remove applied XSeg masks from the extracted faces.")
316
def process_xsegremove(arguments):
317
osex.set_process_lowest_prio()
318
from mainscripts import XSegUtil
319
XSegUtil.remove_xseg (Path(arguments.input_dir) )
320
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
321
p.set_defaults (func=process_xsegremove)
322
323
324
p = xseg_parser.add_parser( "remove_labels", help="Remove XSeg labels from the extracted faces.")
325
def process_xsegremovelabels(arguments):
326
osex.set_process_lowest_prio()
327
from mainscripts import XSegUtil
328
XSegUtil.remove_xseg_labels (Path(arguments.input_dir) )
329
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
330
p.set_defaults (func=process_xsegremovelabels)
331
332
333
p = xseg_parser.add_parser( "fetch", help="Copies faces containing XSeg polygons in <input_dir>_xseg dir.")
334
335
def process_xsegfetch(arguments):
336
osex.set_process_lowest_prio()
337
from mainscripts import XSegUtil
338
XSegUtil.fetch_xseg (Path(arguments.input_dir) )
339
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir")
340
p.set_defaults (func=process_xsegfetch)
341
342
def bad_args(arguments):
343
parser.print_help()
344
exit(0)
345
parser.set_defaults(func=bad_args)
346
347
arguments = parser.parse_args()
348
arguments.func(arguments)
349
350
if exit_code == 0:
351
print ("Done.")
352
353
exit(exit_code)
354
355
'''
356
import code
357
code.interact(local=dict(globals(), **locals()))
358
'''
359
360