cosmicweb98 commited on
Commit
5de33df
·
verified ·
1 Parent(s): da28985

Upload 30 files

Browse files
roop/FaceSet.py ADDED
@@ -0,0 +1,20 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+
3
+ class FaceSet:
4
+ faces = []
5
+ ref_images = []
6
+ embedding_average = 'None'
7
+ embeddings_backup = None
8
+
9
+ def __init__(self):
10
+ self.faces = []
11
+ self.ref_images = []
12
+ self.embeddings_backup = None
13
+
14
+ def AverageEmbeddings(self):
15
+ if len(self.faces) > 1 and self.embeddings_backup is None:
16
+ self.embeddings_backup = self.faces[0]['embedding']
17
+ embeddings = [face.embedding for face in self.faces]
18
+
19
+ self.faces[0]['embedding'] = np.mean(embeddings, axis=0)
20
+ # try median too?
roop/ProcessEntry.py ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ class ProcessEntry:
2
+ def __init__(self, filename: str, start: int, end: int, fps: float):
3
+ self.filename = filename
4
+ self.finalname = None
5
+ self.startframe = start
6
+ self.endframe = end
7
+ self.fps = fps
roop/ProcessMgr.py ADDED
@@ -0,0 +1,750 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import cv2
3
+ import numpy as np
4
+ import psutil
5
+
6
+ from enum import Enum
7
+ from roop.ProcessOptions import ProcessOptions
8
+
9
+ from roop.face_util import get_first_face, get_all_faces, rotate_anticlockwise, rotate_clockwise, clamp_cut_values
10
+ from roop.utilities import compute_cosine_distance, get_device, str_to_class
11
+ import roop.vr_util as vr
12
+
13
+ from typing import Any, List, Callable
14
+ from roop.typing import Frame, Face
15
+ from concurrent.futures import ThreadPoolExecutor, as_completed
16
+ from threading import Thread, Lock
17
+ from queue import Queue
18
+ from tqdm import tqdm
19
+ from roop.ffmpeg_writer import FFMPEG_VideoWriter
20
+ import roop.globals
21
+
22
+
23
+
24
+ # Poor man's enum to be able to compare to int
25
+ class eNoFaceAction():
26
+ USE_ORIGINAL_FRAME = 0
27
+ RETRY_ROTATED = 1
28
+ SKIP_FRAME = 2
29
+ SKIP_FRAME_IF_DISSIMILAR = 3
30
+
31
+
32
+
33
+ def create_queue(temp_frame_paths: List[str]) -> Queue[str]:
34
+ queue: Queue[str] = Queue()
35
+ for frame_path in temp_frame_paths:
36
+ queue.put(frame_path)
37
+ return queue
38
+
39
+
40
+ def pick_queue(queue: Queue[str], queue_per_future: int) -> List[str]:
41
+ queues = []
42
+ for _ in range(queue_per_future):
43
+ if not queue.empty():
44
+ queues.append(queue.get())
45
+ return queues
46
+
47
+
48
+
49
+ class ProcessMgr():
50
+ input_face_datas = []
51
+ target_face_datas = []
52
+
53
+ imagemask = None
54
+
55
+ processors = []
56
+ options : ProcessOptions = None
57
+
58
+ num_threads = 1
59
+ current_index = 0
60
+ processing_threads = 1
61
+ buffer_wait_time = 0.1
62
+
63
+ lock = Lock()
64
+
65
+ frames_queue = None
66
+ processed_queue = None
67
+
68
+ videowriter= None
69
+
70
+ progress_gradio = None
71
+ total_frames = 0
72
+
73
+
74
+
75
+
76
+ plugins = {
77
+ 'faceswap' : 'FaceSwapInsightFace',
78
+ 'mask_clip2seg' : 'Mask_Clip2Seg',
79
+ 'mask_xseg' : 'Mask_XSeg',
80
+ 'codeformer' : 'Enhance_CodeFormer',
81
+ 'gfpgan' : 'Enhance_GFPGAN',
82
+ 'dmdnet' : 'Enhance_DMDNet',
83
+ 'gpen' : 'Enhance_GPEN',
84
+ 'restoreformer++' : 'Enhance_RestoreFormerPPlus',
85
+ 'colorizer' : 'Frame_Colorizer',
86
+ 'filter_generic' : 'Frame_Filter',
87
+ 'removebg' : 'Frame_Masking',
88
+ 'upscale' : 'Frame_Upscale'
89
+ }
90
+
91
+ def __init__(self, progress):
92
+ if progress is not None:
93
+ self.progress_gradio = progress
94
+
95
+ def reuseOldProcessor(self, name:str):
96
+ for p in self.processors:
97
+ if p.processorname == name:
98
+ return p
99
+
100
+ return None
101
+
102
+
103
+ def initialize(self, input_faces, target_faces, options):
104
+ self.input_face_datas = input_faces
105
+ self.target_face_datas = target_faces
106
+ self.options = options
107
+ devicename = get_device()
108
+
109
+ roop.globals.g_desired_face_analysis=["landmark_3d_68", "landmark_2d_106","detection","recognition"]
110
+ if options.swap_mode == "all_female" or options.swap_mode == "all_male":
111
+ roop.globals.g_desired_face_analysis.append("genderage")
112
+
113
+ for p in self.processors:
114
+ newp = next((x for x in options.processors.keys() if x == p.processorname), None)
115
+ if newp is None:
116
+ p.Release()
117
+ del p
118
+
119
+ newprocessors = []
120
+ for key, extoption in options.processors.items():
121
+ p = self.reuseOldProcessor(key)
122
+ if p is None:
123
+ classname = self.plugins[key]
124
+ module = 'roop.processors.' + classname
125
+ p = str_to_class(module, classname)
126
+ if p is not None:
127
+ extoption.update({"devicename": devicename})
128
+ p.Initialize(extoption)
129
+ newprocessors.append(p)
130
+ else:
131
+ print(f"Not using {module}")
132
+ self.processors = newprocessors
133
+
134
+
135
+
136
+ if isinstance(self.options.imagemask, dict) and self.options.imagemask.get("layers") and len(self.options.imagemask["layers"]) > 0:
137
+ self.options.imagemask = self.options.imagemask.get("layers")[0]
138
+ # Get rid of alpha
139
+ self.options.imagemask = cv2.cvtColor(self.options.imagemask, cv2.COLOR_RGBA2GRAY)
140
+ if np.any(self.options.imagemask):
141
+ mo = self.input_face_datas[0].faces[0].mask_offsets
142
+ self.options.imagemask = self.blur_area(self.options.imagemask, mo[4], mo[5])
143
+ self.options.imagemask = self.options.imagemask.astype(np.float32) / 255
144
+ self.options.imagemask = cv2.cvtColor(self.options.imagemask, cv2.COLOR_GRAY2RGB)
145
+ else:
146
+ self.options.imagemask = None
147
+
148
+ self.options.frame_processing = False
149
+ for p in self.processors:
150
+ if p.type.startswith("frame_"):
151
+ self.options.frame_processing = True
152
+
153
+
154
+
155
+
156
+
157
+
158
+ def run_batch(self, source_files, target_files, threads:int = 1):
159
+ progress_bar_format = '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]'
160
+ self.total_frames = len(source_files)
161
+ self.num_threads = threads
162
+ with tqdm(total=self.total_frames, desc='Processing', unit='frame', dynamic_ncols=True, bar_format=progress_bar_format) as progress:
163
+ with ThreadPoolExecutor(max_workers=threads) as executor:
164
+ futures = []
165
+ queue = create_queue(source_files)
166
+ queue_per_future = max(len(source_files) // threads, 1)
167
+ while not queue.empty():
168
+ future = executor.submit(self.process_frames, source_files, target_files, pick_queue(queue, queue_per_future), lambda: self.update_progress(progress))
169
+ futures.append(future)
170
+ for future in as_completed(futures):
171
+ future.result()
172
+
173
+
174
+ def process_frames(self, source_files: List[str], target_files: List[str], current_files, update: Callable[[], None]) -> None:
175
+ for f in current_files:
176
+ if not roop.globals.processing:
177
+ return
178
+
179
+ # Decode the byte array into an OpenCV image
180
+ temp_frame = cv2.imdecode(np.fromfile(f, dtype=np.uint8), cv2.IMREAD_COLOR)
181
+ if temp_frame is not None:
182
+ if self.options.frame_processing:
183
+ for p in self.processors:
184
+ frame = p.Run(temp_frame)
185
+ resimg = frame
186
+ else:
187
+ resimg = self.process_frame(temp_frame)
188
+ if resimg is not None:
189
+ i = source_files.index(f)
190
+ cv2.imwrite(target_files[i], resimg)
191
+ if update:
192
+ update()
193
+
194
+
195
+
196
+ def read_frames_thread(self, cap, frame_start, frame_end, num_threads):
197
+ num_frame = 0
198
+ total_num = frame_end - frame_start
199
+ if frame_start > 0:
200
+ cap.set(cv2.CAP_PROP_POS_FRAMES,frame_start)
201
+
202
+ while True and roop.globals.processing:
203
+ ret, frame = cap.read()
204
+ if not ret:
205
+ break
206
+
207
+ self.frames_queue[num_frame % num_threads].put(frame, block=True)
208
+ num_frame += 1
209
+ if num_frame == total_num:
210
+ break
211
+
212
+ for i in range(num_threads):
213
+ self.frames_queue[i].put(None)
214
+
215
+
216
+
217
+ def process_videoframes(self, threadindex, progress) -> None:
218
+ while True:
219
+ frame = self.frames_queue[threadindex].get()
220
+ if frame is None:
221
+ self.processing_threads -= 1
222
+ self.processed_queue[threadindex].put((False, None))
223
+ return
224
+ else:
225
+ if self.options.frame_processing:
226
+ for p in self.processors:
227
+ frame = p.Run(frame)
228
+ resimg = frame
229
+ else:
230
+ resimg = self.process_frame(frame)
231
+ self.processed_queue[threadindex].put((True, resimg))
232
+ del frame
233
+ progress()
234
+
235
+
236
+ def write_frames_thread(self):
237
+ nextindex = 0
238
+ num_producers = self.num_threads
239
+
240
+ while True:
241
+ process, frame = self.processed_queue[nextindex % self.num_threads].get()
242
+ nextindex += 1
243
+ if frame is not None:
244
+ self.videowriter.write_frame(frame)
245
+ del frame
246
+ elif process == False:
247
+ num_producers -= 1
248
+ if num_producers < 1:
249
+ return
250
+
251
+
252
+
253
+ def run_batch_inmem(self, source_video, target_video, frame_start, frame_end, fps, threads:int = 1, skip_audio=False):
254
+ cap = cv2.VideoCapture(source_video)
255
+ # frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
256
+ frame_count = (frame_end - frame_start) + 1
257
+ width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
258
+ height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
259
+
260
+ processed_resolution = None
261
+ for p in self.processors:
262
+ if hasattr(p, 'getProcessedResolution'):
263
+ processed_resolution = p.getProcessedResolution(width, height)
264
+ print(f"Processed resolution: {processed_resolution}")
265
+ if processed_resolution is not None:
266
+ width = processed_resolution[0]
267
+ height = processed_resolution[1]
268
+
269
+
270
+ self.total_frames = frame_count
271
+ self.num_threads = threads
272
+
273
+ self.processing_threads = self.num_threads
274
+ self.frames_queue = []
275
+ self.processed_queue = []
276
+ for _ in range(threads):
277
+ self.frames_queue.append(Queue(1))
278
+ self.processed_queue.append(Queue(1))
279
+
280
+ self.videowriter = FFMPEG_VideoWriter(target_video, (width, height), fps, codec=roop.globals.video_encoder, crf=roop.globals.video_quality, audiofile=None)
281
+
282
+ readthread = Thread(target=self.read_frames_thread, args=(cap, frame_start, frame_end, threads))
283
+ readthread.start()
284
+
285
+ writethread = Thread(target=self.write_frames_thread)
286
+ writethread.start()
287
+
288
+ progress_bar_format = '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]'
289
+ with tqdm(total=self.total_frames, desc='Processing', unit='frames', dynamic_ncols=True, bar_format=progress_bar_format) as progress:
290
+ with ThreadPoolExecutor(thread_name_prefix='swap_proc', max_workers=self.num_threads) as executor:
291
+ futures = []
292
+
293
+ for threadindex in range(threads):
294
+ future = executor.submit(self.process_videoframes, threadindex, lambda: self.update_progress(progress))
295
+ futures.append(future)
296
+
297
+ for future in as_completed(futures):
298
+ future.result()
299
+ # wait for the task to complete
300
+ readthread.join()
301
+ writethread.join()
302
+ cap.release()
303
+ self.videowriter.close()
304
+ self.frames_queue.clear()
305
+ self.processed_queue.clear()
306
+
307
+
308
+
309
+
310
+ def update_progress(self, progress: Any = None) -> None:
311
+ process = psutil.Process(os.getpid())
312
+ memory_usage = process.memory_info().rss / 1024 / 1024 / 1024
313
+ progress.set_postfix({
314
+ 'memory_usage': '{:.2f}'.format(memory_usage).zfill(5) + 'GB',
315
+ 'execution_threads': self.num_threads
316
+ })
317
+ progress.update(1)
318
+ if self.progress_gradio is not None:
319
+ self.progress_gradio((progress.n, self.total_frames), desc='Processing', total=self.total_frames, unit='frames')
320
+
321
+
322
+
323
+ def process_frame(self, frame:Frame):
324
+ if len(self.input_face_datas) < 1 and not self.options.show_face_masking:
325
+ return frame
326
+ temp_frame = frame.copy()
327
+ num_swapped, temp_frame = self.swap_faces(frame, temp_frame)
328
+ if num_swapped > 0:
329
+ if roop.globals.no_face_action == eNoFaceAction.SKIP_FRAME_IF_DISSIMILAR:
330
+ if len(self.input_face_datas) > num_swapped:
331
+ return None
332
+ return temp_frame
333
+ if roop.globals.no_face_action == eNoFaceAction.USE_ORIGINAL_FRAME:
334
+ return frame
335
+ if roop.globals.no_face_action == eNoFaceAction.SKIP_FRAME:
336
+ #This only works with in-mem processing, as it simply skips the frame.
337
+ #For 'extract frames' it simply leaves the unprocessed frame unprocessed and it gets used in the final output by ffmpeg.
338
+ #If we could delete that frame here, that'd work but that might cause ffmpeg to fail unless the frames are renamed, and I don't think we have the info on what frame it actually is?????
339
+ #alternatively, it could mark all the necessary frames for deletion, delete them at the end, then rename the remaining frames that might work?
340
+ return None
341
+ else:
342
+ return self.retry_rotated(frame)
343
+
344
+ def retry_rotated(self, frame):
345
+ copyframe = frame.copy()
346
+ copyframe = rotate_clockwise(copyframe)
347
+ temp_frame = copyframe.copy()
348
+ num_swapped, temp_frame = self.swap_faces(copyframe, temp_frame)
349
+ if num_swapped > 0:
350
+ return rotate_anticlockwise(temp_frame)
351
+
352
+ copyframe = frame.copy()
353
+ copyframe = rotate_anticlockwise(copyframe)
354
+ temp_frame = copyframe.copy()
355
+ num_swapped, temp_frame = self.swap_faces(copyframe, temp_frame)
356
+ if num_swapped > 0:
357
+ return rotate_clockwise(temp_frame)
358
+ del copyframe
359
+ return frame
360
+
361
+
362
+
363
+ def swap_faces(self, frame, temp_frame):
364
+ num_faces_found = 0
365
+
366
+ if self.options.swap_mode == "first":
367
+ face = get_first_face(frame)
368
+
369
+ if face is None:
370
+ return num_faces_found, frame
371
+
372
+ num_faces_found += 1
373
+ temp_frame = self.process_face(self.options.selected_index, face, temp_frame)
374
+ else:
375
+ faces = get_all_faces(frame)
376
+ if faces is None:
377
+ return num_faces_found, frame
378
+
379
+ if self.options.swap_mode == "all":
380
+ for face in faces:
381
+ num_faces_found += 1
382
+ temp_frame = self.process_face(self.options.selected_index, face, temp_frame)
383
+ del face
384
+
385
+ elif self.options.swap_mode == "selected":
386
+ num_targetfaces = len(self.target_face_datas)
387
+ use_index = num_targetfaces == 1
388
+ for i,tf in enumerate(self.target_face_datas):
389
+ for face in faces:
390
+ if compute_cosine_distance(tf.embedding, face.embedding) <= self.options.face_distance_threshold:
391
+ if i < len(self.input_face_datas):
392
+ if use_index:
393
+ temp_frame = self.process_face(self.options.selected_index, face, temp_frame)
394
+ else:
395
+ temp_frame = self.process_face(i, face, temp_frame)
396
+ num_faces_found += 1
397
+ del face
398
+ if not roop.globals.vr_mode and num_faces_found == num_targetfaces:
399
+ break
400
+ elif self.options.swap_mode == "all_female" or self.options.swap_mode == "all_male":
401
+ gender = 'F' if self.options.swap_mode == "all_female" else 'M'
402
+ for face in faces:
403
+ if face.sex == gender:
404
+ num_faces_found += 1
405
+ temp_frame = self.process_face(self.options.selected_index, face, temp_frame)
406
+ del face
407
+
408
+ if roop.globals.vr_mode and num_faces_found % 2 > 0:
409
+ # stereo image, there has to be an even number of faces
410
+ num_faces_found = 0
411
+ return num_faces_found, frame
412
+ if num_faces_found == 0:
413
+ return num_faces_found, frame
414
+
415
+ #maskprocessor = next((x for x in self.processors if x.type == 'mask'), None)
416
+
417
+ if self.options.imagemask is not None and self.options.imagemask.shape == frame.shape:
418
+ temp_frame = self.simple_blend_with_mask(temp_frame, frame, self.options.imagemask)
419
+ return num_faces_found, temp_frame
420
+
421
+
422
+ def rotation_action(self, original_face:Face, frame:Frame):
423
+ (height, width) = frame.shape[:2]
424
+
425
+ bounding_box_width = original_face.bbox[2] - original_face.bbox[0]
426
+ bounding_box_height = original_face.bbox[3] - original_face.bbox[1]
427
+ horizontal_face = bounding_box_width > bounding_box_height
428
+
429
+ center_x = width // 2.0
430
+ start_x = original_face.bbox[0]
431
+ end_x = original_face.bbox[2]
432
+ bbox_center_x = start_x + (bounding_box_width // 2.0)
433
+
434
+ # need to leverage the array of landmarks as decribed here:
435
+ # https://github.com/deepinsight/insightface/tree/master/alignment/coordinate_reg
436
+ # basically, we should be able to check for the relative position of eyes and nose
437
+ # then use that to determine which way the face is actually facing when in a horizontal position
438
+ # and use that to determine the correct rotation_action
439
+
440
+ forehead_x = original_face.landmark_2d_106[72][0]
441
+ chin_x = original_face.landmark_2d_106[0][0]
442
+
443
+ if horizontal_face:
444
+ if chin_x < forehead_x:
445
+ # this is someone lying down with their face like this (:
446
+ return "rotate_anticlockwise"
447
+ elif forehead_x < chin_x:
448
+ # this is someone lying down with their face like this :)
449
+ return "rotate_clockwise"
450
+ if bbox_center_x >= center_x:
451
+ # this is someone lying down with their face in the right hand side of the frame
452
+ return "rotate_anticlockwise"
453
+ if bbox_center_x < center_x:
454
+ # this is someone lying down with their face in the left hand side of the frame
455
+ return "rotate_clockwise"
456
+
457
+ return None
458
+
459
+
460
+ def auto_rotate_frame(self, original_face, frame:Frame):
461
+ target_face = original_face
462
+ original_frame = frame
463
+
464
+ rotation_action = self.rotation_action(original_face, frame)
465
+
466
+ if rotation_action == "rotate_anticlockwise":
467
+ #face is horizontal, rotating frame anti-clockwise and getting face bounding box from rotated frame
468
+ frame = rotate_anticlockwise(frame)
469
+ elif rotation_action == "rotate_clockwise":
470
+ #face is horizontal, rotating frame clockwise and getting face bounding box from rotated frame
471
+ frame = rotate_clockwise(frame)
472
+
473
+ return target_face, frame, rotation_action
474
+
475
+
476
+ def auto_unrotate_frame(self, frame:Frame, rotation_action):
477
+ if rotation_action == "rotate_anticlockwise":
478
+ return rotate_clockwise(frame)
479
+ elif rotation_action == "rotate_clockwise":
480
+ return rotate_anticlockwise(frame)
481
+
482
+ return frame
483
+
484
+
485
+
486
+ def process_face(self,face_index, target_face:Face, frame:Frame):
487
+ from roop.face_util import align_crop
488
+
489
+ enhanced_frame = None
490
+ if(len(self.input_face_datas) > 0):
491
+ inputface = self.input_face_datas[face_index].faces[0]
492
+ else:
493
+ inputface = None
494
+
495
+ rotation_action = None
496
+ if roop.globals.autorotate_faces:
497
+ # check for sideways rotation of face
498
+ rotation_action = self.rotation_action(target_face, frame)
499
+ if rotation_action is not None:
500
+ (startX, startY, endX, endY) = target_face["bbox"].astype("int")
501
+ width = endX - startX
502
+ height = endY - startY
503
+ offs = int(max(width,height) * 0.25)
504
+ rotcutframe,startX, startY, endX, endY = self.cutout(frame, startX - offs, startY - offs, endX + offs, endY + offs)
505
+ if rotation_action == "rotate_anticlockwise":
506
+ rotcutframe = rotate_anticlockwise(rotcutframe)
507
+ elif rotation_action == "rotate_clockwise":
508
+ rotcutframe = rotate_clockwise(rotcutframe)
509
+ # rotate image and re-detect face to correct wonky landmarks
510
+ rotface = get_first_face(rotcutframe)
511
+ if rotface is None:
512
+ rotation_action = None
513
+ else:
514
+ saved_frame = frame.copy()
515
+ frame = rotcutframe
516
+ target_face = rotface
517
+
518
+
519
+
520
+ # if roop.globals.vr_mode:
521
+ # bbox = target_face.bbox
522
+ # [orig_width, orig_height, _] = frame.shape
523
+
524
+ # # Convert bounding box to ints
525
+ # x1, y1, x2, y2 = map(int, bbox)
526
+
527
+ # # Determine the center of the bounding box
528
+ # x_center = (x1 + x2) / 2
529
+ # y_center = (y1 + y2) / 2
530
+
531
+ # # Normalize coordinates to range [-1, 1]
532
+ # x_center_normalized = x_center / (orig_width / 2) - 1
533
+ # y_center_normalized = y_center / (orig_width / 2) - 1
534
+
535
+ # # Convert normalized coordinates to spherical (theta, phi)
536
+ # theta = x_center_normalized * 180 # Theta ranges from -180 to 180 degrees
537
+ # phi = -y_center_normalized * 90 # Phi ranges from -90 to 90 degrees
538
+
539
+ # img = vr.GetPerspective(frame, 90, theta, phi, 1280, 1280) # Generate perspective image
540
+
541
+
542
+ """ Code ported/adapted from Facefusion which borrowed the idea from Rope:
543
+ Kind of subsampling the cutout and aligned face image and faceswapping slices of it up to
544
+ the desired output resolution. This works around the current resolution limitations without using enhancers.
545
+ """
546
+ model_output_size = 128
547
+ subsample_size = self.options.subsample_size
548
+ subsample_total = subsample_size // model_output_size
549
+ aligned_img, M = align_crop(frame, target_face.kps, subsample_size)
550
+
551
+ fake_frame = aligned_img
552
+ target_face.matrix = M
553
+
554
+ for p in self.processors:
555
+ if p.type == 'swap':
556
+ swap_result_frames = []
557
+ subsample_frames = self.implode_pixel_boost(aligned_img, model_output_size, subsample_total)
558
+ for sliced_frame in subsample_frames:
559
+ for _ in range(0,self.options.num_swap_steps):
560
+ sliced_frame = self.prepare_crop_frame(sliced_frame)
561
+ sliced_frame = p.Run(inputface, target_face, sliced_frame)
562
+ sliced_frame = self.normalize_swap_frame(sliced_frame)
563
+ swap_result_frames.append(sliced_frame)
564
+ fake_frame = self.explode_pixel_boost(swap_result_frames, model_output_size, subsample_total, subsample_size)
565
+ fake_frame = fake_frame.astype(np.uint8)
566
+ scale_factor = 0.0
567
+ elif p.type == 'mask':
568
+ fake_frame = self.process_mask(p, aligned_img, fake_frame)
569
+ else:
570
+ enhanced_frame, scale_factor = p.Run(self.input_face_datas[face_index], target_face, fake_frame)
571
+
572
+ upscale = 512
573
+ orig_width = fake_frame.shape[1]
574
+ if orig_width != upscale:
575
+ fake_frame = cv2.resize(fake_frame, (upscale, upscale), cv2.INTER_CUBIC)
576
+ mask_offsets = (0,0,0,0,1,20) if inputface is None else inputface.mask_offsets
577
+
578
+
579
+ if enhanced_frame is None:
580
+ scale_factor = int(upscale / orig_width)
581
+ result = self.paste_upscale(fake_frame, fake_frame, target_face.matrix, frame, scale_factor, mask_offsets)
582
+ else:
583
+ result = self.paste_upscale(fake_frame, enhanced_frame, target_face.matrix, frame, scale_factor, mask_offsets)
584
+
585
+ if rotation_action is not None:
586
+ fake_frame = self.auto_unrotate_frame(result, rotation_action)
587
+ return self.paste_simple(fake_frame, saved_frame, startX, startY)
588
+
589
+ return result
590
+
591
+
592
+
593
+
594
+ def cutout(self, frame:Frame, start_x, start_y, end_x, end_y):
595
+ if start_x < 0:
596
+ start_x = 0
597
+ if start_y < 0:
598
+ start_y = 0
599
+ if end_x > frame.shape[1]:
600
+ end_x = frame.shape[1]
601
+ if end_y > frame.shape[0]:
602
+ end_y = frame.shape[0]
603
+ return frame[start_y:end_y, start_x:end_x], start_x, start_y, end_x, end_y
604
+
605
+ def paste_simple(self, src:Frame, dest:Frame, start_x, start_y):
606
+ end_x = start_x + src.shape[1]
607
+ end_y = start_y + src.shape[0]
608
+
609
+ start_x, end_x, start_y, end_y = clamp_cut_values(start_x, end_x, start_y, end_y, dest)
610
+ dest[start_y:end_y, start_x:end_x] = src
611
+ return dest
612
+
613
+ def simple_blend_with_mask(self, image1, image2, mask):
614
+ # Blend the images
615
+ blended_image = image1.astype(np.float32) * (1.0 - mask) + image2.astype(np.float32) * mask
616
+ return blended_image.astype(np.uint8)
617
+
618
+
619
+ def paste_upscale(self, fake_face, upsk_face, M, target_img, scale_factor, mask_offsets):
620
+ M_scale = M * scale_factor
621
+ IM = cv2.invertAffineTransform(M_scale)
622
+
623
+ face_matte = np.full((target_img.shape[0],target_img.shape[1]), 255, dtype=np.uint8)
624
+ # Generate white square sized as a upsk_face
625
+ img_matte = np.zeros((upsk_face.shape[0],upsk_face.shape[1]), dtype=np.uint8)
626
+
627
+ w = img_matte.shape[1]
628
+ h = img_matte.shape[0]
629
+
630
+ top = int(mask_offsets[0] * h)
631
+ bottom = int(h - (mask_offsets[1] * h))
632
+ left = int(mask_offsets[2] * w)
633
+ right = int(w - (mask_offsets[3] * w))
634
+ img_matte[top:bottom,left:right] = 255
635
+
636
+ # Transform white square back to target_img
637
+ img_matte = cv2.warpAffine(img_matte, IM, (target_img.shape[1], target_img.shape[0]), flags=cv2.INTER_NEAREST, borderValue=0.0)
638
+ ##Blacken the edges of face_matte by 1 pixels (so the mask in not expanded on the image edges)
639
+ img_matte[:1,:] = img_matte[-1:,:] = img_matte[:,:1] = img_matte[:,-1:] = 0
640
+
641
+ img_matte = self.blur_area(img_matte, mask_offsets[4], mask_offsets[5])
642
+ #Normalize images to float values and reshape
643
+ img_matte = img_matte.astype(np.float32)/255
644
+ face_matte = face_matte.astype(np.float32)/255
645
+ img_matte = np.minimum(face_matte, img_matte)
646
+ if self.options.show_face_area_overlay:
647
+ # Additional steps for green overlay
648
+ green_overlay = np.zeros_like(target_img)
649
+ green_color = [0, 255, 0] # RGB for green
650
+ for i in range(3): # Apply green color where img_matte is not zero
651
+ green_overlay[:, :, i] = np.where(img_matte > 0, green_color[i], 0) ##Transform upcaled face back to target_img
652
+ img_matte = np.reshape(img_matte, [img_matte.shape[0],img_matte.shape[1],1])
653
+ paste_face = cv2.warpAffine(upsk_face, IM, (target_img.shape[1], target_img.shape[0]), borderMode=cv2.BORDER_REPLICATE)
654
+ if upsk_face is not fake_face:
655
+ fake_face = cv2.warpAffine(fake_face, IM, (target_img.shape[1], target_img.shape[0]), borderMode=cv2.BORDER_REPLICATE)
656
+ paste_face = cv2.addWeighted(paste_face, self.options.blend_ratio, fake_face, 1.0 - self.options.blend_ratio, 0)
657
+
658
+ # Re-assemble image
659
+ paste_face = img_matte * paste_face
660
+ paste_face = paste_face + (1-img_matte) * target_img.astype(np.float32)
661
+ if self.options.show_face_area_overlay:
662
+ # Overlay the green overlay on the final image
663
+ paste_face = cv2.addWeighted(paste_face.astype(np.uint8), 1 - 0.5, green_overlay, 0.5, 0)
664
+ return paste_face.astype(np.uint8)
665
+
666
+
667
+ def blur_area(self, img_matte, num_erosion_iterations, blur_amount):
668
+ # Detect the affine transformed white area
669
+ mask_h_inds, mask_w_inds = np.where(img_matte==255)
670
+ # Calculate the size (and diagonal size) of transformed white area width and height boundaries
671
+ mask_h = np.max(mask_h_inds) - np.min(mask_h_inds)
672
+ mask_w = np.max(mask_w_inds) - np.min(mask_w_inds)
673
+ mask_size = int(np.sqrt(mask_h*mask_w))
674
+ # Calculate the kernel size for eroding img_matte by kernel (insightface empirical guess for best size was max(mask_size//10,10))
675
+ # k = max(mask_size//12, 8)
676
+ k = max(mask_size//(blur_amount // 2) , blur_amount // 2)
677
+ kernel = np.ones((k,k),np.uint8)
678
+ img_matte = cv2.erode(img_matte,kernel,iterations = num_erosion_iterations)
679
+ #Calculate the kernel size for blurring img_matte by blur_size (insightface empirical guess for best size was max(mask_size//20, 5))
680
+ # k = max(mask_size//24, 4)
681
+ k = max(mask_size//blur_amount, blur_amount//5)
682
+ kernel_size = (k, k)
683
+ blur_size = tuple(2*i+1 for i in kernel_size)
684
+ return cv2.GaussianBlur(img_matte, blur_size, 0)
685
+
686
+
687
+ def prepare_crop_frame(self, swap_frame):
688
+ model_type = 'inswapper'
689
+ model_mean = [0.0, 0.0, 0.0]
690
+ model_standard_deviation = [1.0, 1.0, 1.0]
691
+
692
+ if model_type == 'ghost':
693
+ swap_frame = swap_frame[:, :, ::-1] / 127.5 - 1
694
+ else:
695
+ swap_frame = swap_frame[:, :, ::-1] / 255.0
696
+ swap_frame = (swap_frame - model_mean) / model_standard_deviation
697
+ swap_frame = swap_frame.transpose(2, 0, 1)
698
+ swap_frame = np.expand_dims(swap_frame, axis = 0).astype(np.float32)
699
+ return swap_frame
700
+
701
+
702
+ def normalize_swap_frame(self, swap_frame):
703
+ model_type = 'inswapper'
704
+ swap_frame = swap_frame.transpose(1, 2, 0)
705
+
706
+ if model_type == 'ghost':
707
+ swap_frame = (swap_frame * 127.5 + 127.5).round()
708
+ else:
709
+ swap_frame = (swap_frame * 255.0).round()
710
+ swap_frame = swap_frame[:, :, ::-1]
711
+ return swap_frame
712
+
713
+ def implode_pixel_boost(self, aligned_face_frame, model_size, pixel_boost_total : int):
714
+ subsample_frame = aligned_face_frame.reshape(model_size, pixel_boost_total, model_size, pixel_boost_total, 3)
715
+ subsample_frame = subsample_frame.transpose(1, 3, 0, 2, 4).reshape(pixel_boost_total ** 2, model_size, model_size, 3)
716
+ return subsample_frame
717
+
718
+
719
+ def explode_pixel_boost(self, subsample_frame, model_size, pixel_boost_total, pixel_boost_size):
720
+ final_frame = np.stack(subsample_frame, axis = 0).reshape(pixel_boost_total, pixel_boost_total, model_size, model_size, 3)
721
+ final_frame = final_frame.transpose(2, 0, 3, 1, 4).reshape(pixel_boost_size, pixel_boost_size, 3)
722
+ return final_frame
723
+
724
+ def process_mask(self, processor, frame:Frame, target:Frame):
725
+ img_mask = processor.Run(frame, self.options.masking_text)
726
+ img_mask = cv2.resize(img_mask, (target.shape[1], target.shape[0]))
727
+ img_mask = np.reshape(img_mask, [img_mask.shape[0],img_mask.shape[1],1])
728
+
729
+ if self.options.show_face_masking:
730
+ result = (1 - img_mask) * frame.astype(np.float32)
731
+ return np.uint8(result)
732
+
733
+
734
+ target = target.astype(np.float32)
735
+ result = (1-img_mask) * target
736
+ result += img_mask * frame.astype(np.float32)
737
+ return np.uint8(result)
738
+
739
+
740
+
741
+
742
+ def unload_models():
743
+ pass
744
+
745
+
746
+ def release_resources(self):
747
+ for p in self.processors:
748
+ p.Release()
749
+ self.processors.clear()
750
+
roop/ProcessOptions.py ADDED
@@ -0,0 +1,14 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ class ProcessOptions:
2
+
3
+ def __init__(self, processordefines:dict, face_distance, blend_ratio, swap_mode, selected_index, masking_text, imagemask, num_steps, subsample_size, show_face_area, show_mask=False):
4
+ self.processors = processordefines
5
+ self.face_distance_threshold = face_distance
6
+ self.blend_ratio = blend_ratio
7
+ self.swap_mode = swap_mode
8
+ self.selected_index = selected_index
9
+ self.masking_text = masking_text
10
+ self.imagemask = imagemask
11
+ self.num_swap_steps = num_steps
12
+ self.show_face_area_overlay = show_face_area
13
+ self.show_face_masking = show_mask
14
+ self.subsample_size = subsample_size
roop/__init__.py ADDED
File without changes
roop/capturer.py ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Optional
2
+ import cv2
3
+ import numpy as np
4
+
5
+ from roop.typing import Frame
6
+
7
+ def get_image_frame(filename: str):
8
+ try:
9
+ return cv2.imdecode(np.fromfile(filename, dtype=np.uint8), cv2.IMREAD_COLOR)
10
+ except:
11
+ print(f"Exception reading {filename}")
12
+ return None
13
+
14
+
15
+ def get_video_frame(video_path: str, frame_number: int = 0) -> Optional[Frame]:
16
+ capture = cv2.VideoCapture(video_path)
17
+ frame_total = capture.get(cv2.CAP_PROP_FRAME_COUNT)
18
+ capture.set(cv2.CAP_PROP_POS_FRAMES, min(frame_total, frame_number - 1))
19
+ has_frame, frame = capture.read()
20
+ capture.release()
21
+ if has_frame:
22
+ return frame
23
+ return None
24
+
25
+
26
+ def get_video_frame_total(video_path: str) -> int:
27
+ capture = cv2.VideoCapture(video_path)
28
+ video_frame_total = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
29
+ capture.release()
30
+ return video_frame_total
roop/core.py ADDED
@@ -0,0 +1,380 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+
3
+ import os
4
+ import sys
5
+ import shutil
6
+ # single thread doubles cuda performance - needs to be set before torch import
7
+ if any(arg.startswith('--execution-provider') for arg in sys.argv):
8
+ os.environ['OMP_NUM_THREADS'] = '1'
9
+
10
+ import warnings
11
+ from typing import List
12
+ import platform
13
+ import signal
14
+ import torch
15
+ import onnxruntime
16
+ import pathlib
17
+
18
+ from time import time
19
+
20
+ import roop.globals
21
+ import roop.metadata
22
+ import roop.utilities as util
23
+ import roop.util_ffmpeg as ffmpeg
24
+ import ui.main as main
25
+ from settings import Settings
26
+ from roop.face_util import extract_face_images
27
+ from roop.ProcessEntry import ProcessEntry
28
+ from roop.ProcessMgr import ProcessMgr
29
+ from roop.ProcessOptions import ProcessOptions
30
+ from roop.capturer import get_video_frame_total
31
+
32
+
33
+ clip_text = None
34
+
35
+ call_display_ui = None
36
+
37
+ process_mgr = None
38
+
39
+
40
+ if 'ROCMExecutionProvider' in roop.globals.execution_providers:
41
+ del torch
42
+
43
+ warnings.filterwarnings('ignore', category=FutureWarning, module='insightface')
44
+ warnings.filterwarnings('ignore', category=UserWarning, module='torchvision')
45
+
46
+
47
+ def parse_args() -> None:
48
+ signal.signal(signal.SIGINT, lambda signal_number, frame: destroy())
49
+ roop.globals.headless = False
50
+ # Always enable all processors when using GUI
51
+ if len(sys.argv) > 1:
52
+ print('No CLI args supported - use Settings Tab instead')
53
+ roop.globals.frame_processors = ['face_swapper', 'face_enhancer']
54
+
55
+
56
+ def encode_execution_providers(execution_providers: List[str]) -> List[str]:
57
+ return [execution_provider.replace('ExecutionProvider', '').lower() for execution_provider in execution_providers]
58
+
59
+
60
+ def decode_execution_providers(execution_providers: List[str]) -> List[str]:
61
+ return [provider for provider, encoded_execution_provider in zip(onnxruntime.get_available_providers(), encode_execution_providers(onnxruntime.get_available_providers()))
62
+ if any(execution_provider in encoded_execution_provider for execution_provider in execution_providers)]
63
+
64
+
65
+ def suggest_max_memory() -> int:
66
+ if platform.system().lower() == 'darwin':
67
+ return 4
68
+ return 16
69
+
70
+
71
+ def suggest_execution_providers() -> List[str]:
72
+ return encode_execution_providers(onnxruntime.get_available_providers())
73
+
74
+
75
+ def suggest_execution_threads() -> int:
76
+ if 'DmlExecutionProvider' in roop.globals.execution_providers:
77
+ return 1
78
+ if 'ROCMExecutionProvider' in roop.globals.execution_providers:
79
+ return 1
80
+ return 8
81
+
82
+
83
+ def limit_resources() -> None:
84
+ # limit memory usage
85
+ if roop.globals.max_memory:
86
+ memory = roop.globals.max_memory * 1024 ** 3
87
+ if platform.system().lower() == 'darwin':
88
+ memory = roop.globals.max_memory * 1024 ** 6
89
+ if platform.system().lower() == 'windows':
90
+ import ctypes
91
+ kernel32 = ctypes.windll.kernel32 # type: ignore[attr-defined]
92
+ kernel32.SetProcessWorkingSetSize(-1, ctypes.c_size_t(memory), ctypes.c_size_t(memory))
93
+ else:
94
+ import resource
95
+ resource.setrlimit(resource.RLIMIT_DATA, (memory, memory))
96
+
97
+
98
+
99
+ def release_resources() -> None:
100
+ import gc
101
+ global process_mgr
102
+
103
+ if process_mgr is not None:
104
+ process_mgr.release_resources()
105
+ process_mgr = None
106
+
107
+ gc.collect()
108
+ # if 'CUDAExecutionProvider' in roop.globals.execution_providers and torch.cuda.is_available():
109
+ # with torch.cuda.device('cuda'):
110
+ # torch.cuda.empty_cache()
111
+ # torch.cuda.ipc_collect()
112
+
113
+
114
+ def pre_check() -> bool:
115
+ if sys.version_info < (3, 9):
116
+ update_status('Python version is not supported - please upgrade to 3.9 or higher.')
117
+ return False
118
+
119
+ download_directory_path = util.resolve_relative_path('../models')
120
+ util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/inswapper_128.onnx'])
121
+ util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/GFPGANv1.4.onnx'])
122
+ util.conditional_download(download_directory_path, ['https://github.com/csxmli2016/DMDNet/releases/download/v1/DMDNet.pth'])
123
+ util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/GPEN-BFR-512.onnx'])
124
+ util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/restoreformer_plus_plus.onnx'])
125
+ util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/xseg.onnx'])
126
+ download_directory_path = util.resolve_relative_path('../models/CLIP')
127
+ util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/rd64-uni-refined.pth'])
128
+ download_directory_path = util.resolve_relative_path('../models/CodeFormer')
129
+ util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/CodeFormerv0.1.onnx'])
130
+ download_directory_path = util.resolve_relative_path('../models/Frame')
131
+ util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/deoldify_artistic.onnx'])
132
+ util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/deoldify_stable.onnx'])
133
+ util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/isnet-general-use.onnx'])
134
+ util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/real_esrgan_x4.onnx'])
135
+ util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/real_esrgan_x2.onnx'])
136
+ util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/lsdir_x4.onnx'])
137
+
138
+ if not shutil.which('ffmpeg'):
139
+ update_status('ffmpeg is not installed.')
140
+ return True
141
+
142
+ def set_display_ui(function):
143
+ global call_display_ui
144
+
145
+ call_display_ui = function
146
+
147
+
148
+ def update_status(message: str) -> None:
149
+ global call_display_ui
150
+
151
+ print(message)
152
+ if call_display_ui is not None:
153
+ call_display_ui(message)
154
+
155
+
156
+
157
+
158
+ def start() -> None:
159
+ if roop.globals.headless:
160
+ print('Headless mode currently unsupported - starting UI!')
161
+ # faces = extract_face_images(roop.globals.source_path, (False, 0))
162
+ # roop.globals.INPUT_FACES.append(faces[roop.globals.source_face_index])
163
+ # faces = extract_face_images(roop.globals.target_path, (False, util.has_image_extension(roop.globals.target_path)))
164
+ # roop.globals.TARGET_FACES.append(faces[roop.globals.target_face_index])
165
+ # if 'face_enhancer' in roop.globals.frame_processors:
166
+ # roop.globals.selected_enhancer = 'GFPGAN'
167
+
168
+ batch_process_regular(None, False, None)
169
+
170
+
171
+ def get_processing_plugins(masking_engine):
172
+ processors = { "faceswap": {}}
173
+ if masking_engine is not None:
174
+ processors.update({masking_engine: {}})
175
+
176
+ if roop.globals.selected_enhancer == 'GFPGAN':
177
+ processors.update({"gfpgan": {}})
178
+ elif roop.globals.selected_enhancer == 'Codeformer':
179
+ processors.update({"codeformer": {}})
180
+ elif roop.globals.selected_enhancer == 'DMDNet':
181
+ processors.update({"dmdnet": {}})
182
+ elif roop.globals.selected_enhancer == 'GPEN':
183
+ processors.update({"gpen": {}})
184
+ elif roop.globals.selected_enhancer == 'Restoreformer++':
185
+ processors.update({"restoreformer++": {}})
186
+ return processors
187
+
188
+
189
+ def live_swap(frame, options):
190
+ global process_mgr
191
+
192
+ if frame is None:
193
+ return frame
194
+
195
+ if process_mgr is None:
196
+ process_mgr = ProcessMgr(None)
197
+
198
+ # if len(roop.globals.INPUT_FACESETS) <= selected_index:
199
+ # selected_index = 0
200
+ process_mgr.initialize(roop.globals.INPUT_FACESETS, roop.globals.TARGET_FACES, options)
201
+ newframe = process_mgr.process_frame(frame)
202
+ if newframe is None:
203
+ return frame
204
+ return newframe
205
+
206
+
207
+ def batch_process_regular(files:list[ProcessEntry], masking_engine:str, new_clip_text:str, use_new_method, imagemask, num_swap_steps, progress, selected_index = 0) -> None:
208
+ global clip_text, process_mgr
209
+
210
+ release_resources()
211
+ limit_resources()
212
+ if process_mgr is None:
213
+ process_mgr = ProcessMgr(progress)
214
+ mask = imagemask["layers"][0] if imagemask is not None else None
215
+ if len(roop.globals.INPUT_FACESETS) <= selected_index:
216
+ selected_index = 0
217
+ options = ProcessOptions(get_processing_plugins(masking_engine), roop.globals.distance_threshold, roop.globals.blend_ratio,
218
+ roop.globals.face_swap_mode, selected_index, new_clip_text, mask, num_swap_steps,
219
+ roop.globals.subsample_size, False)
220
+ process_mgr.initialize(roop.globals.INPUT_FACESETS, roop.globals.TARGET_FACES, options)
221
+ batch_process(files, use_new_method)
222
+ return
223
+
224
+ def batch_process_with_options(files:list[ProcessEntry], options, progress):
225
+ global clip_text, process_mgr
226
+
227
+ release_resources()
228
+ limit_resources()
229
+ if process_mgr is None:
230
+ process_mgr = ProcessMgr(progress)
231
+ process_mgr.initialize(roop.globals.INPUT_FACESETS, roop.globals.TARGET_FACES, options)
232
+ roop.globals.keep_frames = False
233
+ roop.globals.wait_after_extraction = False
234
+ roop.globals.skip_audio = False
235
+ batch_process(files, True)
236
+
237
+
238
+
239
+ def batch_process(files:list[ProcessEntry], use_new_method) -> None:
240
+ global clip_text, process_mgr
241
+
242
+ roop.globals.processing = True
243
+
244
+ # limit threads for some providers
245
+ max_threads = suggest_execution_threads()
246
+ if max_threads == 1:
247
+ roop.globals.execution_threads = 1
248
+
249
+ imagefiles:list[ProcessEntry] = []
250
+ videofiles:list[ProcessEntry] = []
251
+
252
+ update_status('Sorting videos/images')
253
+
254
+
255
+ for index, f in enumerate(files):
256
+ fullname = f.filename
257
+ if util.has_image_extension(fullname):
258
+ destination = util.get_destfilename_from_path(fullname, roop.globals.output_path, f'.{roop.globals.CFG.output_image_format}')
259
+ destination = util.replace_template(destination, index=index)
260
+ pathlib.Path(os.path.dirname(destination)).mkdir(parents=True, exist_ok=True)
261
+ f.finalname = destination
262
+ imagefiles.append(f)
263
+
264
+ elif util.is_video(fullname) or util.has_extension(fullname, ['gif']):
265
+ destination = util.get_destfilename_from_path(fullname, roop.globals.output_path, f'__temp.{roop.globals.CFG.output_video_format}')
266
+ f.finalname = destination
267
+ videofiles.append(f)
268
+
269
+
270
+
271
+ if(len(imagefiles) > 0):
272
+ update_status('Processing image(s)')
273
+ origimages = []
274
+ fakeimages = []
275
+ for f in imagefiles:
276
+ origimages.append(f.filename)
277
+ fakeimages.append(f.finalname)
278
+
279
+ process_mgr.run_batch(origimages, fakeimages, roop.globals.execution_threads)
280
+ origimages.clear()
281
+ fakeimages.clear()
282
+
283
+ if(len(videofiles) > 0):
284
+ for index,v in enumerate(videofiles):
285
+ if not roop.globals.processing:
286
+ end_processing('Processing stopped!')
287
+ return
288
+ fps = v.fps if v.fps > 0 else util.detect_fps(v.filename)
289
+ if v.endframe == 0:
290
+ v.endframe = get_video_frame_total(v.filename)
291
+
292
+ update_status(f'Creating {os.path.basename(v.finalname)} with {fps} FPS...')
293
+ start_processing = time()
294
+ if roop.globals.keep_frames or not use_new_method:
295
+ util.create_temp(v.filename)
296
+ update_status('Extracting frames...')
297
+ ffmpeg.extract_frames(v.filename,v.startframe,v.endframe, fps)
298
+ if not roop.globals.processing:
299
+ end_processing('Processing stopped!')
300
+ return
301
+
302
+ temp_frame_paths = util.get_temp_frame_paths(v.filename)
303
+ process_mgr.run_batch(temp_frame_paths, temp_frame_paths, roop.globals.execution_threads)
304
+ if not roop.globals.processing:
305
+ end_processing('Processing stopped!')
306
+ return
307
+ if roop.globals.wait_after_extraction:
308
+ extract_path = os.path.dirname(temp_frame_paths[0])
309
+ util.open_folder(extract_path)
310
+ input("Press any key to continue...")
311
+ print("Resorting frames to create video")
312
+ util.sort_rename_frames(extract_path)
313
+
314
+ ffmpeg.create_video(v.filename, v.finalname, fps)
315
+ if not roop.globals.keep_frames:
316
+ util.delete_temp_frames(temp_frame_paths[0])
317
+ else:
318
+ if util.has_extension(v.filename, ['gif']):
319
+ skip_audio = True
320
+ else:
321
+ skip_audio = roop.globals.skip_audio
322
+ process_mgr.run_batch_inmem(v.filename, v.finalname, v.startframe, v.endframe, fps,roop.globals.execution_threads, skip_audio)
323
+
324
+ if not roop.globals.processing:
325
+ end_processing('Processing stopped!')
326
+ return
327
+
328
+ video_file_name = v.finalname
329
+ if os.path.isfile(video_file_name):
330
+ destination = ''
331
+ if util.has_extension(v.filename, ['gif']):
332
+ gifname = util.get_destfilename_from_path(v.filename, roop.globals.output_path, '.gif')
333
+ destination = util.replace_template(gifname, index=index)
334
+ pathlib.Path(os.path.dirname(destination)).mkdir(parents=True, exist_ok=True)
335
+
336
+ update_status('Creating final GIF')
337
+ ffmpeg.create_gif_from_video(video_file_name, destination)
338
+ if os.path.isfile(destination):
339
+ os.remove(video_file_name)
340
+ else:
341
+ skip_audio = roop.globals.skip_audio
342
+ destination = util.replace_template(video_file_name, index=index)
343
+ pathlib.Path(os.path.dirname(destination)).mkdir(parents=True, exist_ok=True)
344
+
345
+ if not skip_audio:
346
+ ffmpeg.restore_audio(video_file_name, v.filename, v.startframe, v.endframe, destination)
347
+ if os.path.isfile(destination):
348
+ os.remove(video_file_name)
349
+ else:
350
+ shutil.move(video_file_name, destination)
351
+ update_status(f'\nProcessing {os.path.basename(destination)} took {time() - start_processing} secs')
352
+
353
+ else:
354
+ update_status(f'Failed processing {os.path.basename(v.finalname)}!')
355
+ end_processing('Finished')
356
+
357
+
358
+ def end_processing(msg:str):
359
+ update_status(msg)
360
+ roop.globals.target_folder_path = None
361
+ release_resources()
362
+
363
+
364
+ def destroy() -> None:
365
+ if roop.globals.target_path:
366
+ util.clean_temp(roop.globals.target_path)
367
+ release_resources()
368
+ sys.exit()
369
+
370
+
371
+ def run() -> None:
372
+ parse_args()
373
+ if not pre_check():
374
+ return
375
+ roop.globals.CFG = Settings('config.yaml')
376
+ roop.globals.execution_threads = roop.globals.CFG.max_threads
377
+ roop.globals.video_encoder = roop.globals.CFG.output_video_codec
378
+ roop.globals.video_quality = roop.globals.CFG.video_quality
379
+ roop.globals.max_memory = roop.globals.CFG.memory_limit if roop.globals.CFG.memory_limit > 0 else None
380
+ main.run()
roop/face_util.py ADDED
@@ -0,0 +1,309 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import threading
2
+ from typing import Any
3
+ import insightface
4
+
5
+ import roop.globals
6
+ from roop.typing import Frame, Face
7
+
8
+ import cv2
9
+ import numpy as np
10
+ from skimage import transform as trans
11
+ from roop.capturer import get_video_frame
12
+ from roop.utilities import resolve_relative_path, conditional_download
13
+
14
+ FACE_ANALYSER = None
15
+ THREAD_LOCK_ANALYSER = threading.Lock()
16
+ THREAD_LOCK_SWAPPER = threading.Lock()
17
+ FACE_SWAPPER = None
18
+
19
+
20
+ def get_face_analyser() -> Any:
21
+ global FACE_ANALYSER
22
+
23
+ with THREAD_LOCK_ANALYSER:
24
+ if FACE_ANALYSER is None or roop.globals.g_current_face_analysis != roop.globals.g_desired_face_analysis:
25
+ model_path = resolve_relative_path('..')
26
+ # removed genderage
27
+ allowed_modules = roop.globals.g_desired_face_analysis
28
+ roop.globals.g_current_face_analysis = roop.globals.g_desired_face_analysis
29
+ if roop.globals.CFG.force_cpu:
30
+ print("Forcing CPU for Face Analysis")
31
+ FACE_ANALYSER = insightface.app.FaceAnalysis(
32
+ name="buffalo_l",
33
+ root=model_path, providers=["CPUExecutionProvider"],allowed_modules=allowed_modules
34
+ )
35
+ else:
36
+ FACE_ANALYSER = insightface.app.FaceAnalysis(
37
+ name="buffalo_l", root=model_path, providers=roop.globals.execution_providers,allowed_modules=allowed_modules
38
+ )
39
+ FACE_ANALYSER.prepare(
40
+ ctx_id=0,
41
+ det_size=(640, 640) if roop.globals.default_det_size else (320, 320),
42
+ )
43
+ return FACE_ANALYSER
44
+
45
+
46
+ def get_first_face(frame: Frame) -> Any:
47
+ try:
48
+ faces = get_face_analyser().get(frame)
49
+ return min(faces, key=lambda x: x.bbox[0])
50
+ # return sorted(faces, reverse=True, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]))[0]
51
+ except:
52
+ return None
53
+
54
+
55
+ def get_all_faces(frame: Frame) -> Any:
56
+ try:
57
+ faces = get_face_analyser().get(frame)
58
+ return sorted(faces, key=lambda x: x.bbox[0])
59
+ except:
60
+ return None
61
+
62
+
63
+ def extract_face_images(source_filename, video_info, extra_padding=-1.0):
64
+ face_data = []
65
+ source_image = None
66
+
67
+ if video_info[0]:
68
+ frame = get_video_frame(source_filename, video_info[1])
69
+ if frame is not None:
70
+ source_image = frame
71
+ else:
72
+ return face_data
73
+ else:
74
+ source_image = cv2.imdecode(np.fromfile(source_filename, dtype=np.uint8), cv2.IMREAD_COLOR)
75
+
76
+ faces = get_all_faces(source_image)
77
+ if faces is None:
78
+ return face_data
79
+
80
+ i = 0
81
+ for face in faces:
82
+ (startX, startY, endX, endY) = face["bbox"].astype("int")
83
+ startX, endX, startY, endY = clamp_cut_values(startX, endX, startY, endY, source_image)
84
+ if extra_padding > 0.0:
85
+ if source_image.shape[:2] == (512, 512):
86
+ i += 1
87
+ face_data.append([face, source_image])
88
+ continue
89
+
90
+ found = False
91
+ for i in range(1, 3):
92
+ (startX, startY, endX, endY) = face["bbox"].astype("int")
93
+ startX, endX, startY, endY = clamp_cut_values(startX, endX, startY, endY, source_image)
94
+ cutout_padding = extra_padding
95
+ # top needs extra room for detection
96
+ padding = int((endY - startY) * cutout_padding)
97
+ oldY = startY
98
+ startY -= padding
99
+
100
+ factor = 0.25 if i == 1 else 0.5
101
+ cutout_padding = factor
102
+ padding = int((endY - oldY) * cutout_padding)
103
+ endY += padding
104
+ padding = int((endX - startX) * cutout_padding)
105
+ startX -= padding
106
+ endX += padding
107
+ startX, endX, startY, endY = clamp_cut_values(
108
+ startX, endX, startY, endY, source_image
109
+ )
110
+ face_temp = source_image[startY:endY, startX:endX]
111
+ face_temp = resize_image_keep_content(face_temp)
112
+ testfaces = get_all_faces(face_temp)
113
+ if testfaces is not None and len(testfaces) > 0:
114
+ i += 1
115
+ face_data.append([testfaces[0], face_temp])
116
+ found = True
117
+ break
118
+
119
+ if not found:
120
+ print("No face found after resizing, this shouldn't happen!")
121
+ continue
122
+
123
+ face_temp = source_image[startY:endY, startX:endX]
124
+ if face_temp.size < 1:
125
+ continue
126
+
127
+ i += 1
128
+ face_data.append([face, face_temp])
129
+ return face_data
130
+
131
+
132
+ def clamp_cut_values(startX, endX, startY, endY, image):
133
+ if startX < 0:
134
+ startX = 0
135
+ if endX > image.shape[1]:
136
+ endX = image.shape[1]
137
+ if startY < 0:
138
+ startY = 0
139
+ if endY > image.shape[0]:
140
+ endY = image.shape[0]
141
+ return startX, endX, startY, endY
142
+
143
+
144
+
145
+ def face_offset_top(face: Face, offset):
146
+ face["bbox"][1] += offset
147
+ face["bbox"][3] += offset
148
+ lm106 = face.landmark_2d_106
149
+ add = np.full_like(lm106, [0, offset])
150
+ face["landmark_2d_106"] = lm106 + add
151
+ return face
152
+
153
+
154
+ def resize_image_keep_content(image, new_width=512, new_height=512):
155
+ dim = None
156
+ (h, w) = image.shape[:2]
157
+ if h > w:
158
+ r = new_height / float(h)
159
+ dim = (int(w * r), new_height)
160
+ else:
161
+ # Calculate the ratio of the width and construct the dimensions
162
+ r = new_width / float(w)
163
+ dim = (new_width, int(h * r))
164
+ image = cv2.resize(image, dim, interpolation=cv2.INTER_AREA)
165
+ (h, w) = image.shape[:2]
166
+ if h == new_height and w == new_width:
167
+ return image
168
+ resize_img = np.zeros(shape=(new_height, new_width, 3), dtype=image.dtype)
169
+ offs = (new_width - w) if h == new_height else (new_height - h)
170
+ startoffs = int(offs // 2) if offs % 2 == 0 else int(offs // 2) + 1
171
+ offs = int(offs // 2)
172
+
173
+ if h == new_height:
174
+ resize_img[0:new_height, startoffs : new_width - offs] = image
175
+ else:
176
+ resize_img[startoffs : new_height - offs, 0:new_width] = image
177
+ return resize_img
178
+
179
+
180
+ def rotate_image_90(image, rotate=True):
181
+ if rotate:
182
+ return np.rot90(image)
183
+ else:
184
+ return np.rot90(image, 1, (1, 0))
185
+
186
+
187
+ def rotate_anticlockwise(frame):
188
+ return rotate_image_90(frame)
189
+
190
+
191
+ def rotate_clockwise(frame):
192
+ return rotate_image_90(frame, False)
193
+
194
+
195
+ def rotate_image_180(image):
196
+ return np.flip(image, 0)
197
+
198
+
199
+ # alignment code from insightface https://github.com/deepinsight/insightface/blob/master/python-package/insightface/utils/face_align.py
200
+
201
+ arcface_dst = np.array(
202
+ [
203
+ [38.2946, 51.6963],
204
+ [73.5318, 51.5014],
205
+ [56.0252, 71.7366],
206
+ [41.5493, 92.3655],
207
+ [70.7299, 92.2041],
208
+ ],
209
+ dtype=np.float32,
210
+ )
211
+
212
+
213
+ def estimate_norm(lmk, image_size=112):
214
+ assert lmk.shape == (5, 2)
215
+ if image_size % 112 == 0:
216
+ ratio = float(image_size) / 112.0
217
+ diff_x = 0
218
+ elif image_size % 128 == 0:
219
+ ratio = float(image_size) / 128.0
220
+ diff_x = 8.0 * ratio
221
+ elif image_size % 512 == 0:
222
+ ratio = float(image_size) / 512.0
223
+ diff_x = 32.0 * ratio
224
+
225
+ dst = arcface_dst * ratio
226
+ dst[:, 0] += diff_x
227
+ tform = trans.SimilarityTransform()
228
+ tform.estimate(lmk, dst)
229
+ M = tform.params[0:2, :]
230
+ return M
231
+
232
+
233
+
234
+ # aligned, M = norm_crop2(f[1], face.kps, 512)
235
+ def align_crop(img, landmark, image_size=112, mode="arcface"):
236
+ M = estimate_norm(landmark, image_size)
237
+ warped = cv2.warpAffine(img, M, (image_size, image_size), borderValue=0.0)
238
+ return warped, M
239
+
240
+
241
+ def square_crop(im, S):
242
+ if im.shape[0] > im.shape[1]:
243
+ height = S
244
+ width = int(float(im.shape[1]) / im.shape[0] * S)
245
+ scale = float(S) / im.shape[0]
246
+ else:
247
+ width = S
248
+ height = int(float(im.shape[0]) / im.shape[1] * S)
249
+ scale = float(S) / im.shape[1]
250
+ resized_im = cv2.resize(im, (width, height))
251
+ det_im = np.zeros((S, S, 3), dtype=np.uint8)
252
+ det_im[: resized_im.shape[0], : resized_im.shape[1], :] = resized_im
253
+ return det_im, scale
254
+
255
+
256
+ def transform(data, center, output_size, scale, rotation):
257
+ scale_ratio = scale
258
+ rot = float(rotation) * np.pi / 180.0
259
+ # translation = (output_size/2-center[0]*scale_ratio, output_size/2-center[1]*scale_ratio)
260
+ t1 = trans.SimilarityTransform(scale=scale_ratio)
261
+ cx = center[0] * scale_ratio
262
+ cy = center[1] * scale_ratio
263
+ t2 = trans.SimilarityTransform(translation=(-1 * cx, -1 * cy))
264
+ t3 = trans.SimilarityTransform(rotation=rot)
265
+ t4 = trans.SimilarityTransform(translation=(output_size / 2, output_size / 2))
266
+ t = t1 + t2 + t3 + t4
267
+ M = t.params[0:2]
268
+ cropped = cv2.warpAffine(data, M, (output_size, output_size), borderValue=0.0)
269
+ return cropped, M
270
+
271
+
272
+ def trans_points2d(pts, M):
273
+ new_pts = np.zeros(shape=pts.shape, dtype=np.float32)
274
+ for i in range(pts.shape[0]):
275
+ pt = pts[i]
276
+ new_pt = np.array([pt[0], pt[1], 1.0], dtype=np.float32)
277
+ new_pt = np.dot(M, new_pt)
278
+ # print('new_pt', new_pt.shape, new_pt)
279
+ new_pts[i] = new_pt[0:2]
280
+
281
+ return new_pts
282
+
283
+
284
+ def trans_points3d(pts, M):
285
+ scale = np.sqrt(M[0][0] * M[0][0] + M[0][1] * M[0][1])
286
+ # print(scale)
287
+ new_pts = np.zeros(shape=pts.shape, dtype=np.float32)
288
+ for i in range(pts.shape[0]):
289
+ pt = pts[i]
290
+ new_pt = np.array([pt[0], pt[1], 1.0], dtype=np.float32)
291
+ new_pt = np.dot(M, new_pt)
292
+ # print('new_pt', new_pt.shape, new_pt)
293
+ new_pts[i][0:2] = new_pt[0:2]
294
+ new_pts[i][2] = pts[i][2] * scale
295
+
296
+ return new_pts
297
+
298
+
299
+ def trans_points(pts, M):
300
+ if pts.shape[1] == 2:
301
+ return trans_points2d(pts, M)
302
+ else:
303
+ return trans_points3d(pts, M)
304
+
305
+ def create_blank_image(width, height):
306
+ img = np.zeros((height, width, 4), dtype=np.uint8)
307
+ img[:] = [0,0,0,0]
308
+ return img
309
+
roop/ffmpeg_writer.py ADDED
@@ -0,0 +1,218 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ FFMPEG_Writer - write set of frames to video file
3
+
4
+ original from
5
+ https://github.com/Zulko/moviepy/blob/master/moviepy/video/io/ffmpeg_writer.py
6
+
7
+ removed unnecessary dependencies
8
+
9
+ The MIT License (MIT)
10
+
11
+ Copyright (c) 2015 Zulko
12
+ Copyright (c) 2023 Janvarev Vladislav
13
+ """
14
+
15
+ import os
16
+ import subprocess as sp
17
+
18
+ PIPE = -1
19
+ STDOUT = -2
20
+ DEVNULL = -3
21
+
22
+ FFMPEG_BINARY = "ffmpeg"
23
+
24
+ class FFMPEG_VideoWriter:
25
+ """ A class for FFMPEG-based video writing.
26
+
27
+ A class to write videos using ffmpeg. ffmpeg will write in a large
28
+ choice of formats.
29
+
30
+ Parameters
31
+ -----------
32
+
33
+ filename
34
+ Any filename like 'video.mp4' etc. but if you want to avoid
35
+ complications it is recommended to use the generic extension
36
+ '.avi' for all your videos.
37
+
38
+ size
39
+ Size (width,height) of the output video in pixels.
40
+
41
+ fps
42
+ Frames per second in the output video file.
43
+
44
+ codec
45
+ FFMPEG codec. It seems that in terms of quality the hierarchy is
46
+ 'rawvideo' = 'png' > 'mpeg4' > 'libx264'
47
+ 'png' manages the same lossless quality as 'rawvideo' but yields
48
+ smaller files. Type ``ffmpeg -codecs`` in a terminal to get a list
49
+ of accepted codecs.
50
+
51
+ Note for default 'libx264': by default the pixel format yuv420p
52
+ is used. If the video dimensions are not both even (e.g. 720x405)
53
+ another pixel format is used, and this can cause problem in some
54
+ video readers.
55
+
56
+ audiofile
57
+ Optional: The name of an audio file that will be incorporated
58
+ to the video.
59
+
60
+ preset
61
+ Sets the time that FFMPEG will take to compress the video. The slower,
62
+ the better the compression rate. Possibilities are: ultrafast,superfast,
63
+ veryfast, faster, fast, medium (default), slow, slower, veryslow,
64
+ placebo.
65
+
66
+ bitrate
67
+ Only relevant for codecs which accept a bitrate. "5000k" offers
68
+ nice results in general.
69
+
70
+ """
71
+
72
+ def __init__(self, filename, size, fps, codec="libx265", crf=14, audiofile=None,
73
+ preset="medium", bitrate=None,
74
+ logfile=None, threads=None, ffmpeg_params=None):
75
+
76
+ if logfile is None:
77
+ logfile = sp.PIPE
78
+
79
+ self.filename = filename
80
+ self.codec = codec
81
+ self.ext = self.filename.split(".")[-1]
82
+ w = size[0] - 1 if size[0] % 2 != 0 else size[0]
83
+ h = size[1] - 1 if size[1] % 2 != 0 else size[1]
84
+
85
+
86
+ # order is important
87
+ cmd = [
88
+ FFMPEG_BINARY,
89
+ '-hide_banner',
90
+ '-hwaccel', 'auto',
91
+ '-y',
92
+ '-loglevel', 'error' if logfile == sp.PIPE else 'info',
93
+ '-f', 'rawvideo',
94
+ '-vcodec', 'rawvideo',
95
+ '-s', '%dx%d' % (size[0], size[1]),
96
+ #'-pix_fmt', 'rgba' if withmask else 'rgb24',
97
+ '-pix_fmt', 'bgr24',
98
+ '-r', str(fps),
99
+ '-an', '-i', '-'
100
+ ]
101
+
102
+ if audiofile is not None:
103
+ cmd.extend([
104
+ '-i', audiofile,
105
+ '-acodec', 'copy'
106
+ ])
107
+
108
+ cmd.extend([
109
+ '-vcodec', codec,
110
+ '-crf', str(crf)
111
+ #'-preset', preset,
112
+ ])
113
+ if ffmpeg_params is not None:
114
+ cmd.extend(ffmpeg_params)
115
+ if bitrate is not None:
116
+ cmd.extend([
117
+ '-b', bitrate
118
+ ])
119
+
120
+ # scale to a resolution divisible by 2 if not even
121
+ cmd.extend(['-vf', f'scale={w}:{h}' if w != size[0] or h != size[1] else 'colorspace=bt709:iall=bt601-6-625:fast=1'])
122
+
123
+ if threads is not None:
124
+ cmd.extend(["-threads", str(threads)])
125
+
126
+ cmd.extend([
127
+ '-pix_fmt', 'yuv420p',
128
+
129
+ ])
130
+ cmd.extend([
131
+ filename
132
+ ])
133
+
134
+ test = str(cmd)
135
+ print(test)
136
+
137
+ popen_params = {"stdout": DEVNULL,
138
+ "stderr": logfile,
139
+ "stdin": sp.PIPE}
140
+
141
+ # This was added so that no extra unwanted window opens on windows
142
+ # when the child process is created
143
+ if os.name == "nt":
144
+ popen_params["creationflags"] = 0x08000000 # CREATE_NO_WINDOW
145
+
146
+ self.proc = sp.Popen(cmd, **popen_params)
147
+
148
+
149
+ def write_frame(self, img_array):
150
+ """ Writes one frame in the file."""
151
+ try:
152
+ #if PY3:
153
+ self.proc.stdin.write(img_array.tobytes())
154
+ # else:
155
+ # self.proc.stdin.write(img_array.tostring())
156
+ except IOError as err:
157
+ _, ffmpeg_error = self.proc.communicate()
158
+ error = (str(err) + ("\n\nroop unleashed error: FFMPEG encountered "
159
+ "the following error while writing file %s:"
160
+ "\n\n %s" % (self.filename, str(ffmpeg_error))))
161
+
162
+ if b"Unknown encoder" in ffmpeg_error:
163
+
164
+ error = error+("\n\nThe video export "
165
+ "failed because FFMPEG didn't find the specified "
166
+ "codec for video encoding (%s). Please install "
167
+ "this codec or change the codec when calling "
168
+ "write_videofile. For instance:\n"
169
+ " >>> clip.write_videofile('myvid.webm', codec='libvpx')")%(self.codec)
170
+
171
+ elif b"incorrect codec parameters ?" in ffmpeg_error:
172
+
173
+ error = error+("\n\nThe video export "
174
+ "failed, possibly because the codec specified for "
175
+ "the video (%s) is not compatible with the given "
176
+ "extension (%s). Please specify a valid 'codec' "
177
+ "argument in write_videofile. This would be 'libx264' "
178
+ "or 'mpeg4' for mp4, 'libtheora' for ogv, 'libvpx for webm. "
179
+ "Another possible reason is that the audio codec was not "
180
+ "compatible with the video codec. For instance the video "
181
+ "extensions 'ogv' and 'webm' only allow 'libvorbis' (default) as a"
182
+ "video codec."
183
+ )%(self.codec, self.ext)
184
+
185
+ elif b"encoder setup failed" in ffmpeg_error:
186
+
187
+ error = error+("\n\nThe video export "
188
+ "failed, possibly because the bitrate you specified "
189
+ "was too high or too low for the video codec.")
190
+
191
+ elif b"Invalid encoder type" in ffmpeg_error:
192
+
193
+ error = error + ("\n\nThe video export failed because the codec "
194
+ "or file extension you provided is not a video")
195
+
196
+
197
+ raise IOError(error)
198
+
199
+ def close(self):
200
+ if self.proc:
201
+ self.proc.stdin.close()
202
+ if self.proc.stderr is not None:
203
+ self.proc.stderr.close()
204
+ self.proc.wait()
205
+
206
+ self.proc = None
207
+
208
+ # Support the Context Manager protocol, to ensure that resources are cleaned up.
209
+
210
+ def __enter__(self):
211
+ return self
212
+
213
+ def __exit__(self, exc_type, exc_value, traceback):
214
+ self.close()
215
+
216
+
217
+
218
+
roop/globals.py ADDED
@@ -0,0 +1,54 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from settings import Settings
2
+ from typing import List
3
+
4
+ source_path = None
5
+ target_path = None
6
+ output_path = None
7
+ target_folder_path = None
8
+
9
+ frame_processors: List[str] = []
10
+ keep_fps = None
11
+ keep_frames = None
12
+ autorotate_faces = None
13
+ vr_mode = None
14
+ skip_audio = None
15
+ wait_after_extraction = None
16
+ many_faces = None
17
+ use_batch = None
18
+ source_face_index = 0
19
+ target_face_index = 0
20
+ face_position = None
21
+ video_encoder = None
22
+ video_quality = None
23
+ max_memory = None
24
+ execution_providers: List[str] = []
25
+ execution_threads = None
26
+ headless = None
27
+ log_level = 'error'
28
+ selected_enhancer = None
29
+ subsample_size = 128
30
+ face_swap_mode = None
31
+ blend_ratio = 0.5
32
+ distance_threshold = 0.65
33
+ default_det_size = True
34
+
35
+ no_face_action = 0
36
+
37
+ processing = False
38
+
39
+ g_current_face_analysis = None
40
+ g_desired_face_analysis = None
41
+
42
+ FACE_ENHANCER = None
43
+
44
+ INPUT_FACESETS = []
45
+ TARGET_FACES = []
46
+
47
+
48
+ IMAGE_CHAIN_PROCESSOR = None
49
+ VIDEO_CHAIN_PROCESSOR = None
50
+ BATCH_IMAGE_CHAIN_PROCESSOR = None
51
+
52
+ CFG: Settings = None
53
+
54
+
roop/metadata.py ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ name = 'roop unleashed'
2
+ version = '4.1.1'
roop/processors/Enhance_CodeFormer.py ADDED
@@ -0,0 +1,75 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Any, List, Callable
2
+ import cv2
3
+ import numpy as np
4
+ import onnxruntime
5
+ import roop.globals
6
+
7
+ from roop.typing import Face, Frame, FaceSet
8
+ from roop.utilities import resolve_relative_path
9
+
10
+
11
+ # THREAD_LOCK = threading.Lock()
12
+
13
+
14
+ class Enhance_CodeFormer():
15
+ model_codeformer = None
16
+
17
+ plugin_options:dict = None
18
+
19
+ processorname = 'codeformer'
20
+ type = 'enhance'
21
+
22
+
23
+ def Initialize(self, plugin_options:dict):
24
+ if self.plugin_options is not None:
25
+ if self.plugin_options["devicename"] != plugin_options["devicename"]:
26
+ self.Release()
27
+
28
+ self.plugin_options = plugin_options
29
+ if self.model_codeformer is None:
30
+ # replace Mac mps with cpu for the moment
31
+ self.devicename = self.plugin_options["devicename"].replace('mps', 'cpu')
32
+ model_path = resolve_relative_path('../models/CodeFormer/CodeFormerv0.1.onnx')
33
+ self.model_codeformer = onnxruntime.InferenceSession(model_path, None, providers=roop.globals.execution_providers)
34
+ self.model_inputs = self.model_codeformer.get_inputs()
35
+ model_outputs = self.model_codeformer.get_outputs()
36
+ self.io_binding = self.model_codeformer.io_binding()
37
+ self.io_binding.bind_cpu_input(self.model_inputs[1].name, np.array([0.5]))
38
+ self.io_binding.bind_output(model_outputs[0].name, self.devicename)
39
+
40
+
41
+ def Run(self, source_faceset: FaceSet, target_face: Face, temp_frame: Frame) -> Frame:
42
+ input_size = temp_frame.shape[1]
43
+ # preprocess
44
+ temp_frame = cv2.resize(temp_frame, (512, 512), cv2.INTER_CUBIC)
45
+ temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_BGR2RGB)
46
+ temp_frame = temp_frame.astype('float32') / 255.0
47
+ temp_frame = (temp_frame - 0.5) / 0.5
48
+ temp_frame = np.expand_dims(temp_frame, axis=0).transpose(0, 3, 1, 2)
49
+
50
+ self.io_binding.bind_cpu_input(self.model_inputs[0].name, temp_frame.astype(np.float32))
51
+ self.model_codeformer.run_with_iobinding(self.io_binding)
52
+ ort_outs = self.io_binding.copy_outputs_to_cpu()
53
+ result = ort_outs[0][0]
54
+ del ort_outs
55
+
56
+ # post-process
57
+ result = result.transpose((1, 2, 0))
58
+
59
+ un_min = -1.0
60
+ un_max = 1.0
61
+ result = np.clip(result, un_min, un_max)
62
+ result = (result - un_min) / (un_max - un_min)
63
+
64
+ result = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
65
+ result = (result * 255.0).round()
66
+ scale_factor = int(result.shape[1] / input_size)
67
+ return result.astype(np.uint8), scale_factor
68
+
69
+
70
+ def Release(self):
71
+ del self.model_codeformer
72
+ self.model_codeformer = None
73
+ del self.io_binding
74
+ self.io_binding = None
75
+
roop/processors/Enhance_DMDNet.py ADDED
@@ -0,0 +1,898 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Any, List, Callable
2
+ import cv2
3
+ import numpy as np
4
+ import torch
5
+ import torch.nn as nn
6
+ import torch.nn.functional as F
7
+ import torch.nn.utils.spectral_norm as SpectralNorm
8
+ import threading
9
+ from torchvision.ops import roi_align
10
+
11
+ from math import sqrt
12
+
13
+ from torchvision.transforms.functional import normalize
14
+
15
+ from roop.typing import Face, Frame, FaceSet
16
+
17
+
18
+ THREAD_LOCK_DMDNET = threading.Lock()
19
+
20
+
21
+ class Enhance_DMDNet():
22
+ plugin_options:dict = None
23
+ model_dmdnet = None
24
+ torchdevice = None
25
+
26
+ processorname = 'dmdnet'
27
+ type = 'enhance'
28
+
29
+
30
+ def Initialize(self, plugin_options:dict):
31
+ if self.plugin_options is not None:
32
+ if self.plugin_options["devicename"] != plugin_options["devicename"]:
33
+ self.Release()
34
+
35
+ self.plugin_options = plugin_options
36
+ if self.model_dmdnet is None:
37
+ self.model_dmdnet = self.create(self.plugin_options["devicename"])
38
+
39
+
40
+ # temp_frame already cropped+aligned, bbox not
41
+ def Run(self, source_faceset: FaceSet, target_face: Face, temp_frame: Frame) -> Frame:
42
+ input_size = temp_frame.shape[1]
43
+
44
+ result = self.enhance_face(source_faceset, temp_frame, target_face)
45
+ scale_factor = int(result.shape[1] / input_size)
46
+ return result.astype(np.uint8), scale_factor
47
+
48
+
49
+ def Release(self):
50
+ self.model_gfpgan = None
51
+
52
+
53
+ # https://stackoverflow.com/a/67174339
54
+ def landmarks106_to_68(self, pt106):
55
+ map106to68=[1,10,12,14,16,3,5,7,0,23,21,19,32,30,28,26,17,
56
+ 43,48,49,51,50,
57
+ 102,103,104,105,101,
58
+ 72,73,74,86,78,79,80,85,84,
59
+ 35,41,42,39,37,36,
60
+ 89,95,96,93,91,90,
61
+ 52,64,63,71,67,68,61,58,59,53,56,55,65,66,62,70,69,57,60,54
62
+ ]
63
+
64
+ pt68 = []
65
+ for i in range(68):
66
+ index = map106to68[i]
67
+ pt68.append(pt106[index])
68
+ return pt68
69
+
70
+
71
+
72
+
73
+ def check_bbox(self, imgs, boxes):
74
+ boxes = boxes.view(-1, 4, 4)
75
+ colors = [(0, 255, 0), (0, 255, 0), (255, 255, 0), (255, 0, 0)]
76
+ i = 0
77
+ for img, box in zip(imgs, boxes):
78
+ img = (img + 1)/2 * 255
79
+ img2 = img.permute(1, 2, 0).float().cpu().flip(2).numpy().copy()
80
+ for idx, point in enumerate(box):
81
+ cv2.rectangle(img2, (int(point[0]), int(point[1])), (int(point[2]), int(point[3])), color=colors[idx], thickness=2)
82
+ cv2.imwrite('dmdnet_{:02d}.png'.format(i), img2)
83
+ i += 1
84
+
85
+
86
+ def trans_points2d(self, pts, M):
87
+ new_pts = np.zeros(shape=pts.shape, dtype=np.float32)
88
+ for i in range(pts.shape[0]):
89
+ pt = pts[i]
90
+ new_pt = np.array([pt[0], pt[1], 1.0], dtype=np.float32)
91
+ new_pt = np.dot(M, new_pt)
92
+ new_pts[i] = new_pt[0:2]
93
+
94
+ return new_pts
95
+
96
+
97
+ def enhance_face(self, ref_faceset: FaceSet, temp_frame, face: Face):
98
+ # preprocess
99
+ start_x, start_y, end_x, end_y = map(int, face['bbox'])
100
+ lm106 = face.landmark_2d_106
101
+ lq_landmarks = np.asarray(self.landmarks106_to_68(lm106))
102
+
103
+ if temp_frame.shape[0] != 512 or temp_frame.shape[1] != 512:
104
+ # scale to 512x512
105
+ scale_factor = 512 / temp_frame.shape[1]
106
+
107
+ M = face.matrix * scale_factor
108
+
109
+ lq_landmarks = self.trans_points2d(lq_landmarks, M)
110
+ temp_frame = cv2.resize(temp_frame, (512,512), interpolation = cv2.INTER_AREA)
111
+
112
+ if temp_frame.ndim == 2:
113
+ temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_GRAY2RGB) # GGG
114
+ # else:
115
+ # temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_BGR2RGB) # RGB
116
+
117
+ lq = read_img_tensor(temp_frame)
118
+
119
+ LQLocs = get_component_location(lq_landmarks)
120
+ # self.check_bbox(lq, LQLocs.unsqueeze(0))
121
+
122
+ # specific, change 1000 to 1 to activate
123
+ if len(ref_faceset.faces) > 1:
124
+ SpecificImgs = []
125
+ SpecificLocs = []
126
+ for i,face in enumerate(ref_faceset.faces):
127
+ lm106 = face.landmark_2d_106
128
+ lq_landmarks = np.asarray(self.landmarks106_to_68(lm106))
129
+ ref_image = ref_faceset.ref_images[i]
130
+ if ref_image.shape[0] != 512 or ref_image.shape[1] != 512:
131
+ # scale to 512x512
132
+ scale_factor = 512 / ref_image.shape[1]
133
+
134
+ M = face.matrix * scale_factor
135
+
136
+ lq_landmarks = self.trans_points2d(lq_landmarks, M)
137
+ ref_image = cv2.resize(ref_image, (512,512), interpolation = cv2.INTER_AREA)
138
+
139
+ if ref_image.ndim == 2:
140
+ temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_GRAY2RGB) # GGG
141
+ # else:
142
+ # temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_BGR2RGB) # RGB
143
+
144
+ ref_tensor = read_img_tensor(ref_image)
145
+ ref_locs = get_component_location(lq_landmarks)
146
+ # self.check_bbox(ref_tensor, ref_locs.unsqueeze(0))
147
+
148
+ SpecificImgs.append(ref_tensor)
149
+ SpecificLocs.append(ref_locs.unsqueeze(0))
150
+
151
+ SpecificImgs = torch.cat(SpecificImgs, dim=0)
152
+ SpecificLocs = torch.cat(SpecificLocs, dim=0)
153
+ # check_bbox(SpecificImgs, SpecificLocs)
154
+ SpMem256, SpMem128, SpMem64 = self.model_dmdnet.generate_specific_dictionary(sp_imgs = SpecificImgs.to(self.torchdevice), sp_locs = SpecificLocs)
155
+ SpMem256Para = {}
156
+ SpMem128Para = {}
157
+ SpMem64Para = {}
158
+ for k, v in SpMem256.items():
159
+ SpMem256Para[k] = v
160
+ for k, v in SpMem128.items():
161
+ SpMem128Para[k] = v
162
+ for k, v in SpMem64.items():
163
+ SpMem64Para[k] = v
164
+ else:
165
+ # generic
166
+ SpMem256Para, SpMem128Para, SpMem64Para = None, None, None
167
+
168
+ with torch.no_grad():
169
+ with THREAD_LOCK_DMDNET:
170
+ try:
171
+ GenericResult, SpecificResult = self.model_dmdnet(lq = lq.to(self.torchdevice), loc = LQLocs.unsqueeze(0), sp_256 = SpMem256Para, sp_128 = SpMem128Para, sp_64 = SpMem64Para)
172
+ except Exception as e:
173
+ print(f'Error {e} there may be something wrong with the detected component locations.')
174
+ return temp_frame
175
+
176
+ if SpecificResult is not None:
177
+ save_specific = SpecificResult * 0.5 + 0.5
178
+ save_specific = save_specific.squeeze(0).permute(1, 2, 0).flip(2) # RGB->BGR
179
+ save_specific = np.clip(save_specific.float().cpu().numpy(), 0, 1) * 255.0
180
+ temp_frame = save_specific.astype("uint8")
181
+ if False:
182
+ save_generic = GenericResult * 0.5 + 0.5
183
+ save_generic = save_generic.squeeze(0).permute(1, 2, 0).flip(2) # RGB->BGR
184
+ save_generic = np.clip(save_generic.float().cpu().numpy(), 0, 1) * 255.0
185
+ check_lq = lq * 0.5 + 0.5
186
+ check_lq = check_lq.squeeze(0).permute(1, 2, 0).flip(2) # RGB->BGR
187
+ check_lq = np.clip(check_lq.float().cpu().numpy(), 0, 1) * 255.0
188
+ cv2.imwrite('dmdnet_comparison.png', cv2.cvtColor(np.hstack((check_lq, save_generic, save_specific)),cv2.COLOR_RGB2BGR))
189
+ else:
190
+ save_generic = GenericResult * 0.5 + 0.5
191
+ save_generic = save_generic.squeeze(0).permute(1, 2, 0).flip(2) # RGB->BGR
192
+ save_generic = np.clip(save_generic.float().cpu().numpy(), 0, 1) * 255.0
193
+ temp_frame = save_generic.astype("uint8")
194
+ temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_RGB2BGR) # RGB
195
+ return temp_frame
196
+
197
+
198
+
199
+ def create(self, devicename):
200
+ self.torchdevice = torch.device(devicename)
201
+ model_dmdnet = DMDNet().to(self.torchdevice)
202
+ weights = torch.load('./models/DMDNet.pth')
203
+ model_dmdnet.load_state_dict(weights, strict=True)
204
+
205
+ model_dmdnet.eval()
206
+ num_params = 0
207
+ for param in model_dmdnet.parameters():
208
+ num_params += param.numel()
209
+ return model_dmdnet
210
+
211
+ # print('{:>8s} : {}'.format('Using device', device))
212
+ # print('{:>8s} : {:.2f}M'.format('Model params', num_params/1e6))
213
+
214
+
215
+
216
+ def read_img_tensor(Img=None): #rgb -1~1
217
+ Img = Img.transpose((2, 0, 1))/255.0
218
+ Img = torch.from_numpy(Img).float()
219
+ normalize(Img, [0.5,0.5,0.5], [0.5,0.5,0.5], inplace=True)
220
+ ImgTensor = Img.unsqueeze(0)
221
+ return ImgTensor
222
+
223
+
224
+ def get_component_location(Landmarks, re_read=False):
225
+ if re_read:
226
+ ReadLandmark = []
227
+ with open(Landmarks,'r') as f:
228
+ for line in f:
229
+ tmp = [float(i) for i in line.split(' ') if i != '\n']
230
+ ReadLandmark.append(tmp)
231
+ ReadLandmark = np.array(ReadLandmark) #
232
+ Landmarks = np.reshape(ReadLandmark, [-1, 2]) # 68*2
233
+ Map_LE_B = list(np.hstack((range(17,22), range(36,42))))
234
+ Map_RE_B = list(np.hstack((range(22,27), range(42,48))))
235
+ Map_LE = list(range(36,42))
236
+ Map_RE = list(range(42,48))
237
+ Map_NO = list(range(29,36))
238
+ Map_MO = list(range(48,68))
239
+
240
+ Landmarks[Landmarks>504]=504
241
+ Landmarks[Landmarks<8]=8
242
+
243
+ #left eye
244
+ Mean_LE = np.mean(Landmarks[Map_LE],0)
245
+ L_LE1 = Mean_LE[1] - np.min(Landmarks[Map_LE_B,1])
246
+ L_LE1 = L_LE1 * 1.3
247
+ L_LE2 = L_LE1 / 1.9
248
+ L_LE_xy = L_LE1 + L_LE2
249
+ L_LE_lt = [L_LE_xy/2, L_LE1]
250
+ L_LE_rb = [L_LE_xy/2, L_LE2]
251
+ Location_LE = np.hstack((Mean_LE - L_LE_lt + 1, Mean_LE + L_LE_rb)).astype(int)
252
+
253
+ #right eye
254
+ Mean_RE = np.mean(Landmarks[Map_RE],0)
255
+ L_RE1 = Mean_RE[1] - np.min(Landmarks[Map_RE_B,1])
256
+ L_RE1 = L_RE1 * 1.3
257
+ L_RE2 = L_RE1 / 1.9
258
+ L_RE_xy = L_RE1 + L_RE2
259
+ L_RE_lt = [L_RE_xy/2, L_RE1]
260
+ L_RE_rb = [L_RE_xy/2, L_RE2]
261
+ Location_RE = np.hstack((Mean_RE - L_RE_lt + 1, Mean_RE + L_RE_rb)).astype(int)
262
+
263
+ #nose
264
+ Mean_NO = np.mean(Landmarks[Map_NO],0)
265
+ L_NO1 =( np.max([Mean_NO[0] - Landmarks[31][0], Landmarks[35][0] - Mean_NO[0]])) * 1.25
266
+ L_NO2 = (Landmarks[33][1] - Mean_NO[1]) * 1.1
267
+ L_NO_xy = L_NO1 * 2
268
+ L_NO_lt = [L_NO_xy/2, L_NO_xy - L_NO2]
269
+ L_NO_rb = [L_NO_xy/2, L_NO2]
270
+ Location_NO = np.hstack((Mean_NO - L_NO_lt + 1, Mean_NO + L_NO_rb)).astype(int)
271
+
272
+ #mouth
273
+ Mean_MO = np.mean(Landmarks[Map_MO],0)
274
+ L_MO = np.max((np.max(np.max(Landmarks[Map_MO],0) - np.min(Landmarks[Map_MO],0))/2,16)) * 1.1
275
+ MO_O = Mean_MO - L_MO + 1
276
+ MO_T = Mean_MO + L_MO
277
+ MO_T[MO_T>510]=510
278
+ Location_MO = np.hstack((MO_O, MO_T)).astype(int)
279
+ return torch.cat([torch.FloatTensor(Location_LE).unsqueeze(0), torch.FloatTensor(Location_RE).unsqueeze(0), torch.FloatTensor(Location_NO).unsqueeze(0), torch.FloatTensor(Location_MO).unsqueeze(0)], dim=0)
280
+
281
+
282
+
283
+
284
+ def calc_mean_std_4D(feat, eps=1e-5):
285
+ # eps is a small value added to the variance to avoid divide-by-zero.
286
+ size = feat.size()
287
+ assert (len(size) == 4)
288
+ N, C = size[:2]
289
+ feat_var = feat.view(N, C, -1).var(dim=2) + eps
290
+ feat_std = feat_var.sqrt().view(N, C, 1, 1)
291
+ feat_mean = feat.view(N, C, -1).mean(dim=2).view(N, C, 1, 1)
292
+ return feat_mean, feat_std
293
+
294
+ def adaptive_instance_normalization_4D(content_feat, style_feat): # content_feat is ref feature, style is degradate feature
295
+ size = content_feat.size()
296
+ style_mean, style_std = calc_mean_std_4D(style_feat)
297
+ content_mean, content_std = calc_mean_std_4D(content_feat)
298
+ normalized_feat = (content_feat - content_mean.expand(size)) / content_std.expand(size)
299
+ return normalized_feat * style_std.expand(size) + style_mean.expand(size)
300
+
301
+
302
+ def convU(in_channels, out_channels,conv_layer, norm_layer, kernel_size=3, stride=1,dilation=1, bias=True):
303
+ return nn.Sequential(
304
+ SpectralNorm(conv_layer(in_channels, out_channels, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=((kernel_size-1)//2)*dilation, bias=bias)),
305
+ nn.LeakyReLU(0.2),
306
+ SpectralNorm(conv_layer(out_channels, out_channels, kernel_size=kernel_size, stride=stride, dilation=dilation, padding=((kernel_size-1)//2)*dilation, bias=bias)),
307
+ )
308
+
309
+
310
+ class MSDilateBlock(nn.Module):
311
+ def __init__(self, in_channels,conv_layer=nn.Conv2d, norm_layer=nn.BatchNorm2d, kernel_size=3, dilation=[1,1,1,1], bias=True):
312
+ super(MSDilateBlock, self).__init__()
313
+ self.conv1 = convU(in_channels, in_channels,conv_layer, norm_layer, kernel_size,dilation=dilation[0], bias=bias)
314
+ self.conv2 = convU(in_channels, in_channels,conv_layer, norm_layer, kernel_size,dilation=dilation[1], bias=bias)
315
+ self.conv3 = convU(in_channels, in_channels,conv_layer, norm_layer, kernel_size,dilation=dilation[2], bias=bias)
316
+ self.conv4 = convU(in_channels, in_channels,conv_layer, norm_layer, kernel_size,dilation=dilation[3], bias=bias)
317
+ self.convi = SpectralNorm(conv_layer(in_channels*4, in_channels, kernel_size=kernel_size, stride=1, padding=(kernel_size-1)//2, bias=bias))
318
+ def forward(self, x):
319
+ conv1 = self.conv1(x)
320
+ conv2 = self.conv2(x)
321
+ conv3 = self.conv3(x)
322
+ conv4 = self.conv4(x)
323
+ cat = torch.cat([conv1, conv2, conv3, conv4], 1)
324
+ out = self.convi(cat) + x
325
+ return out
326
+
327
+
328
+ class AdaptiveInstanceNorm(nn.Module):
329
+ def __init__(self, in_channel):
330
+ super().__init__()
331
+ self.norm = nn.InstanceNorm2d(in_channel)
332
+
333
+ def forward(self, input, style):
334
+ style_mean, style_std = calc_mean_std_4D(style)
335
+ out = self.norm(input)
336
+ size = input.size()
337
+ out = style_std.expand(size) * out + style_mean.expand(size)
338
+ return out
339
+
340
+ class NoiseInjection(nn.Module):
341
+ def __init__(self, channel):
342
+ super().__init__()
343
+ self.weight = nn.Parameter(torch.zeros(1, channel, 1, 1))
344
+ def forward(self, image, noise):
345
+ if noise is None:
346
+ b, c, h, w = image.shape
347
+ noise = image.new_empty(b, 1, h, w).normal_()
348
+ return image + self.weight * noise
349
+
350
+ class StyledUpBlock(nn.Module):
351
+ def __init__(self, in_channel, out_channel, kernel_size=3, padding=1,upsample=False, noise_inject=False):
352
+ super().__init__()
353
+
354
+ self.noise_inject = noise_inject
355
+ if upsample:
356
+ self.conv1 = nn.Sequential(
357
+ nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
358
+ SpectralNorm(nn.Conv2d(in_channel, out_channel, kernel_size, padding=padding)),
359
+ nn.LeakyReLU(0.2),
360
+ )
361
+ else:
362
+ self.conv1 = nn.Sequential(
363
+ SpectralNorm(nn.Conv2d(in_channel, out_channel, kernel_size, padding=padding)),
364
+ nn.LeakyReLU(0.2),
365
+ SpectralNorm(nn.Conv2d(out_channel, out_channel, kernel_size, padding=padding)),
366
+ )
367
+ self.convup = nn.Sequential(
368
+ nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False),
369
+ SpectralNorm(nn.Conv2d(out_channel, out_channel, kernel_size, padding=padding)),
370
+ nn.LeakyReLU(0.2),
371
+ SpectralNorm(nn.Conv2d(out_channel, out_channel, kernel_size, padding=padding)),
372
+ )
373
+ if self.noise_inject:
374
+ self.noise1 = NoiseInjection(out_channel)
375
+
376
+ self.lrelu1 = nn.LeakyReLU(0.2)
377
+
378
+ self.ScaleModel1 = nn.Sequential(
379
+ SpectralNorm(nn.Conv2d(in_channel,out_channel,3, 1, 1)),
380
+ nn.LeakyReLU(0.2),
381
+ SpectralNorm(nn.Conv2d(out_channel, out_channel, 3, 1, 1))
382
+ )
383
+ self.ShiftModel1 = nn.Sequential(
384
+ SpectralNorm(nn.Conv2d(in_channel,out_channel,3, 1, 1)),
385
+ nn.LeakyReLU(0.2),
386
+ SpectralNorm(nn.Conv2d(out_channel, out_channel, 3, 1, 1)),
387
+ )
388
+
389
+ def forward(self, input, style):
390
+ out = self.conv1(input)
391
+ out = self.lrelu1(out)
392
+ Shift1 = self.ShiftModel1(style)
393
+ Scale1 = self.ScaleModel1(style)
394
+ out = out * Scale1 + Shift1
395
+ if self.noise_inject:
396
+ out = self.noise1(out, noise=None)
397
+ outup = self.convup(out)
398
+ return outup
399
+
400
+
401
+ ####################################################################
402
+ ###############Face Dictionary Generator
403
+ ####################################################################
404
+ def AttentionBlock(in_channel):
405
+ return nn.Sequential(
406
+ SpectralNorm(nn.Conv2d(in_channel, in_channel, 3, 1, 1)),
407
+ nn.LeakyReLU(0.2),
408
+ SpectralNorm(nn.Conv2d(in_channel, in_channel, 3, 1, 1)),
409
+ )
410
+
411
+ class DilateResBlock(nn.Module):
412
+ def __init__(self, dim, dilation=[5,3] ):
413
+ super(DilateResBlock, self).__init__()
414
+ self.Res = nn.Sequential(
415
+ SpectralNorm(nn.Conv2d(dim, dim, 3, 1, ((3-1)//2)*dilation[0], dilation[0])),
416
+ nn.LeakyReLU(0.2),
417
+ SpectralNorm(nn.Conv2d(dim, dim, 3, 1, ((3-1)//2)*dilation[1], dilation[1])),
418
+ )
419
+ def forward(self, x):
420
+ out = x + self.Res(x)
421
+ return out
422
+
423
+
424
+ class KeyValue(nn.Module):
425
+ def __init__(self, indim, keydim, valdim):
426
+ super(KeyValue, self).__init__()
427
+ self.Key = nn.Sequential(
428
+ SpectralNorm(nn.Conv2d(indim, keydim, kernel_size=(3,3), padding=(1,1), stride=1)),
429
+ nn.LeakyReLU(0.2),
430
+ SpectralNorm(nn.Conv2d(keydim, keydim, kernel_size=(3,3), padding=(1,1), stride=1)),
431
+ )
432
+ self.Value = nn.Sequential(
433
+ SpectralNorm(nn.Conv2d(indim, valdim, kernel_size=(3,3), padding=(1,1), stride=1)),
434
+ nn.LeakyReLU(0.2),
435
+ SpectralNorm(nn.Conv2d(valdim, valdim, kernel_size=(3,3), padding=(1,1), stride=1)),
436
+ )
437
+ def forward(self, x):
438
+ return self.Key(x), self.Value(x)
439
+
440
+ class MaskAttention(nn.Module):
441
+ def __init__(self, indim):
442
+ super(MaskAttention, self).__init__()
443
+ self.conv1 = nn.Sequential(
444
+ SpectralNorm(nn.Conv2d(indim, indim//3, kernel_size=(3,3), padding=(1,1), stride=1)),
445
+ nn.LeakyReLU(0.2),
446
+ SpectralNorm(nn.Conv2d(indim//3, indim//3, kernel_size=(3,3), padding=(1,1), stride=1)),
447
+ )
448
+ self.conv2 = nn.Sequential(
449
+ SpectralNorm(nn.Conv2d(indim, indim//3, kernel_size=(3,3), padding=(1,1), stride=1)),
450
+ nn.LeakyReLU(0.2),
451
+ SpectralNorm(nn.Conv2d(indim//3, indim//3, kernel_size=(3,3), padding=(1,1), stride=1)),
452
+ )
453
+ self.conv3 = nn.Sequential(
454
+ SpectralNorm(nn.Conv2d(indim, indim//3, kernel_size=(3,3), padding=(1,1), stride=1)),
455
+ nn.LeakyReLU(0.2),
456
+ SpectralNorm(nn.Conv2d(indim//3, indim//3, kernel_size=(3,3), padding=(1,1), stride=1)),
457
+ )
458
+ self.convCat = nn.Sequential(
459
+ SpectralNorm(nn.Conv2d(indim//3 * 3, indim, kernel_size=(3,3), padding=(1,1), stride=1)),
460
+ nn.LeakyReLU(0.2),
461
+ SpectralNorm(nn.Conv2d(indim, indim, kernel_size=(3,3), padding=(1,1), stride=1)),
462
+ )
463
+ def forward(self, x, y, z):
464
+ c1 = self.conv1(x)
465
+ c2 = self.conv2(y)
466
+ c3 = self.conv3(z)
467
+ return self.convCat(torch.cat([c1,c2,c3], dim=1))
468
+
469
+ class Query(nn.Module):
470
+ def __init__(self, indim, quedim):
471
+ super(Query, self).__init__()
472
+ self.Query = nn.Sequential(
473
+ SpectralNorm(nn.Conv2d(indim, quedim, kernel_size=(3,3), padding=(1,1), stride=1)),
474
+ nn.LeakyReLU(0.2),
475
+ SpectralNorm(nn.Conv2d(quedim, quedim, kernel_size=(3,3), padding=(1,1), stride=1)),
476
+ )
477
+ def forward(self, x):
478
+ return self.Query(x)
479
+
480
+ def roi_align_self(input, location, target_size):
481
+ test = (target_size.item(),target_size.item())
482
+ return torch.cat([F.interpolate(input[i:i+1,:,location[i,1]:location[i,3],location[i,0]:location[i,2]],test,mode='bilinear',align_corners=False) for i in range(input.size(0))],0)
483
+
484
+ class FeatureExtractor(nn.Module):
485
+ def __init__(self, ngf = 64, key_scale = 4):#
486
+ super().__init__()
487
+
488
+ self.key_scale = 4
489
+ self.part_sizes = np.array([80,80,50,110]) #
490
+ self.feature_sizes = np.array([256,128,64]) #
491
+
492
+ self.conv1 = nn.Sequential(
493
+ SpectralNorm(nn.Conv2d(3, ngf, 3, 2, 1)),
494
+ nn.LeakyReLU(0.2),
495
+ SpectralNorm(nn.Conv2d(ngf, ngf, 3, 1, 1)),
496
+ )
497
+ self.conv2 = nn.Sequential(
498
+ SpectralNorm(nn.Conv2d(ngf, ngf, 3, 1, 1)),
499
+ nn.LeakyReLU(0.2),
500
+ SpectralNorm(nn.Conv2d(ngf, ngf, 3, 1, 1))
501
+ )
502
+ self.res1 = DilateResBlock(ngf, [5,3])
503
+ self.res2 = DilateResBlock(ngf, [5,3])
504
+
505
+
506
+ self.conv3 = nn.Sequential(
507
+ SpectralNorm(nn.Conv2d(ngf, ngf*2, 3, 2, 1)),
508
+ nn.LeakyReLU(0.2),
509
+ SpectralNorm(nn.Conv2d(ngf*2, ngf*2, 3, 1, 1)),
510
+ )
511
+ self.conv4 = nn.Sequential(
512
+ SpectralNorm(nn.Conv2d(ngf*2, ngf*2, 3, 1, 1)),
513
+ nn.LeakyReLU(0.2),
514
+ SpectralNorm(nn.Conv2d(ngf*2, ngf*2, 3, 1, 1))
515
+ )
516
+ self.res3 = DilateResBlock(ngf*2, [3,1])
517
+ self.res4 = DilateResBlock(ngf*2, [3,1])
518
+
519
+ self.conv5 = nn.Sequential(
520
+ SpectralNorm(nn.Conv2d(ngf*2, ngf*4, 3, 2, 1)),
521
+ nn.LeakyReLU(0.2),
522
+ SpectralNorm(nn.Conv2d(ngf*4, ngf*4, 3, 1, 1)),
523
+ )
524
+ self.conv6 = nn.Sequential(
525
+ SpectralNorm(nn.Conv2d(ngf*4, ngf*4, 3, 1, 1)),
526
+ nn.LeakyReLU(0.2),
527
+ SpectralNorm(nn.Conv2d(ngf*4, ngf*4, 3, 1, 1))
528
+ )
529
+ self.res5 = DilateResBlock(ngf*4, [1,1])
530
+ self.res6 = DilateResBlock(ngf*4, [1,1])
531
+
532
+ self.LE_256_Q = Query(ngf, ngf // self.key_scale)
533
+ self.RE_256_Q = Query(ngf, ngf // self.key_scale)
534
+ self.MO_256_Q = Query(ngf, ngf // self.key_scale)
535
+ self.LE_128_Q = Query(ngf * 2, ngf * 2 // self.key_scale)
536
+ self.RE_128_Q = Query(ngf * 2, ngf * 2 // self.key_scale)
537
+ self.MO_128_Q = Query(ngf * 2, ngf * 2 // self.key_scale)
538
+ self.LE_64_Q = Query(ngf * 4, ngf * 4 // self.key_scale)
539
+ self.RE_64_Q = Query(ngf * 4, ngf * 4 // self.key_scale)
540
+ self.MO_64_Q = Query(ngf * 4, ngf * 4 // self.key_scale)
541
+
542
+
543
+ def forward(self, img, locs):
544
+ le_location = locs[:,0,:].int().cpu().numpy()
545
+ re_location = locs[:,1,:].int().cpu().numpy()
546
+ no_location = locs[:,2,:].int().cpu().numpy()
547
+ mo_location = locs[:,3,:].int().cpu().numpy()
548
+
549
+
550
+ f1_0 = self.conv1(img)
551
+ f1_1 = self.res1(f1_0)
552
+ f2_0 = self.conv2(f1_1)
553
+ f2_1 = self.res2(f2_0)
554
+
555
+ f3_0 = self.conv3(f2_1)
556
+ f3_1 = self.res3(f3_0)
557
+ f4_0 = self.conv4(f3_1)
558
+ f4_1 = self.res4(f4_0)
559
+
560
+ f5_0 = self.conv5(f4_1)
561
+ f5_1 = self.res5(f5_0)
562
+ f6_0 = self.conv6(f5_1)
563
+ f6_1 = self.res6(f6_0)
564
+
565
+
566
+ ####ROI Align
567
+ le_part_256 = roi_align_self(f2_1.clone(), le_location//2, self.part_sizes[0]//2)
568
+ re_part_256 = roi_align_self(f2_1.clone(), re_location//2, self.part_sizes[1]//2)
569
+ mo_part_256 = roi_align_self(f2_1.clone(), mo_location//2, self.part_sizes[3]//2)
570
+
571
+ le_part_128 = roi_align_self(f4_1.clone(), le_location//4, self.part_sizes[0]//4)
572
+ re_part_128 = roi_align_self(f4_1.clone(), re_location//4, self.part_sizes[1]//4)
573
+ mo_part_128 = roi_align_self(f4_1.clone(), mo_location//4, self.part_sizes[3]//4)
574
+
575
+ le_part_64 = roi_align_self(f6_1.clone(), le_location//8, self.part_sizes[0]//8)
576
+ re_part_64 = roi_align_self(f6_1.clone(), re_location//8, self.part_sizes[1]//8)
577
+ mo_part_64 = roi_align_self(f6_1.clone(), mo_location//8, self.part_sizes[3]//8)
578
+
579
+
580
+ le_256_q = self.LE_256_Q(le_part_256)
581
+ re_256_q = self.RE_256_Q(re_part_256)
582
+ mo_256_q = self.MO_256_Q(mo_part_256)
583
+
584
+ le_128_q = self.LE_128_Q(le_part_128)
585
+ re_128_q = self.RE_128_Q(re_part_128)
586
+ mo_128_q = self.MO_128_Q(mo_part_128)
587
+
588
+ le_64_q = self.LE_64_Q(le_part_64)
589
+ re_64_q = self.RE_64_Q(re_part_64)
590
+ mo_64_q = self.MO_64_Q(mo_part_64)
591
+
592
+ return {'f256': f2_1, 'f128': f4_1, 'f64': f6_1,\
593
+ 'le256': le_part_256, 're256': re_part_256, 'mo256': mo_part_256, \
594
+ 'le128': le_part_128, 're128': re_part_128, 'mo128': mo_part_128, \
595
+ 'le64': le_part_64, 're64': re_part_64, 'mo64': mo_part_64, \
596
+ 'le_256_q': le_256_q, 're_256_q': re_256_q, 'mo_256_q': mo_256_q,\
597
+ 'le_128_q': le_128_q, 're_128_q': re_128_q, 'mo_128_q': mo_128_q,\
598
+ 'le_64_q': le_64_q, 're_64_q': re_64_q, 'mo_64_q': mo_64_q}
599
+
600
+
601
+ class DMDNet(nn.Module):
602
+ def __init__(self, ngf = 64, banks_num = 128):
603
+ super().__init__()
604
+ self.part_sizes = np.array([80,80,50,110]) # size for 512
605
+ self.feature_sizes = np.array([256,128,64]) # size for 512
606
+
607
+ self.banks_num = banks_num
608
+ self.key_scale = 4
609
+
610
+ self.E_lq = FeatureExtractor(key_scale = self.key_scale)
611
+ self.E_hq = FeatureExtractor(key_scale = self.key_scale)
612
+
613
+ self.LE_256_KV = KeyValue(ngf, ngf // self.key_scale, ngf)
614
+ self.RE_256_KV = KeyValue(ngf, ngf // self.key_scale, ngf)
615
+ self.MO_256_KV = KeyValue(ngf, ngf // self.key_scale, ngf)
616
+
617
+ self.LE_128_KV = KeyValue(ngf * 2 , ngf * 2 // self.key_scale, ngf * 2)
618
+ self.RE_128_KV = KeyValue(ngf * 2 , ngf * 2 // self.key_scale, ngf * 2)
619
+ self.MO_128_KV = KeyValue(ngf * 2 , ngf * 2 // self.key_scale, ngf * 2)
620
+
621
+ self.LE_64_KV = KeyValue(ngf * 4 , ngf * 4 // self.key_scale, ngf * 4)
622
+ self.RE_64_KV = KeyValue(ngf * 4 , ngf * 4 // self.key_scale, ngf * 4)
623
+ self.MO_64_KV = KeyValue(ngf * 4 , ngf * 4 // self.key_scale, ngf * 4)
624
+
625
+
626
+ self.LE_256_Attention = AttentionBlock(64)
627
+ self.RE_256_Attention = AttentionBlock(64)
628
+ self.MO_256_Attention = AttentionBlock(64)
629
+
630
+ self.LE_128_Attention = AttentionBlock(128)
631
+ self.RE_128_Attention = AttentionBlock(128)
632
+ self.MO_128_Attention = AttentionBlock(128)
633
+
634
+ self.LE_64_Attention = AttentionBlock(256)
635
+ self.RE_64_Attention = AttentionBlock(256)
636
+ self.MO_64_Attention = AttentionBlock(256)
637
+
638
+ self.LE_256_Mask = MaskAttention(64)
639
+ self.RE_256_Mask = MaskAttention(64)
640
+ self.MO_256_Mask = MaskAttention(64)
641
+
642
+ self.LE_128_Mask = MaskAttention(128)
643
+ self.RE_128_Mask = MaskAttention(128)
644
+ self.MO_128_Mask = MaskAttention(128)
645
+
646
+ self.LE_64_Mask = MaskAttention(256)
647
+ self.RE_64_Mask = MaskAttention(256)
648
+ self.MO_64_Mask = MaskAttention(256)
649
+
650
+ self.MSDilate = MSDilateBlock(ngf*4, dilation = [4,3,2,1])
651
+
652
+ self.up1 = StyledUpBlock(ngf*4, ngf*2, noise_inject=False) #
653
+ self.up2 = StyledUpBlock(ngf*2, ngf, noise_inject=False) #
654
+ self.up3 = StyledUpBlock(ngf, ngf, noise_inject=False) #
655
+ self.up4 = nn.Sequential(
656
+ SpectralNorm(nn.Conv2d(ngf, ngf, 3, 1, 1)),
657
+ nn.LeakyReLU(0.2),
658
+ UpResBlock(ngf),
659
+ UpResBlock(ngf),
660
+ SpectralNorm(nn.Conv2d(ngf, 3, kernel_size=3, stride=1, padding=1)),
661
+ nn.Tanh()
662
+ )
663
+
664
+ # define generic memory, revise register_buffer to register_parameter for backward update
665
+ self.register_buffer('le_256_mem_key', torch.randn(128,16,40,40))
666
+ self.register_buffer('re_256_mem_key', torch.randn(128,16,40,40))
667
+ self.register_buffer('mo_256_mem_key', torch.randn(128,16,55,55))
668
+ self.register_buffer('le_256_mem_value', torch.randn(128,64,40,40))
669
+ self.register_buffer('re_256_mem_value', torch.randn(128,64,40,40))
670
+ self.register_buffer('mo_256_mem_value', torch.randn(128,64,55,55))
671
+
672
+
673
+ self.register_buffer('le_128_mem_key', torch.randn(128,32,20,20))
674
+ self.register_buffer('re_128_mem_key', torch.randn(128,32,20,20))
675
+ self.register_buffer('mo_128_mem_key', torch.randn(128,32,27,27))
676
+ self.register_buffer('le_128_mem_value', torch.randn(128,128,20,20))
677
+ self.register_buffer('re_128_mem_value', torch.randn(128,128,20,20))
678
+ self.register_buffer('mo_128_mem_value', torch.randn(128,128,27,27))
679
+
680
+ self.register_buffer('le_64_mem_key', torch.randn(128,64,10,10))
681
+ self.register_buffer('re_64_mem_key', torch.randn(128,64,10,10))
682
+ self.register_buffer('mo_64_mem_key', torch.randn(128,64,13,13))
683
+ self.register_buffer('le_64_mem_value', torch.randn(128,256,10,10))
684
+ self.register_buffer('re_64_mem_value', torch.randn(128,256,10,10))
685
+ self.register_buffer('mo_64_mem_value', torch.randn(128,256,13,13))
686
+
687
+
688
+ def readMem(self, k, v, q):
689
+ sim = F.conv2d(q, k)
690
+ score = F.softmax(sim/sqrt(sim.size(1)), dim=1) #B * S * 1 * 1 6*128
691
+ sb,sn,sw,sh = score.size()
692
+ s_m = score.view(sb, -1).unsqueeze(1)#2*1*M
693
+ vb,vn,vw,vh = v.size()
694
+ v_in = v.view(vb, -1).repeat(sb,1,1)#2*M*(c*w*h)
695
+ mem_out = torch.bmm(s_m, v_in).squeeze(1).view(sb, vn, vw,vh)
696
+ max_inds = torch.argmax(score, dim=1).squeeze()
697
+ return mem_out, max_inds
698
+
699
+
700
+ def memorize(self, img, locs):
701
+ fs = self.E_hq(img, locs)
702
+ LE256_key, LE256_value = self.LE_256_KV(fs['le256'])
703
+ RE256_key, RE256_value = self.RE_256_KV(fs['re256'])
704
+ MO256_key, MO256_value = self.MO_256_KV(fs['mo256'])
705
+
706
+ LE128_key, LE128_value = self.LE_128_KV(fs['le128'])
707
+ RE128_key, RE128_value = self.RE_128_KV(fs['re128'])
708
+ MO128_key, MO128_value = self.MO_128_KV(fs['mo128'])
709
+
710
+ LE64_key, LE64_value = self.LE_64_KV(fs['le64'])
711
+ RE64_key, RE64_value = self.RE_64_KV(fs['re64'])
712
+ MO64_key, MO64_value = self.MO_64_KV(fs['mo64'])
713
+
714
+ Mem256 = {'LE256Key': LE256_key, 'LE256Value': LE256_value, 'RE256Key': RE256_key, 'RE256Value': RE256_value,'MO256Key': MO256_key, 'MO256Value': MO256_value}
715
+ Mem128 = {'LE128Key': LE128_key, 'LE128Value': LE128_value, 'RE128Key': RE128_key, 'RE128Value': RE128_value,'MO128Key': MO128_key, 'MO128Value': MO128_value}
716
+ Mem64 = {'LE64Key': LE64_key, 'LE64Value': LE64_value, 'RE64Key': RE64_key, 'RE64Value': RE64_value,'MO64Key': MO64_key, 'MO64Value': MO64_value}
717
+
718
+ FS256 = {'LE256F':fs['le256'], 'RE256F':fs['re256'], 'MO256F':fs['mo256']}
719
+ FS128 = {'LE128F':fs['le128'], 'RE128F':fs['re128'], 'MO128F':fs['mo128']}
720
+ FS64 = {'LE64F':fs['le64'], 'RE64F':fs['re64'], 'MO64F':fs['mo64']}
721
+
722
+ return Mem256, Mem128, Mem64
723
+
724
+ def enhancer(self, fs_in, sp_256=None, sp_128=None, sp_64=None):
725
+ le_256_q = fs_in['le_256_q']
726
+ re_256_q = fs_in['re_256_q']
727
+ mo_256_q = fs_in['mo_256_q']
728
+
729
+ le_128_q = fs_in['le_128_q']
730
+ re_128_q = fs_in['re_128_q']
731
+ mo_128_q = fs_in['mo_128_q']
732
+
733
+ le_64_q = fs_in['le_64_q']
734
+ re_64_q = fs_in['re_64_q']
735
+ mo_64_q = fs_in['mo_64_q']
736
+
737
+
738
+ ####for 256
739
+ le_256_mem_g, le_256_inds = self.readMem(self.le_256_mem_key, self.le_256_mem_value, le_256_q)
740
+ re_256_mem_g, re_256_inds = self.readMem(self.re_256_mem_key, self.re_256_mem_value, re_256_q)
741
+ mo_256_mem_g, mo_256_inds = self.readMem(self.mo_256_mem_key, self.mo_256_mem_value, mo_256_q)
742
+
743
+ le_128_mem_g, le_128_inds = self.readMem(self.le_128_mem_key, self.le_128_mem_value, le_128_q)
744
+ re_128_mem_g, re_128_inds = self.readMem(self.re_128_mem_key, self.re_128_mem_value, re_128_q)
745
+ mo_128_mem_g, mo_128_inds = self.readMem(self.mo_128_mem_key, self.mo_128_mem_value, mo_128_q)
746
+
747
+ le_64_mem_g, le_64_inds = self.readMem(self.le_64_mem_key, self.le_64_mem_value, le_64_q)
748
+ re_64_mem_g, re_64_inds = self.readMem(self.re_64_mem_key, self.re_64_mem_value, re_64_q)
749
+ mo_64_mem_g, mo_64_inds = self.readMem(self.mo_64_mem_key, self.mo_64_mem_value, mo_64_q)
750
+
751
+ if sp_256 is not None and sp_128 is not None and sp_64 is not None:
752
+ le_256_mem_s, _ = self.readMem(sp_256['LE256Key'], sp_256['LE256Value'], le_256_q)
753
+ re_256_mem_s, _ = self.readMem(sp_256['RE256Key'], sp_256['RE256Value'], re_256_q)
754
+ mo_256_mem_s, _ = self.readMem(sp_256['MO256Key'], sp_256['MO256Value'], mo_256_q)
755
+ le_256_mask = self.LE_256_Mask(fs_in['le256'],le_256_mem_s,le_256_mem_g)
756
+ le_256_mem = le_256_mask*le_256_mem_s + (1-le_256_mask)*le_256_mem_g
757
+ re_256_mask = self.RE_256_Mask(fs_in['re256'],re_256_mem_s,re_256_mem_g)
758
+ re_256_mem = re_256_mask*re_256_mem_s + (1-re_256_mask)*re_256_mem_g
759
+ mo_256_mask = self.MO_256_Mask(fs_in['mo256'],mo_256_mem_s,mo_256_mem_g)
760
+ mo_256_mem = mo_256_mask*mo_256_mem_s + (1-mo_256_mask)*mo_256_mem_g
761
+
762
+ le_128_mem_s, _ = self.readMem(sp_128['LE128Key'], sp_128['LE128Value'], le_128_q)
763
+ re_128_mem_s, _ = self.readMem(sp_128['RE128Key'], sp_128['RE128Value'], re_128_q)
764
+ mo_128_mem_s, _ = self.readMem(sp_128['MO128Key'], sp_128['MO128Value'], mo_128_q)
765
+ le_128_mask = self.LE_128_Mask(fs_in['le128'],le_128_mem_s,le_128_mem_g)
766
+ le_128_mem = le_128_mask*le_128_mem_s + (1-le_128_mask)*le_128_mem_g
767
+ re_128_mask = self.RE_128_Mask(fs_in['re128'],re_128_mem_s,re_128_mem_g)
768
+ re_128_mem = re_128_mask*re_128_mem_s + (1-re_128_mask)*re_128_mem_g
769
+ mo_128_mask = self.MO_128_Mask(fs_in['mo128'],mo_128_mem_s,mo_128_mem_g)
770
+ mo_128_mem = mo_128_mask*mo_128_mem_s + (1-mo_128_mask)*mo_128_mem_g
771
+
772
+ le_64_mem_s, _ = self.readMem(sp_64['LE64Key'], sp_64['LE64Value'], le_64_q)
773
+ re_64_mem_s, _ = self.readMem(sp_64['RE64Key'], sp_64['RE64Value'], re_64_q)
774
+ mo_64_mem_s, _ = self.readMem(sp_64['MO64Key'], sp_64['MO64Value'], mo_64_q)
775
+ le_64_mask = self.LE_64_Mask(fs_in['le64'],le_64_mem_s,le_64_mem_g)
776
+ le_64_mem = le_64_mask*le_64_mem_s + (1-le_64_mask)*le_64_mem_g
777
+ re_64_mask = self.RE_64_Mask(fs_in['re64'],re_64_mem_s,re_64_mem_g)
778
+ re_64_mem = re_64_mask*re_64_mem_s + (1-re_64_mask)*re_64_mem_g
779
+ mo_64_mask = self.MO_64_Mask(fs_in['mo64'],mo_64_mem_s,mo_64_mem_g)
780
+ mo_64_mem = mo_64_mask*mo_64_mem_s + (1-mo_64_mask)*mo_64_mem_g
781
+ else:
782
+ le_256_mem = le_256_mem_g
783
+ re_256_mem = re_256_mem_g
784
+ mo_256_mem = mo_256_mem_g
785
+ le_128_mem = le_128_mem_g
786
+ re_128_mem = re_128_mem_g
787
+ mo_128_mem = mo_128_mem_g
788
+ le_64_mem = le_64_mem_g
789
+ re_64_mem = re_64_mem_g
790
+ mo_64_mem = mo_64_mem_g
791
+
792
+ le_256_mem_norm = adaptive_instance_normalization_4D(le_256_mem, fs_in['le256'])
793
+ re_256_mem_norm = adaptive_instance_normalization_4D(re_256_mem, fs_in['re256'])
794
+ mo_256_mem_norm = adaptive_instance_normalization_4D(mo_256_mem, fs_in['mo256'])
795
+
796
+ ####for 128
797
+ le_128_mem_norm = adaptive_instance_normalization_4D(le_128_mem, fs_in['le128'])
798
+ re_128_mem_norm = adaptive_instance_normalization_4D(re_128_mem, fs_in['re128'])
799
+ mo_128_mem_norm = adaptive_instance_normalization_4D(mo_128_mem, fs_in['mo128'])
800
+
801
+ ####for 64
802
+ le_64_mem_norm = adaptive_instance_normalization_4D(le_64_mem, fs_in['le64'])
803
+ re_64_mem_norm = adaptive_instance_normalization_4D(re_64_mem, fs_in['re64'])
804
+ mo_64_mem_norm = adaptive_instance_normalization_4D(mo_64_mem, fs_in['mo64'])
805
+
806
+
807
+ EnMem256 = {'LE256Norm': le_256_mem_norm, 'RE256Norm': re_256_mem_norm, 'MO256Norm': mo_256_mem_norm}
808
+ EnMem128 = {'LE128Norm': le_128_mem_norm, 'RE128Norm': re_128_mem_norm, 'MO128Norm': mo_128_mem_norm}
809
+ EnMem64 = {'LE64Norm': le_64_mem_norm, 'RE64Norm': re_64_mem_norm, 'MO64Norm': mo_64_mem_norm}
810
+ Ind256 = {'LE': le_256_inds, 'RE': re_256_inds, 'MO': mo_256_inds}
811
+ Ind128 = {'LE': le_128_inds, 'RE': re_128_inds, 'MO': mo_128_inds}
812
+ Ind64 = {'LE': le_64_inds, 'RE': re_64_inds, 'MO': mo_64_inds}
813
+ return EnMem256, EnMem128, EnMem64, Ind256, Ind128, Ind64
814
+
815
+ def reconstruct(self, fs_in, locs, memstar):
816
+ le_256_mem_norm, re_256_mem_norm, mo_256_mem_norm = memstar[0]['LE256Norm'], memstar[0]['RE256Norm'], memstar[0]['MO256Norm']
817
+ le_128_mem_norm, re_128_mem_norm, mo_128_mem_norm = memstar[1]['LE128Norm'], memstar[1]['RE128Norm'], memstar[1]['MO128Norm']
818
+ le_64_mem_norm, re_64_mem_norm, mo_64_mem_norm = memstar[2]['LE64Norm'], memstar[2]['RE64Norm'], memstar[2]['MO64Norm']
819
+
820
+ le_256_final = self.LE_256_Attention(le_256_mem_norm - fs_in['le256']) * le_256_mem_norm + fs_in['le256']
821
+ re_256_final = self.RE_256_Attention(re_256_mem_norm - fs_in['re256']) * re_256_mem_norm + fs_in['re256']
822
+ mo_256_final = self.MO_256_Attention(mo_256_mem_norm - fs_in['mo256']) * mo_256_mem_norm + fs_in['mo256']
823
+
824
+ le_128_final = self.LE_128_Attention(le_128_mem_norm - fs_in['le128']) * le_128_mem_norm + fs_in['le128']
825
+ re_128_final = self.RE_128_Attention(re_128_mem_norm - fs_in['re128']) * re_128_mem_norm + fs_in['re128']
826
+ mo_128_final = self.MO_128_Attention(mo_128_mem_norm - fs_in['mo128']) * mo_128_mem_norm + fs_in['mo128']
827
+
828
+ le_64_final = self.LE_64_Attention(le_64_mem_norm - fs_in['le64']) * le_64_mem_norm + fs_in['le64']
829
+ re_64_final = self.RE_64_Attention(re_64_mem_norm - fs_in['re64']) * re_64_mem_norm + fs_in['re64']
830
+ mo_64_final = self.MO_64_Attention(mo_64_mem_norm - fs_in['mo64']) * mo_64_mem_norm + fs_in['mo64']
831
+
832
+
833
+ le_location = locs[:,0,:]
834
+ re_location = locs[:,1,:]
835
+ mo_location = locs[:,3,:]
836
+
837
+ # Somehow with latest Torch it doesn't like numpy wrappers anymore
838
+
839
+ # le_location = le_location.cpu().int().numpy()
840
+ # re_location = re_location.cpu().int().numpy()
841
+ # mo_location = mo_location.cpu().int().numpy()
842
+ le_location = le_location.cpu().int()
843
+ re_location = re_location.cpu().int()
844
+ mo_location = mo_location.cpu().int()
845
+
846
+ up_in_256 = fs_in['f256'].clone()# * 0
847
+ up_in_128 = fs_in['f128'].clone()# * 0
848
+ up_in_64 = fs_in['f64'].clone()# * 0
849
+
850
+ for i in range(fs_in['f256'].size(0)):
851
+ up_in_256[i:i+1,:,le_location[i,1]//2:le_location[i,3]//2,le_location[i,0]//2:le_location[i,2]//2] = F.interpolate(le_256_final[i:i+1,:,:,:].clone(), (le_location[i,3]//2-le_location[i,1]//2,le_location[i,2]//2-le_location[i,0]//2),mode='bilinear',align_corners=False)
852
+ up_in_256[i:i+1,:,re_location[i,1]//2:re_location[i,3]//2,re_location[i,0]//2:re_location[i,2]//2] = F.interpolate(re_256_final[i:i+1,:,:,:].clone(), (re_location[i,3]//2-re_location[i,1]//2,re_location[i,2]//2-re_location[i,0]//2),mode='bilinear',align_corners=False)
853
+ up_in_256[i:i+1,:,mo_location[i,1]//2:mo_location[i,3]//2,mo_location[i,0]//2:mo_location[i,2]//2] = F.interpolate(mo_256_final[i:i+1,:,:,:].clone(), (mo_location[i,3]//2-mo_location[i,1]//2,mo_location[i,2]//2-mo_location[i,0]//2),mode='bilinear',align_corners=False)
854
+
855
+ up_in_128[i:i+1,:,le_location[i,1]//4:le_location[i,3]//4,le_location[i,0]//4:le_location[i,2]//4] = F.interpolate(le_128_final[i:i+1,:,:,:].clone(), (le_location[i,3]//4-le_location[i,1]//4,le_location[i,2]//4-le_location[i,0]//4),mode='bilinear',align_corners=False)
856
+ up_in_128[i:i+1,:,re_location[i,1]//4:re_location[i,3]//4,re_location[i,0]//4:re_location[i,2]//4] = F.interpolate(re_128_final[i:i+1,:,:,:].clone(), (re_location[i,3]//4-re_location[i,1]//4,re_location[i,2]//4-re_location[i,0]//4),mode='bilinear',align_corners=False)
857
+ up_in_128[i:i+1,:,mo_location[i,1]//4:mo_location[i,3]//4,mo_location[i,0]//4:mo_location[i,2]//4] = F.interpolate(mo_128_final[i:i+1,:,:,:].clone(), (mo_location[i,3]//4-mo_location[i,1]//4,mo_location[i,2]//4-mo_location[i,0]//4),mode='bilinear',align_corners=False)
858
+
859
+ up_in_64[i:i+1,:,le_location[i,1]//8:le_location[i,3]//8,le_location[i,0]//8:le_location[i,2]//8] = F.interpolate(le_64_final[i:i+1,:,:,:].clone(), (le_location[i,3]//8-le_location[i,1]//8,le_location[i,2]//8-le_location[i,0]//8),mode='bilinear',align_corners=False)
860
+ up_in_64[i:i+1,:,re_location[i,1]//8:re_location[i,3]//8,re_location[i,0]//8:re_location[i,2]//8] = F.interpolate(re_64_final[i:i+1,:,:,:].clone(), (re_location[i,3]//8-re_location[i,1]//8,re_location[i,2]//8-re_location[i,0]//8),mode='bilinear',align_corners=False)
861
+ up_in_64[i:i+1,:,mo_location[i,1]//8:mo_location[i,3]//8,mo_location[i,0]//8:mo_location[i,2]//8] = F.interpolate(mo_64_final[i:i+1,:,:,:].clone(), (mo_location[i,3]//8-mo_location[i,1]//8,mo_location[i,2]//8-mo_location[i,0]//8),mode='bilinear',align_corners=False)
862
+
863
+ ms_in_64 = self.MSDilate(fs_in['f64'].clone())
864
+ fea_up1 = self.up1(ms_in_64, up_in_64)
865
+ fea_up2 = self.up2(fea_up1, up_in_128) #
866
+ fea_up3 = self.up3(fea_up2, up_in_256) #
867
+ output = self.up4(fea_up3) #
868
+ return output
869
+
870
+ def generate_specific_dictionary(self, sp_imgs=None, sp_locs=None):
871
+ return self.memorize(sp_imgs, sp_locs)
872
+
873
+ def forward(self, lq=None, loc=None, sp_256 = None, sp_128 = None, sp_64 = None):
874
+ try:
875
+ fs_in = self.E_lq(lq, loc) # low quality images
876
+ except Exception as e:
877
+ print(e)
878
+
879
+ GeMemNorm256, GeMemNorm128, GeMemNorm64, Ind256, Ind128, Ind64 = self.enhancer(fs_in)
880
+ GeOut = self.reconstruct(fs_in, loc, memstar = [GeMemNorm256, GeMemNorm128, GeMemNorm64])
881
+ if sp_256 is not None and sp_128 is not None and sp_64 is not None:
882
+ GSMemNorm256, GSMemNorm128, GSMemNorm64, _, _, _ = self.enhancer(fs_in, sp_256, sp_128, sp_64)
883
+ GSOut = self.reconstruct(fs_in, loc, memstar = [GSMemNorm256, GSMemNorm128, GSMemNorm64])
884
+ else:
885
+ GSOut = None
886
+ return GeOut, GSOut
887
+
888
+ class UpResBlock(nn.Module):
889
+ def __init__(self, dim, conv_layer = nn.Conv2d, norm_layer = nn.BatchNorm2d):
890
+ super(UpResBlock, self).__init__()
891
+ self.Model = nn.Sequential(
892
+ SpectralNorm(conv_layer(dim, dim, 3, 1, 1)),
893
+ nn.LeakyReLU(0.2),
894
+ SpectralNorm(conv_layer(dim, dim, 3, 1, 1)),
895
+ )
896
+ def forward(self, x):
897
+ out = x + self.Model(x)
898
+ return out
roop/processors/Enhance_GFPGAN.py ADDED
@@ -0,0 +1,77 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Any, List, Callable
2
+ import cv2
3
+ import numpy as np
4
+ import onnxruntime
5
+ import roop.globals
6
+
7
+ from roop.typing import Face, Frame, FaceSet
8
+ from roop.utilities import resolve_relative_path
9
+
10
+
11
+ # THREAD_LOCK = threading.Lock()
12
+
13
+
14
+ class Enhance_GFPGAN():
15
+ plugin_options:dict = None
16
+
17
+ model_gfpgan = None
18
+ name = None
19
+ devicename = None
20
+
21
+ processorname = 'gfpgan'
22
+ type = 'enhance'
23
+
24
+
25
+ def Initialize(self, plugin_options:dict):
26
+ if self.plugin_options is not None:
27
+ if self.plugin_options["devicename"] != plugin_options["devicename"]:
28
+ self.Release()
29
+
30
+ self.plugin_options = plugin_options
31
+ if self.model_gfpgan is None:
32
+ model_path = resolve_relative_path('../models/GFPGANv1.4.onnx')
33
+ self.model_gfpgan = onnxruntime.InferenceSession(model_path, None, providers=roop.globals.execution_providers)
34
+ # replace Mac mps with cpu for the moment
35
+ self.devicename = self.plugin_options["devicename"].replace('mps', 'cpu')
36
+
37
+ self.name = self.model_gfpgan.get_inputs()[0].name
38
+
39
+ def Run(self, source_faceset: FaceSet, target_face: Face, temp_frame: Frame) -> Frame:
40
+ # preprocess
41
+ input_size = temp_frame.shape[1]
42
+ temp_frame = cv2.resize(temp_frame, (512, 512), cv2.INTER_CUBIC)
43
+
44
+ temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_BGR2RGB)
45
+ temp_frame = temp_frame.astype('float32') / 255.0
46
+ temp_frame = (temp_frame - 0.5) / 0.5
47
+ temp_frame = np.expand_dims(temp_frame, axis=0).transpose(0, 3, 1, 2)
48
+
49
+ io_binding = self.model_gfpgan.io_binding()
50
+ io_binding.bind_cpu_input("input", temp_frame)
51
+ io_binding.bind_output("1288", self.devicename)
52
+ self.model_gfpgan.run_with_iobinding(io_binding)
53
+ ort_outs = io_binding.copy_outputs_to_cpu()
54
+ result = ort_outs[0][0]
55
+
56
+ # post-process
57
+ result = np.clip(result, -1, 1)
58
+ result = (result + 1) / 2
59
+ result = result.transpose(1, 2, 0) * 255.0
60
+ result = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
61
+ scale_factor = int(result.shape[1] / input_size)
62
+ return result.astype(np.uint8), scale_factor
63
+
64
+
65
+ def Release(self):
66
+ self.model_gfpgan = None
67
+
68
+
69
+
70
+
71
+
72
+
73
+
74
+
75
+
76
+
77
+
roop/processors/Enhance_GPEN.py ADDED
@@ -0,0 +1,63 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Any, List, Callable
2
+ import cv2
3
+ import numpy as np
4
+ import onnxruntime
5
+ import roop.globals
6
+
7
+ from roop.typing import Face, Frame, FaceSet
8
+ from roop.utilities import resolve_relative_path
9
+
10
+
11
+ class Enhance_GPEN():
12
+ plugin_options:dict = None
13
+
14
+ model_gpen = None
15
+ name = None
16
+ devicename = None
17
+
18
+ processorname = 'gpen'
19
+ type = 'enhance'
20
+
21
+
22
+ def Initialize(self, plugin_options:dict):
23
+ if self.plugin_options is not None:
24
+ if self.plugin_options["devicename"] != plugin_options["devicename"]:
25
+ self.Release()
26
+
27
+ self.plugin_options = plugin_options
28
+ if self.model_gpen is None:
29
+ model_path = resolve_relative_path('../models/GPEN-BFR-512.onnx')
30
+ self.model_gpen = onnxruntime.InferenceSession(model_path, None, providers=roop.globals.execution_providers)
31
+ # replace Mac mps with cpu for the moment
32
+ self.devicename = self.plugin_options["devicename"].replace('mps', 'cpu')
33
+
34
+ self.name = self.model_gpen.get_inputs()[0].name
35
+
36
+ def Run(self, source_faceset: FaceSet, target_face: Face, temp_frame: Frame) -> Frame:
37
+ # preprocess
38
+ input_size = temp_frame.shape[1]
39
+ temp_frame = cv2.resize(temp_frame, (512, 512), cv2.INTER_CUBIC)
40
+
41
+ temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_BGR2RGB)
42
+ temp_frame = temp_frame.astype('float32') / 255.0
43
+ temp_frame = (temp_frame - 0.5) / 0.5
44
+ temp_frame = np.expand_dims(temp_frame, axis=0).transpose(0, 3, 1, 2)
45
+
46
+ io_binding = self.model_gpen.io_binding()
47
+ io_binding.bind_cpu_input("input", temp_frame)
48
+ io_binding.bind_output("output", self.devicename)
49
+ self.model_gpen.run_with_iobinding(io_binding)
50
+ ort_outs = io_binding.copy_outputs_to_cpu()
51
+ result = ort_outs[0][0]
52
+
53
+ # post-process
54
+ result = np.clip(result, -1, 1)
55
+ result = (result + 1) / 2
56
+ result = result.transpose(1, 2, 0) * 255.0
57
+ result = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
58
+ scale_factor = int(result.shape[1] / input_size)
59
+ return result.astype(np.uint8), scale_factor
60
+
61
+
62
+ def Release(self):
63
+ self.model_gpen = None
roop/processors/Enhance_RestoreFormerPPlus.py ADDED
@@ -0,0 +1,64 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Any, List, Callable
2
+ import cv2
3
+ import numpy as np
4
+ import onnxruntime
5
+ import roop.globals
6
+
7
+ from roop.typing import Face, Frame, FaceSet
8
+ from roop.utilities import resolve_relative_path
9
+
10
+ class Enhance_RestoreFormerPPlus():
11
+ plugin_options:dict = None
12
+ model_restoreformerpplus = None
13
+ devicename = None
14
+ name = None
15
+
16
+ processorname = 'restoreformer++'
17
+ type = 'enhance'
18
+
19
+
20
+ def Initialize(self, plugin_options:dict):
21
+ if self.plugin_options is not None:
22
+ if self.plugin_options["devicename"] != plugin_options["devicename"]:
23
+ self.Release()
24
+
25
+ self.plugin_options = plugin_options
26
+ if self.model_restoreformerpplus is None:
27
+ # replace Mac mps with cpu for the moment
28
+ self.devicename = self.plugin_options["devicename"].replace('mps', 'cpu')
29
+ model_path = resolve_relative_path('../models/restoreformer_plus_plus.onnx')
30
+ self.model_restoreformerpplus = onnxruntime.InferenceSession(model_path, None, providers=roop.globals.execution_providers)
31
+ self.model_inputs = self.model_restoreformerpplus.get_inputs()
32
+ model_outputs = self.model_restoreformerpplus.get_outputs()
33
+ self.io_binding = self.model_restoreformerpplus.io_binding()
34
+ self.io_binding.bind_output(model_outputs[0].name, self.devicename)
35
+
36
+ def Run(self, source_faceset: FaceSet, target_face: Face, temp_frame: Frame) -> Frame:
37
+ # preprocess
38
+ input_size = temp_frame.shape[1]
39
+ temp_frame = cv2.resize(temp_frame, (512, 512), cv2.INTER_CUBIC)
40
+ temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_BGR2RGB)
41
+ temp_frame = temp_frame.astype('float32') / 255.0
42
+ temp_frame = (temp_frame - 0.5) / 0.5
43
+ temp_frame = np.expand_dims(temp_frame, axis=0).transpose(0, 3, 1, 2)
44
+
45
+ self.io_binding.bind_cpu_input(self.model_inputs[0].name, temp_frame) # .astype(np.float32)
46
+ self.model_restoreformerpplus.run_with_iobinding(self.io_binding)
47
+ ort_outs = self.io_binding.copy_outputs_to_cpu()
48
+ result = ort_outs[0][0]
49
+ del ort_outs
50
+
51
+ result = np.clip(result, -1, 1)
52
+ result = (result + 1) / 2
53
+ result = result.transpose(1, 2, 0) * 255.0
54
+ result = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
55
+ scale_factor = int(result.shape[1] / input_size)
56
+ return result.astype(np.uint8), scale_factor
57
+
58
+
59
+ def Release(self):
60
+ del self.model_restoreformerpplus
61
+ self.model_restoreformerpplus = None
62
+ del self.io_binding
63
+ self.io_binding = None
64
+
roop/processors/FaceSwapInsightFace.py ADDED
@@ -0,0 +1,61 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import roop.globals
2
+ import cv2
3
+ import numpy as np
4
+ import onnx
5
+ import onnxruntime
6
+
7
+ from roop.typing import Face, Frame
8
+ from roop.utilities import resolve_relative_path
9
+
10
+
11
+
12
+ class FaceSwapInsightFace():
13
+ plugin_options:dict = None
14
+ model_swap_insightface = None
15
+
16
+ processorname = 'faceswap'
17
+ type = 'swap'
18
+
19
+
20
+ def Initialize(self, plugin_options:dict):
21
+ if self.plugin_options is not None:
22
+ if self.plugin_options["devicename"] != plugin_options["devicename"]:
23
+ self.Release()
24
+
25
+ self.plugin_options = plugin_options
26
+ if self.model_swap_insightface is None:
27
+ model_path = resolve_relative_path('../models/inswapper_128.onnx')
28
+ graph = onnx.load(model_path).graph
29
+ self.emap = onnx.numpy_helper.to_array(graph.initializer[-1])
30
+ self.devicename = self.plugin_options["devicename"].replace('mps', 'cpu')
31
+ self.input_mean = 0.0
32
+ self.input_std = 255.0
33
+ #cuda_options = {"arena_extend_strategy": "kSameAsRequested", 'cudnn_conv_algo_search': 'DEFAULT'}
34
+ sess_options = onnxruntime.SessionOptions()
35
+ sess_options.enable_cpu_mem_arena = False
36
+ self.model_swap_insightface = onnxruntime.InferenceSession(model_path, sess_options, providers=roop.globals.execution_providers)
37
+
38
+
39
+
40
+ def Run(self, source_face: Face, target_face: Face, temp_frame: Frame) -> Frame:
41
+ latent = source_face.normed_embedding.reshape((1,-1))
42
+ latent = np.dot(latent, self.emap)
43
+ latent /= np.linalg.norm(latent)
44
+ io_binding = self.model_swap_insightface.io_binding()
45
+ io_binding.bind_cpu_input("target", temp_frame)
46
+ io_binding.bind_cpu_input("source", latent)
47
+ io_binding.bind_output("output", self.devicename)
48
+ self.model_swap_insightface.run_with_iobinding(io_binding)
49
+ ort_outs = io_binding.copy_outputs_to_cpu()[0]
50
+ return ort_outs[0]
51
+
52
+
53
+ def Release(self):
54
+ del self.model_swap_insightface
55
+ self.model_swap_insightface = None
56
+
57
+
58
+
59
+
60
+
61
+
roop/processors/Frame_Colorizer.py ADDED
@@ -0,0 +1,70 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import cv2
2
+ import numpy as np
3
+ import onnxruntime
4
+ import roop.globals
5
+
6
+ from roop.utilities import resolve_relative_path
7
+ from roop.typing import Frame
8
+
9
+ class Frame_Colorizer():
10
+ plugin_options:dict = None
11
+ model_colorizer = None
12
+ devicename = None
13
+ prev_type = None
14
+
15
+ processorname = 'deoldify'
16
+ type = 'frame_colorizer'
17
+
18
+
19
+ def Initialize(self, plugin_options:dict):
20
+ if self.plugin_options is not None:
21
+ if self.plugin_options["devicename"] != plugin_options["devicename"]:
22
+ self.Release()
23
+
24
+ self.plugin_options = plugin_options
25
+ if self.prev_type is not None and self.prev_type != self.plugin_options["subtype"]:
26
+ self.Release()
27
+ self.prev_type = self.plugin_options["subtype"]
28
+ if self.model_colorizer is None:
29
+ # replace Mac mps with cpu for the moment
30
+ self.devicename = self.plugin_options["devicename"].replace('mps', 'cpu')
31
+ if self.prev_type == "deoldify_artistic":
32
+ model_path = resolve_relative_path('../models/Frame/deoldify_artistic.onnx')
33
+ elif self.prev_type == "deoldify_stable":
34
+ model_path = resolve_relative_path('../models/Frame/deoldify_stable.onnx')
35
+
36
+ onnxruntime.set_default_logger_severity(3)
37
+ self.model_colorizer = onnxruntime.InferenceSession(model_path, None, providers=roop.globals.execution_providers)
38
+ self.model_inputs = self.model_colorizer.get_inputs()
39
+ model_outputs = self.model_colorizer.get_outputs()
40
+ self.io_binding = self.model_colorizer.io_binding()
41
+ self.io_binding.bind_output(model_outputs[0].name, self.devicename)
42
+
43
+ def Run(self, input_frame: Frame) -> Frame:
44
+ temp_frame = cv2.cvtColor(input_frame, cv2.COLOR_BGR2GRAY)
45
+ temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_GRAY2RGB)
46
+ temp_frame = cv2.resize(temp_frame, (256, 256))
47
+ temp_frame = temp_frame.transpose((2, 0, 1))
48
+ temp_frame = np.expand_dims(temp_frame, axis=0).astype(np.float32)
49
+ self.io_binding.bind_cpu_input(self.model_inputs[0].name, temp_frame)
50
+ self.model_colorizer.run_with_iobinding(self.io_binding)
51
+ ort_outs = self.io_binding.copy_outputs_to_cpu()
52
+ result = ort_outs[0][0]
53
+ del ort_outs
54
+ colorized_frame = result.transpose(1, 2, 0)
55
+ colorized_frame = cv2.resize(colorized_frame, (input_frame.shape[1], input_frame.shape[0]))
56
+ temp_blue_channel, _, _ = cv2.split(input_frame)
57
+ colorized_frame = cv2.cvtColor(colorized_frame, cv2.COLOR_BGR2RGB).astype(np.uint8)
58
+ colorized_frame = cv2.cvtColor(colorized_frame, cv2.COLOR_BGR2LAB)
59
+ _, color_green_channel, color_red_channel = cv2.split(colorized_frame)
60
+ colorized_frame = cv2.merge((temp_blue_channel, color_green_channel, color_red_channel))
61
+ colorized_frame = cv2.cvtColor(colorized_frame, cv2.COLOR_LAB2BGR)
62
+ return colorized_frame.astype(np.uint8)
63
+
64
+
65
+ def Release(self):
66
+ del self.model_colorizer
67
+ self.model_colorizer = None
68
+ del self.io_binding
69
+ self.io_binding = None
70
+
roop/processors/Frame_Filter.py ADDED
@@ -0,0 +1,105 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import cv2
2
+ import numpy as np
3
+
4
+ from roop.typing import Frame
5
+
6
+ class Frame_Filter():
7
+ processorname = 'generic_filter'
8
+ type = 'frame_processor'
9
+
10
+ plugin_options:dict = None
11
+
12
+ c64_palette = np.array([
13
+ [0, 0, 0],
14
+ [255, 255, 255],
15
+ [0x81, 0x33, 0x38],
16
+ [0x75, 0xce, 0xc8],
17
+ [0x8e, 0x3c, 0x97],
18
+ [0x56, 0xac, 0x4d],
19
+ [0x2e, 0x2c, 0x9b],
20
+ [0xed, 0xf1, 0x71],
21
+ [0x8e, 0x50, 0x29],
22
+ [0x55, 0x38, 0x00],
23
+ [0xc4, 0x6c, 0x71],
24
+ [0x4a, 0x4a, 0x4a],
25
+ [0x7b, 0x7b, 0x7b],
26
+ [0xa9, 0xff, 0x9f],
27
+ [0x70, 0x6d, 0xeb],
28
+ [0xb2, 0xb2, 0xb2]
29
+ ])
30
+
31
+
32
+ def RenderC64Screen(self, image):
33
+ # Simply round the color values to the nearest color in the palette
34
+ image = cv2.resize(image,(320,200))
35
+ palette = self.c64_palette / 255.0 # Normalize palette
36
+ img_normalized = image / 255.0 # Normalize image
37
+
38
+ # Calculate the index in the palette that is closest to each pixel in the image
39
+ indices = np.sqrt(((img_normalized[:, :, None, :] - palette[None, None, :, :]) ** 2).sum(axis=3)).argmin(axis=2)
40
+ # Map the image to the palette colors
41
+ mapped_image = palette[indices]
42
+ return (mapped_image * 255).astype(np.uint8) # Denormalize and return the image
43
+
44
+
45
+ def RenderDetailEnhance(self, image):
46
+ return cv2.detailEnhance(image)
47
+
48
+ def RenderStylize(self, image):
49
+ return cv2.stylization(image)
50
+
51
+ def RenderPencilSketch(self, image):
52
+ imgray, imout = cv2.pencilSketch(image, sigma_s=60, sigma_r=0.07, shade_factor=0.05)
53
+ return imout
54
+
55
+ def RenderCartoon(self, image):
56
+ numDownSamples = 2 # number of downscaling steps
57
+ numBilateralFilters = 7 # number of bilateral filtering steps
58
+
59
+ img_color = image
60
+ for _ in range(numDownSamples):
61
+ img_color = cv2.pyrDown(img_color)
62
+ for _ in range(numBilateralFilters):
63
+ img_color = cv2.bilateralFilter(img_color, 9, 9, 7)
64
+ for _ in range(numDownSamples):
65
+ img_color = cv2.pyrUp(img_color)
66
+ img_gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
67
+ img_blur = cv2.medianBlur(img_gray, 7)
68
+ img_edge = cv2.adaptiveThreshold(img_blur, 255,
69
+ cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 9, 2)
70
+ img_edge = cv2.cvtColor(img_edge, cv2.COLOR_GRAY2RGB)
71
+ if img_color.shape != image.shape:
72
+ img_color = cv2.resize(img_color, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_LINEAR)
73
+ if img_color.shape != img_edge.shape:
74
+ img_edge = cv2.resize(img_edge, (img_color.shape[1], img_color.shape[0]), interpolation=cv2.INTER_LINEAR)
75
+ return cv2.bitwise_and(img_color, img_edge)
76
+
77
+
78
+ def Initialize(self, plugin_options:dict):
79
+ if self.plugin_options is not None:
80
+ if self.plugin_options["devicename"] != plugin_options["devicename"]:
81
+ self.Release()
82
+ self.plugin_options = plugin_options
83
+
84
+ def Run(self, temp_frame: Frame) -> Frame:
85
+ subtype = self.plugin_options["subtype"]
86
+ if subtype == "stylize":
87
+ return self.RenderStylize(temp_frame).astype(np.uint8)
88
+ if subtype == "detailenhance":
89
+ return self.RenderDetailEnhance(temp_frame).astype(np.uint8)
90
+ if subtype == "pencil":
91
+ return self.RenderPencilSketch(temp_frame).astype(np.uint8)
92
+ if subtype == "cartoon":
93
+ return self.RenderCartoon(temp_frame).astype(np.uint8)
94
+ if subtype == "C64":
95
+ return self.RenderC64Screen(temp_frame).astype(np.uint8)
96
+
97
+
98
+ def Release(self):
99
+ pass
100
+
101
+ def getProcessedResolution(self, width, height):
102
+ if self.plugin_options["subtype"] == "C64":
103
+ return (320,200)
104
+ return None
105
+
roop/processors/Frame_Masking.py ADDED
@@ -0,0 +1,71 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import cv2
2
+ import numpy as np
3
+ import onnxruntime
4
+ import roop.globals
5
+
6
+ from roop.utilities import resolve_relative_path
7
+ from roop.typing import Frame
8
+
9
+ class Frame_Masking():
10
+ plugin_options:dict = None
11
+ model_masking = None
12
+ devicename = None
13
+ name = None
14
+
15
+ processorname = 'removebg'
16
+ type = 'frame_masking'
17
+
18
+
19
+ def Initialize(self, plugin_options:dict):
20
+ if self.plugin_options is not None:
21
+ if self.plugin_options["devicename"] != plugin_options["devicename"]:
22
+ self.Release()
23
+
24
+ self.plugin_options = plugin_options
25
+ if self.model_masking is None:
26
+ # replace Mac mps with cpu for the moment
27
+ self.devicename = self.plugin_options["devicename"]
28
+ self.devicename = self.devicename.replace('mps', 'cpu')
29
+ model_path = resolve_relative_path('../models/Frame/isnet-general-use.onnx')
30
+ self.model_masking = onnxruntime.InferenceSession(model_path, None, providers=roop.globals.execution_providers)
31
+ self.model_inputs = self.model_masking.get_inputs()
32
+ model_outputs = self.model_masking.get_outputs()
33
+ self.io_binding = self.model_masking.io_binding()
34
+ self.io_binding.bind_output(model_outputs[0].name, self.devicename)
35
+
36
+ def Run(self, temp_frame: Frame) -> Frame:
37
+ # Pre process:Resize, BGR->RGB, float32 cast
38
+ input_image = cv2.resize(temp_frame, (1024, 1024))
39
+ input_image = cv2.cvtColor(input_image, cv2.COLOR_BGR2RGB)
40
+ mean = [0.5, 0.5, 0.5]
41
+ std = [1.0, 1.0, 1.0]
42
+ input_image = (input_image / 255.0 - mean) / std
43
+ input_image = input_image.transpose(2, 0, 1)
44
+ input_image = np.expand_dims(input_image, axis=0)
45
+ input_image = input_image.astype('float32')
46
+
47
+ self.io_binding.bind_cpu_input(self.model_inputs[0].name, input_image)
48
+ self.model_masking.run_with_iobinding(self.io_binding)
49
+ ort_outs = self.io_binding.copy_outputs_to_cpu()
50
+ result = ort_outs[0][0]
51
+ del ort_outs
52
+ # Post process:squeeze, Sigmoid, Normarize, uint8 cast
53
+ mask = np.squeeze(result[0])
54
+ min_value = np.min(mask)
55
+ max_value = np.max(mask)
56
+ mask = (mask - min_value) / (max_value - min_value)
57
+ #mask = np.where(mask < score_th, 0, 1)
58
+ #mask *= 255
59
+ mask = cv2.resize(mask, (temp_frame.shape[1], temp_frame.shape[0]), interpolation=cv2.INTER_LINEAR)
60
+ mask = np.reshape(mask, [mask.shape[0],mask.shape[1],1])
61
+ result = mask * temp_frame.astype(np.float32)
62
+ return result.astype(np.uint8)
63
+
64
+
65
+
66
+ def Release(self):
67
+ del self.model_masking
68
+ self.model_masking = None
69
+ del self.io_binding
70
+ self.io_binding = None
71
+
roop/processors/Frame_Upscale.py ADDED
@@ -0,0 +1,131 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import cv2
2
+ import numpy as np
3
+ import onnxruntime
4
+ import roop.globals
5
+ import threading
6
+
7
+ from roop.utilities import resolve_relative_path
8
+ from roop.typing import Frame
9
+
10
+ class Frame_Upscale():
11
+ plugin_options:dict = None
12
+ model_upscale = None
13
+ devicename = None
14
+ prev_type = None
15
+
16
+ processorname = 'upscale'
17
+ type = 'frame_enhancer'
18
+
19
+ THREAD_LOCK_UPSCALE = threading.Lock()
20
+
21
+
22
+ def Initialize(self, plugin_options:dict):
23
+ if self.plugin_options is not None:
24
+ if self.plugin_options["devicename"] != plugin_options["devicename"]:
25
+ self.Release()
26
+
27
+ self.plugin_options = plugin_options
28
+ if self.prev_type is not None and self.prev_type != self.plugin_options["subtype"]:
29
+ self.Release()
30
+ self.prev_type = self.plugin_options["subtype"]
31
+ if self.model_upscale is None:
32
+ # replace Mac mps with cpu for the moment
33
+ self.devicename = self.plugin_options["devicename"].replace('mps', 'cpu')
34
+ if self.prev_type == "esrganx4":
35
+ model_path = resolve_relative_path('../models/Frame/real_esrgan_x4.onnx')
36
+ self.scale = 4
37
+ elif self.prev_type == "esrganx2":
38
+ model_path = resolve_relative_path('../models/Frame/real_esrgan_x2.onnx')
39
+ self.scale = 2
40
+ elif self.prev_type == "lsdirx4":
41
+ model_path = resolve_relative_path('../models/Frame/lsdir_x4.onnx')
42
+ self.scale = 4
43
+
44
+ self.model_upscale = onnxruntime.InferenceSession(model_path, None, providers=roop.globals.execution_providers)
45
+ self.model_inputs = self.model_upscale.get_inputs()
46
+ model_outputs = self.model_upscale.get_outputs()
47
+ self.io_binding = self.model_upscale.io_binding()
48
+ self.io_binding.bind_output(model_outputs[0].name, self.devicename)
49
+
50
+ def getProcessedResolution(self, width, height):
51
+ return (width * self.scale, height * self.scale)
52
+
53
+ # borrowed from facefusion -> https://github.com/facefusion/facefusion
54
+ def prepare_tile_frame(self, tile_frame : Frame) -> Frame:
55
+ tile_frame = np.expand_dims(tile_frame[:, :, ::-1], axis = 0)
56
+ tile_frame = tile_frame.transpose(0, 3, 1, 2)
57
+ tile_frame = tile_frame.astype(np.float32) / 255
58
+ return tile_frame
59
+
60
+
61
+ def normalize_tile_frame(self, tile_frame : Frame) -> Frame:
62
+ tile_frame = tile_frame.transpose(0, 2, 3, 1).squeeze(0) * 255
63
+ tile_frame = tile_frame.clip(0, 255).astype(np.uint8)[:, :, ::-1]
64
+ return tile_frame
65
+
66
+ def create_tile_frames(self, input_frame : Frame, size):
67
+ input_frame = np.pad(input_frame, ((size[1], size[1]), (size[1], size[1]), (0, 0)))
68
+ tile_width = size[0] - 2 * size[2]
69
+ pad_size_bottom = size[2] + tile_width - input_frame.shape[0] % tile_width
70
+ pad_size_right = size[2] + tile_width - input_frame.shape[1] % tile_width
71
+ pad_vision_frame = np.pad(input_frame, ((size[2], pad_size_bottom), (size[2], pad_size_right), (0, 0)))
72
+ pad_height, pad_width = pad_vision_frame.shape[:2]
73
+ row_range = range(size[2], pad_height - size[2], tile_width)
74
+ col_range = range(size[2], pad_width - size[2], tile_width)
75
+ tile_frames = []
76
+
77
+ for row_frame in row_range:
78
+ top = row_frame - size[2]
79
+ bottom = row_frame + size[2] + tile_width
80
+ for column_vision_frame in col_range:
81
+ left = column_vision_frame - size[2]
82
+ right = column_vision_frame + size[2] + tile_width
83
+ tile_frames.append(pad_vision_frame[top:bottom, left:right, :])
84
+ return tile_frames, pad_width, pad_height
85
+
86
+
87
+ def merge_tile_frames(self, tile_frames, temp_width : int, temp_height : int, pad_width : int, pad_height : int, size) -> Frame:
88
+ merge_frame = np.zeros((pad_height, pad_width, 3)).astype(np.uint8)
89
+ tile_width = tile_frames[0].shape[1] - 2 * size[2]
90
+ tiles_per_row = min(pad_width // tile_width, len(tile_frames))
91
+
92
+ for index, tile_frame in enumerate(tile_frames):
93
+ tile_frame = tile_frame[size[2]:-size[2], size[2]:-size[2]]
94
+ row_index = index // tiles_per_row
95
+ col_index = index % tiles_per_row
96
+ top = row_index * tile_frame.shape[0]
97
+ bottom = top + tile_frame.shape[0]
98
+ left = col_index * tile_frame.shape[1]
99
+ right = left + tile_frame.shape[1]
100
+ merge_frame[top:bottom, left:right, :] = tile_frame
101
+ merge_frame = merge_frame[size[1] : size[1] + temp_height, size[1]: size[1] + temp_width, :]
102
+ return merge_frame
103
+
104
+
105
+ def Run(self, temp_frame: Frame) -> Frame:
106
+ size = (128, 8, 2)
107
+ temp_height, temp_width = temp_frame.shape[:2]
108
+ upscale_tile_frames, pad_width, pad_height = self.create_tile_frames(temp_frame, size)
109
+
110
+ for index, tile_frame in enumerate(upscale_tile_frames):
111
+ tile_frame = self.prepare_tile_frame(tile_frame)
112
+ with self.THREAD_LOCK_UPSCALE:
113
+ self.io_binding.bind_cpu_input(self.model_inputs[0].name, tile_frame)
114
+ self.model_upscale.run_with_iobinding(self.io_binding)
115
+ ort_outs = self.io_binding.copy_outputs_to_cpu()
116
+ result = ort_outs[0]
117
+ upscale_tile_frames[index] = self.normalize_tile_frame(result)
118
+ final_frame = self.merge_tile_frames(upscale_tile_frames, temp_width * self.scale
119
+ , temp_height * self.scale
120
+ , pad_width * self.scale, pad_height * self.scale
121
+ , (size[0] * self.scale, size[1] * self.scale, size[2] * self.scale))
122
+ return final_frame.astype(np.uint8)
123
+
124
+
125
+
126
+ def Release(self):
127
+ del self.model_upscale
128
+ self.model_upscale = None
129
+ del self.io_binding
130
+ self.io_binding = None
131
+
roop/processors/Mask_Clip2Seg.py ADDED
@@ -0,0 +1,94 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import cv2
2
+ import numpy as np
3
+ import torch
4
+ import threading
5
+ from torchvision import transforms
6
+ from clip.clipseg import CLIPDensePredT
7
+ import numpy as np
8
+
9
+ from roop.typing import Frame
10
+
11
+ THREAD_LOCK_CLIP = threading.Lock()
12
+
13
+
14
+ class Mask_Clip2Seg():
15
+ plugin_options:dict = None
16
+ model_clip = None
17
+
18
+ processorname = 'clip2seg'
19
+ type = 'mask'
20
+
21
+
22
+ def Initialize(self, plugin_options:dict):
23
+ if self.plugin_options is not None:
24
+ if self.plugin_options["devicename"] != plugin_options["devicename"]:
25
+ self.Release()
26
+
27
+ self.plugin_options = plugin_options
28
+ if self.model_clip is None:
29
+ self.model_clip = CLIPDensePredT(version='ViT-B/16', reduce_dim=64, complex_trans_conv=True)
30
+ self.model_clip.eval();
31
+ self.model_clip.load_state_dict(torch.load('models/CLIP/rd64-uni-refined.pth', map_location=torch.device('cpu')), strict=False)
32
+
33
+ device = torch.device(self.plugin_options["devicename"])
34
+ self.model_clip.to(device)
35
+
36
+
37
+ def Run(self, img1, keywords:str) -> Frame:
38
+ if keywords is None or len(keywords) < 1 or img1 is None:
39
+ return img1
40
+
41
+ source_image_small = cv2.resize(img1, (256,256))
42
+
43
+ img_mask = np.full((source_image_small.shape[0],source_image_small.shape[1]), 0, dtype=np.float32)
44
+ mask_border = 1
45
+ l = 0
46
+ t = 0
47
+ r = 1
48
+ b = 1
49
+
50
+ mask_blur = 5
51
+ clip_blur = 5
52
+
53
+ img_mask = cv2.rectangle(img_mask, (mask_border+int(l), mask_border+int(t)),
54
+ (256 - mask_border-int(r), 256-mask_border-int(b)), (255, 255, 255), -1)
55
+ img_mask = cv2.GaussianBlur(img_mask, (mask_blur*2+1,mask_blur*2+1), 0)
56
+ img_mask /= 255
57
+
58
+
59
+ input_image = source_image_small
60
+
61
+ transform = transforms.Compose([
62
+ transforms.ToTensor(),
63
+ transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
64
+ transforms.Resize((256, 256)),
65
+ ])
66
+ img = transform(input_image).unsqueeze(0)
67
+
68
+ thresh = 0.5
69
+ prompts = keywords.split(',')
70
+ with THREAD_LOCK_CLIP:
71
+ with torch.no_grad():
72
+ preds = self.model_clip(img.repeat(len(prompts),1,1,1), prompts)[0]
73
+ clip_mask = torch.sigmoid(preds[0][0])
74
+ for i in range(len(prompts)-1):
75
+ clip_mask += torch.sigmoid(preds[i+1][0])
76
+
77
+ clip_mask = clip_mask.data.cpu().numpy()
78
+ np.clip(clip_mask, 0, 1)
79
+
80
+ clip_mask[clip_mask>thresh] = 1.0
81
+ clip_mask[clip_mask<=thresh] = 0.0
82
+ kernel = np.ones((5, 5), np.float32)
83
+ clip_mask = cv2.dilate(clip_mask, kernel, iterations=1)
84
+ clip_mask = cv2.GaussianBlur(clip_mask, (clip_blur*2+1,clip_blur*2+1), 0)
85
+
86
+ img_mask *= clip_mask
87
+ img_mask[img_mask<0.0] = 0.0
88
+ return img_mask
89
+
90
+
91
+
92
+ def Release(self):
93
+ self.model_clip = None
94
+
roop/processors/Mask_XSeg.py ADDED
@@ -0,0 +1,60 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import cv2
3
+ import onnxruntime
4
+ import threading
5
+ import roop.globals
6
+
7
+ from roop.typing import Frame
8
+ from roop.utilities import resolve_relative_path
9
+
10
+ THREAD_LOCK_CLIP = threading.Lock()
11
+
12
+
13
+ class Mask_XSeg():
14
+ plugin_options:dict = None
15
+
16
+ model_xseg = None
17
+
18
+ processorname = 'mask_xseg'
19
+ type = 'mask'
20
+
21
+
22
+ def Initialize(self, plugin_options:dict):
23
+ if self.plugin_options is not None:
24
+ if self.plugin_options["devicename"] != plugin_options["devicename"]:
25
+ self.Release()
26
+
27
+ self.plugin_options = plugin_options
28
+ if self.model_xseg is None:
29
+ model_path = resolve_relative_path('../models/xseg.onnx')
30
+ onnxruntime.set_default_logger_severity(3)
31
+ self.model_xseg = onnxruntime.InferenceSession(model_path, None, providers=roop.globals.execution_providers)
32
+ self.model_inputs = self.model_xseg.get_inputs()
33
+ self.model_outputs = self.model_xseg.get_outputs()
34
+
35
+ # replace Mac mps with cpu for the moment
36
+ self.devicename = self.plugin_options["devicename"].replace('mps', 'cpu')
37
+
38
+
39
+ def Run(self, img1, keywords:str) -> Frame:
40
+ temp_frame = cv2.resize(img1, (256, 256), cv2.INTER_CUBIC)
41
+ temp_frame = temp_frame.astype('float32') / 255.0
42
+ temp_frame = temp_frame[None, ...]
43
+ io_binding = self.model_xseg.io_binding()
44
+ io_binding.bind_cpu_input(self.model_inputs[0].name, temp_frame)
45
+ io_binding.bind_output(self.model_outputs[0].name, self.devicename)
46
+ self.model_xseg.run_with_iobinding(io_binding)
47
+ ort_outs = io_binding.copy_outputs_to_cpu()
48
+ result = ort_outs[0][0]
49
+ result = np.clip(result, 0, 1.0)
50
+ result[result < 0.1] = 0
51
+ # invert values to mask areas to keep
52
+ result = 1.0 - result
53
+ return result
54
+
55
+
56
+ def Release(self):
57
+ del self.model_xseg
58
+ self.model_xseg = None
59
+
60
+
roop/processors/__init__.py ADDED
File without changes
roop/template_parser.py ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import re
2
+ from datetime import datetime
3
+
4
+ template_functions = {
5
+ "timestamp": lambda data: str(int(datetime.now().timestamp())),
6
+ "i": lambda data: data.get("index", False),
7
+ "file": lambda data: data.get("file", False),
8
+ "date": lambda data: datetime.now().strftime("%Y-%m-%d"),
9
+ "time": lambda data: datetime.now().strftime("%H-%M-%S"),
10
+ }
11
+
12
+
13
+ def parse(text: str, data: dict):
14
+ pattern = r"\{([^}]+)\}"
15
+
16
+ matches = re.findall(pattern, text)
17
+
18
+ for match in matches:
19
+ replacement = template_functions[match](data)
20
+ if replacement is not False:
21
+ text = text.replace(f"{{{match}}}", replacement)
22
+
23
+ return text
roop/typing.py ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Any
2
+
3
+ from insightface.app.common import Face
4
+ from roop.FaceSet import FaceSet
5
+ import numpy
6
+
7
+ Face = Face
8
+ FaceSet = FaceSet
9
+ Frame = numpy.ndarray[Any, Any]
roop/util_ffmpeg.py ADDED
@@ -0,0 +1,110 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ import os
3
+ import subprocess
4
+ import roop.globals
5
+ import roop.utilities as util
6
+
7
+ from typing import List, Any
8
+
9
+ def run_ffmpeg(args: List[str]) -> bool:
10
+ commands = ['ffmpeg', '-hide_banner', '-hwaccel', 'auto', '-y', '-loglevel', roop.globals.log_level]
11
+ commands.extend(args)
12
+ print ("Running ffmpeg")
13
+ try:
14
+ subprocess.check_output(commands, stderr=subprocess.STDOUT)
15
+ return True
16
+ except Exception as e:
17
+ print("Running ffmpeg failed! Commandline:")
18
+ print (" ".join(commands))
19
+ return False
20
+
21
+
22
+
23
+ def cut_video(original_video: str, cut_video: str, start_frame: int, end_frame: int, reencode: bool):
24
+ fps = util.detect_fps(original_video)
25
+ start_time = start_frame / fps
26
+ num_frames = end_frame - start_frame
27
+
28
+ if reencode:
29
+ run_ffmpeg(['-ss', format(start_time, ".2f"), '-i', original_video, '-c:v', roop.globals.video_encoder, '-c:a', 'aac', '-frames:v', str(num_frames), cut_video])
30
+ else:
31
+ run_ffmpeg(['-ss', format(start_time, ".2f"), '-i', original_video, '-frames:v', str(num_frames), '-c:v' ,'copy','-c:a' ,'copy', cut_video])
32
+
33
+ def join_videos(videos: List[str], dest_filename: str, simple: bool):
34
+ if simple:
35
+ txtfilename = util.resolve_relative_path('../temp')
36
+ txtfilename = os.path.join(txtfilename, 'joinvids.txt')
37
+ with open(txtfilename, "w", encoding="utf-8") as f:
38
+ for v in videos:
39
+ v = v.replace('\\', '/')
40
+ f.write(f"file {v}\n")
41
+ commands = ['-f', 'concat', '-safe', '0', '-i', f'{txtfilename}', '-vcodec', 'copy', f'{dest_filename}']
42
+ run_ffmpeg(commands)
43
+
44
+ else:
45
+ inputs = []
46
+ filter = ''
47
+ for i,v in enumerate(videos):
48
+ inputs.append('-i')
49
+ inputs.append(v)
50
+ filter += f'[{i}:v:0][{i}:a:0]'
51
+ run_ffmpeg([" ".join(inputs), '-filter_complex', f'"{filter}concat=n={len(videos)}:v=1:a=1[outv][outa]"', '-map', '"[outv]"', '-map', '"[outa]"', dest_filename])
52
+
53
+ # filter += f'[{i}:v:0][{i}:a:0]'
54
+ # run_ffmpeg([" ".join(inputs), '-filter_complex', f'"{filter}concat=n={len(videos)}:v=1:a=1[outv][outa]"', '-map', '"[outv]"', '-map', '"[outa]"', dest_filename])
55
+
56
+
57
+
58
+ def extract_frames(target_path : str, trim_frame_start, trim_frame_end, fps : float) -> bool:
59
+ util.create_temp(target_path)
60
+ temp_directory_path = util.get_temp_directory_path(target_path)
61
+ commands = ['-i', target_path, '-q:v', '1', '-pix_fmt', 'rgb24', ]
62
+ if trim_frame_start is not None and trim_frame_end is not None:
63
+ commands.extend([ '-vf', 'trim=start_frame=' + str(trim_frame_start) + ':end_frame=' + str(trim_frame_end) + ',fps=' + str(fps) ])
64
+ commands.extend(['-vsync', '0', os.path.join(temp_directory_path, '%06d.' + roop.globals.CFG.output_image_format)])
65
+ return run_ffmpeg(commands)
66
+
67
+
68
+ def create_video(target_path: str, dest_filename: str, fps: float = 24.0, temp_directory_path: str = None) -> None:
69
+ if temp_directory_path is None:
70
+ temp_directory_path = util.get_temp_directory_path(target_path)
71
+ run_ffmpeg(['-r', str(fps), '-i', os.path.join(temp_directory_path, f'%06d.{roop.globals.CFG.output_image_format}'), '-c:v', roop.globals.video_encoder, '-crf', str(roop.globals.video_quality), '-pix_fmt', 'yuv420p', '-vf', 'colorspace=bt709:iall=bt601-6-625:fast=1', '-y', dest_filename])
72
+ return dest_filename
73
+
74
+
75
+ def create_gif_from_video(video_path: str, gif_path):
76
+ from roop.capturer import get_video_frame
77
+
78
+ fps = util.detect_fps(video_path)
79
+ frame = get_video_frame(video_path)
80
+
81
+ run_ffmpeg(['-i', video_path, '-vf', f'fps={fps},scale={frame.shape[0]}:-1:flags=lanczos,split[s0][s1];[s0]palettegen[p];[s1][p]paletteuse', '-loop', '0', gif_path])
82
+
83
+
84
+ def restore_audio(intermediate_video: str, original_video: str, trim_frame_start, trim_frame_end, final_video : str) -> None:
85
+ fps = util.detect_fps(original_video)
86
+ commands = [ '-i', intermediate_video ]
87
+ if trim_frame_start is None and trim_frame_end is None:
88
+ commands.extend([ '-c:a', 'copy' ])
89
+ else:
90
+ # if trim_frame_start is not None:
91
+ # start_time = trim_frame_start / fps
92
+ # commands.extend([ '-ss', format(start_time, ".2f")])
93
+ # else:
94
+ # commands.extend([ '-ss', '0' ])
95
+ # if trim_frame_end is not None:
96
+ # end_time = trim_frame_end / fps
97
+ # commands.extend([ '-to', format(end_time, ".2f")])
98
+ # commands.extend([ '-c:a', 'aac' ])
99
+ if trim_frame_start is not None:
100
+ start_time = trim_frame_start / fps
101
+ commands.extend([ '-ss', format(start_time, ".2f")])
102
+ else:
103
+ commands.extend([ '-ss', '0' ])
104
+ if trim_frame_end is not None:
105
+ end_time = trim_frame_end / fps
106
+ commands.extend([ '-to', format(end_time, ".2f")])
107
+ commands.extend([ '-i', original_video, "-c", "copy" ])
108
+
109
+ commands.extend([ '-map', '0:v:0', '-map', '1:a:0?', '-shortest', final_video ])
110
+ run_ffmpeg(commands)
roop/utilities.py ADDED
@@ -0,0 +1,339 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import glob
2
+ import mimetypes
3
+ import os
4
+ import platform
5
+ import shutil
6
+ import ssl
7
+ import subprocess
8
+ import sys
9
+ import urllib
10
+ import torch
11
+ import gradio
12
+ import tempfile
13
+ import cv2
14
+ import zipfile
15
+ import traceback
16
+
17
+ from pathlib import Path
18
+ from typing import List, Any
19
+ from tqdm import tqdm
20
+ from scipy.spatial import distance
21
+
22
+ import roop.template_parser as template_parser
23
+
24
+ import roop.globals
25
+
26
+ TEMP_FILE = "temp.mp4"
27
+ TEMP_DIRECTORY = "temp"
28
+
29
+ # monkey patch ssl for mac
30
+ if platform.system().lower() == "darwin":
31
+ ssl._create_default_https_context = ssl._create_unverified_context
32
+
33
+
34
+ # https://github.com/facefusion/facefusion/blob/master/facefusion
35
+ def detect_fps(target_path: str) -> float:
36
+ fps = 24.0
37
+ cap = cv2.VideoCapture(target_path)
38
+ if cap.isOpened():
39
+ fps = cap.get(cv2.CAP_PROP_FPS)
40
+ cap.release()
41
+ return fps
42
+
43
+
44
+ # Gradio wants Images in RGB
45
+ def convert_to_gradio(image):
46
+ if image is None:
47
+ return None
48
+ return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
49
+
50
+
51
+ def sort_filenames_ignore_path(filenames):
52
+ """Sorts a list of filenames containing a complete path by their filename,
53
+ while retaining their original path.
54
+
55
+ Args:
56
+ filenames: A list of filenames containing a complete path.
57
+
58
+ Returns:
59
+ A sorted list of filenames containing a complete path.
60
+ """
61
+ filename_path_tuples = [
62
+ (os.path.split(filename)[1], filename) for filename in filenames
63
+ ]
64
+ sorted_filename_path_tuples = sorted(filename_path_tuples, key=lambda x: x[0])
65
+ return [
66
+ filename_path_tuple[1] for filename_path_tuple in sorted_filename_path_tuples
67
+ ]
68
+
69
+
70
+ def sort_rename_frames(path: str):
71
+ filenames = os.listdir(path)
72
+ filenames.sort()
73
+ for i in range(len(filenames)):
74
+ of = os.path.join(path, filenames[i])
75
+ newidx = i + 1
76
+ new_filename = os.path.join(
77
+ path, f"{newidx:06d}." + roop.globals.CFG.output_image_format
78
+ )
79
+ os.rename(of, new_filename)
80
+
81
+
82
+ def get_temp_frame_paths(target_path: str) -> List[str]:
83
+ temp_directory_path = get_temp_directory_path(target_path)
84
+ return glob.glob(
85
+ (
86
+ os.path.join(
87
+ glob.escape(temp_directory_path),
88
+ f"*.{roop.globals.CFG.output_image_format}",
89
+ )
90
+ )
91
+ )
92
+
93
+
94
+ def get_temp_directory_path(target_path: str) -> str:
95
+ target_name, _ = os.path.splitext(os.path.basename(target_path))
96
+ target_directory_path = os.path.dirname(target_path)
97
+ return os.path.join(target_directory_path, TEMP_DIRECTORY, target_name)
98
+
99
+
100
+ def get_temp_output_path(target_path: str) -> str:
101
+ temp_directory_path = get_temp_directory_path(target_path)
102
+ return os.path.join(temp_directory_path, TEMP_FILE)
103
+
104
+
105
+ def normalize_output_path(source_path: str, target_path: str, output_path: str) -> Any:
106
+ if source_path and target_path:
107
+ source_name, _ = os.path.splitext(os.path.basename(source_path))
108
+ target_name, target_extension = os.path.splitext(os.path.basename(target_path))
109
+ if os.path.isdir(output_path):
110
+ return os.path.join(
111
+ output_path, source_name + "-" + target_name + target_extension
112
+ )
113
+ return output_path
114
+
115
+
116
+ def get_destfilename_from_path(
117
+ srcfilepath: str, destfilepath: str, extension: str
118
+ ) -> str:
119
+ fn, ext = os.path.splitext(os.path.basename(srcfilepath))
120
+ if "." in extension:
121
+ return os.path.join(destfilepath, f"{fn}{extension}")
122
+ return os.path.join(destfilepath, f"{fn}{extension}{ext}")
123
+
124
+
125
+ def replace_template(file_path: str, index: int = 0) -> str:
126
+ fn, ext = os.path.splitext(os.path.basename(file_path))
127
+
128
+ # Remove the "__temp" placeholder that was used as a temporary filename
129
+ fn = fn.replace("__temp", "")
130
+
131
+ template = roop.globals.CFG.output_template
132
+ replaced_filename = template_parser.parse(
133
+ template, {"index": str(index), "file": fn}
134
+ )
135
+
136
+ return os.path.join(roop.globals.output_path, f"{replaced_filename}{ext}")
137
+
138
+
139
+ def create_temp(target_path: str) -> None:
140
+ temp_directory_path = get_temp_directory_path(target_path)
141
+ Path(temp_directory_path).mkdir(parents=True, exist_ok=True)
142
+
143
+
144
+ def move_temp(target_path: str, output_path: str) -> None:
145
+ temp_output_path = get_temp_output_path(target_path)
146
+ if os.path.isfile(temp_output_path):
147
+ if os.path.isfile(output_path):
148
+ os.remove(output_path)
149
+ shutil.move(temp_output_path, output_path)
150
+
151
+
152
+ def clean_temp(target_path: str) -> None:
153
+ temp_directory_path = get_temp_directory_path(target_path)
154
+ parent_directory_path = os.path.dirname(temp_directory_path)
155
+ if not roop.globals.keep_frames and os.path.isdir(temp_directory_path):
156
+ shutil.rmtree(temp_directory_path)
157
+ if os.path.exists(parent_directory_path) and not os.listdir(parent_directory_path):
158
+ os.rmdir(parent_directory_path)
159
+
160
+
161
+ def delete_temp_frames(filename: str) -> None:
162
+ dir = os.path.dirname(os.path.dirname(filename))
163
+ shutil.rmtree(dir)
164
+
165
+
166
+ def has_image_extension(image_path: str) -> bool:
167
+ return image_path.lower().endswith(("png", "jpg", "jpeg", "webp"))
168
+
169
+
170
+ def has_extension(filepath: str, extensions: List[str]) -> bool:
171
+ return filepath.lower().endswith(tuple(extensions))
172
+
173
+
174
+ def is_image(image_path: str) -> bool:
175
+ if image_path and os.path.isfile(image_path):
176
+ mimetype, _ = mimetypes.guess_type(image_path)
177
+ return bool(mimetype and mimetype.startswith("image/"))
178
+ return False
179
+
180
+
181
+ def is_video(video_path: str) -> bool:
182
+ if video_path and os.path.isfile(video_path):
183
+ mimetype, _ = mimetypes.guess_type(video_path)
184
+ return bool(mimetype and mimetype.startswith("video/"))
185
+ return False
186
+
187
+
188
+ def conditional_download(download_directory_path: str, urls: List[str]) -> None:
189
+ if not os.path.exists(download_directory_path):
190
+ os.makedirs(download_directory_path)
191
+ for url in urls:
192
+ download_file_path = os.path.join(
193
+ download_directory_path, os.path.basename(url)
194
+ )
195
+ if not os.path.exists(download_file_path):
196
+ request = urllib.request.urlopen(url) # type: ignore[attr-defined]
197
+ total = int(request.headers.get("Content-Length", 0))
198
+ with tqdm(
199
+ total=total,
200
+ desc=f"Downloading {url}",
201
+ unit="B",
202
+ unit_scale=True,
203
+ unit_divisor=1024,
204
+ ) as progress:
205
+ urllib.request.urlretrieve(url, download_file_path, reporthook=lambda count, block_size, total_size: progress.update(block_size)) # type: ignore[attr-defined]
206
+
207
+
208
+ def get_local_files_from_folder(folder: str) -> List[str]:
209
+ if not os.path.exists(folder) or not os.path.isdir(folder):
210
+ return None
211
+ files = [
212
+ os.path.join(folder, f)
213
+ for f in os.listdir(folder)
214
+ if os.path.isfile(os.path.join(folder, f))
215
+ ]
216
+ return files
217
+
218
+
219
+ def resolve_relative_path(path: str) -> str:
220
+ return os.path.abspath(os.path.join(os.path.dirname(__file__), path))
221
+
222
+
223
+ def get_device() -> str:
224
+ if len(roop.globals.execution_providers) < 1:
225
+ roop.globals.execution_providers = ["CPUExecutionProvider"]
226
+
227
+ prov = roop.globals.execution_providers[0]
228
+ if "CoreMLExecutionProvider" in prov:
229
+ return "mps"
230
+ if "CUDAExecutionProvider" in prov or "ROCMExecutionProvider" in prov:
231
+ return "cuda"
232
+ if "OpenVINOExecutionProvider" in prov:
233
+ return "mkl"
234
+ return "cpu"
235
+
236
+
237
+ def str_to_class(module_name, class_name) -> Any:
238
+ from importlib import import_module
239
+
240
+ class_ = None
241
+ try:
242
+ module_ = import_module(module_name)
243
+ try:
244
+ class_ = getattr(module_, class_name)()
245
+ except AttributeError:
246
+ print(f"Class {class_name} does not exist")
247
+ except ImportError:
248
+ print(f"Module {module_name} does not exist")
249
+ return class_
250
+
251
+ def is_installed(name:str) -> bool:
252
+ return shutil.which(name);
253
+
254
+ # Taken from https://stackoverflow.com/a/68842705
255
+ def get_platform() -> str:
256
+ if sys.platform == "linux":
257
+ try:
258
+ proc_version = open("/proc/version").read()
259
+ if "Microsoft" in proc_version:
260
+ return "wsl"
261
+ except:
262
+ pass
263
+ return sys.platform
264
+
265
+ def open_with_default_app(filename:str):
266
+ if filename == None:
267
+ return
268
+ platform = get_platform()
269
+ if platform == "darwin":
270
+ subprocess.call(("open", filename))
271
+ elif platform in ["win64", "win32"]: os.startfile(filename.replace("/", "\\"))
272
+ elif platform == "wsl":
273
+ subprocess.call("cmd.exe /C start".split() + [filename])
274
+ else: # linux variants
275
+ subprocess.call("xdg-open", filename)
276
+
277
+
278
+ def prepare_for_batch(target_files) -> str:
279
+ print("Preparing temp files")
280
+ tempfolder = os.path.join(tempfile.gettempdir(), "rooptmp")
281
+ if os.path.exists(tempfolder):
282
+ shutil.rmtree(tempfolder)
283
+ Path(tempfolder).mkdir(parents=True, exist_ok=True)
284
+ for f in target_files:
285
+ newname = os.path.basename(f.name)
286
+ shutil.move(f.name, os.path.join(tempfolder, newname))
287
+ return tempfolder
288
+
289
+
290
+ def zip(files, zipname):
291
+ with zipfile.ZipFile(zipname, "w") as zip_file:
292
+ for f in files:
293
+ zip_file.write(f, os.path.basename(f))
294
+
295
+
296
+ def unzip(zipfilename: str, target_path: str):
297
+ with zipfile.ZipFile(zipfilename, "r") as zip_file:
298
+ zip_file.extractall(target_path)
299
+
300
+
301
+ def mkdir_with_umask(directory):
302
+ oldmask = os.umask(0)
303
+ # mode needs octal
304
+ os.makedirs(directory, mode=0o775, exist_ok=True)
305
+ os.umask(oldmask)
306
+
307
+
308
+ def open_folder(path: str):
309
+ platform = get_platform()
310
+ try:
311
+ if platform == "darwin":
312
+ subprocess.call(("open", path))
313
+ elif platform in ["win64", "win32"]:
314
+ open_with_default_app(path)
315
+ elif platform == "wsl":
316
+ subprocess.call("cmd.exe /C start".split() + [path])
317
+ else: # linux variants
318
+ subprocess.Popen(["xdg-open", path])
319
+ except Exception as e:
320
+ traceback.print_exc()
321
+ pass
322
+ # import webbrowser
323
+ # webbrowser.open(url)
324
+
325
+
326
+ def create_version_html() -> str:
327
+ python_version = ".".join([str(x) for x in sys.version_info[0:3]])
328
+ versions_html = f"""
329
+ python: <span title="{sys.version}">{python_version}</span>
330
+
331
+ torch: {getattr(torch, '__long_version__',torch.__version__)}
332
+
333
+ gradio: {gradio.__version__}
334
+ """
335
+ return versions_html
336
+
337
+
338
+ def compute_cosine_distance(emb1, emb2) -> float:
339
+ return distance.cosine(emb1, emb2)
roop/virtualcam.py ADDED
@@ -0,0 +1,88 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import cv2
2
+ import roop.globals
3
+ import ui.globals
4
+ import pyvirtualcam
5
+ import threading
6
+ import platform
7
+
8
+
9
+ cam_active = False
10
+ cam_thread = None
11
+ vcam = None
12
+
13
+ def virtualcamera(streamobs, cam_num,width,height):
14
+ from roop.ProcessOptions import ProcessOptions
15
+ from roop.core import live_swap, get_processing_plugins
16
+
17
+ global cam_active
18
+
19
+ #time.sleep(2)
20
+ print('Starting capture')
21
+ cap = cv2.VideoCapture(cam_num, cv2.CAP_DSHOW if platform.system() != 'Darwin' else cv2.CAP_AVFOUNDATION)
22
+ if not cap.isOpened():
23
+ print("Cannot open camera")
24
+ cap.release()
25
+ del cap
26
+ return
27
+
28
+ pref_width = width
29
+ pref_height = height
30
+ pref_fps_in = 30
31
+ cap.set(cv2.CAP_PROP_FRAME_WIDTH, pref_width)
32
+ cap.set(cv2.CAP_PROP_FRAME_HEIGHT, pref_height)
33
+ cap.set(cv2.CAP_PROP_FPS, pref_fps_in)
34
+ cam_active = True
35
+
36
+ # native format UYVY
37
+
38
+ cam = None
39
+ if streamobs:
40
+ print('Detecting virtual cam devices')
41
+ cam = pyvirtualcam.Camera(width=pref_width, height=pref_height, fps=pref_fps_in, fmt=pyvirtualcam.PixelFormat.BGR, print_fps=False)
42
+ if cam:
43
+ print(f'Using virtual camera: {cam.device}')
44
+ print(f'Using {cam.native_fmt}')
45
+ else:
46
+ print(f'Not streaming to virtual camera!')
47
+ subsample_size = roop.globals.subsample_size
48
+
49
+ # always use xseg masking
50
+ options = ProcessOptions(get_processing_plugins("mask_xseg"), roop.globals.distance_threshold, roop.globals.blend_ratio,
51
+ "all", 0, None, None, 1, subsample_size, False)
52
+ while cam_active:
53
+ ret, frame = cap.read()
54
+ if not ret:
55
+ break
56
+
57
+ if len(roop.globals.INPUT_FACESETS) > 0:
58
+ frame = live_swap(frame, options)
59
+ if cam:
60
+ cam.send(frame)
61
+ cam.sleep_until_next_frame()
62
+ ui.globals.ui_camera_frame = frame
63
+
64
+ if cam:
65
+ cam.close()
66
+ cap.release()
67
+ print('Camera stopped')
68
+
69
+
70
+
71
+ def start_virtual_cam(streamobs, cam_number, resolution):
72
+ global cam_thread, cam_active
73
+
74
+ if not cam_active:
75
+ width, height = map(int, resolution.split('x'))
76
+ cam_thread = threading.Thread(target=virtualcamera, args=[streamobs, cam_number, width, height])
77
+ cam_thread.start()
78
+
79
+
80
+
81
+ def stop_virtual_cam():
82
+ global cam_active, cam_thread
83
+
84
+ if cam_active:
85
+ cam_active = False
86
+ cam_thread.join()
87
+
88
+
roop/vr_util.py ADDED
@@ -0,0 +1,57 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import cv2
2
+ import numpy as np
3
+
4
+ # VR Lense Distortion
5
+ # Taken from https://github.com/g0kuvonlange/vrswap
6
+
7
+
8
+ def get_perspective(img, FOV, THETA, PHI, height, width):
9
+ #
10
+ # THETA is left/right angle, PHI is up/down angle, both in degree
11
+ #
12
+ [orig_width, orig_height, _] = img.shape
13
+ equ_h = orig_height
14
+ equ_w = orig_width
15
+ equ_cx = (equ_w - 1) / 2.0
16
+ equ_cy = (equ_h - 1) / 2.0
17
+
18
+ wFOV = FOV
19
+ hFOV = float(height) / width * wFOV
20
+
21
+ w_len = np.tan(np.radians(wFOV / 2.0))
22
+ h_len = np.tan(np.radians(hFOV / 2.0))
23
+
24
+ x_map = np.ones([height, width], np.float32)
25
+ y_map = np.tile(np.linspace(-w_len, w_len, width), [height, 1])
26
+ z_map = -np.tile(np.linspace(-h_len, h_len, height), [width, 1]).T
27
+
28
+ D = np.sqrt(x_map**2 + y_map**2 + z_map**2)
29
+ xyz = np.stack((x_map, y_map, z_map), axis=2) / np.repeat(
30
+ D[:, :, np.newaxis], 3, axis=2
31
+ )
32
+
33
+ y_axis = np.array([0.0, 1.0, 0.0], np.float32)
34
+ z_axis = np.array([0.0, 0.0, 1.0], np.float32)
35
+ [R1, _] = cv2.Rodrigues(z_axis * np.radians(THETA))
36
+ [R2, _] = cv2.Rodrigues(np.dot(R1, y_axis) * np.radians(-PHI))
37
+
38
+ xyz = xyz.reshape([height * width, 3]).T
39
+ xyz = np.dot(R1, xyz)
40
+ xyz = np.dot(R2, xyz).T
41
+ lat = np.arcsin(xyz[:, 2])
42
+ lon = np.arctan2(xyz[:, 1], xyz[:, 0])
43
+
44
+ lon = lon.reshape([height, width]) / np.pi * 180
45
+ lat = -lat.reshape([height, width]) / np.pi * 180
46
+
47
+ lon = lon / 180 * equ_cx + equ_cx
48
+ lat = lat / 90 * equ_cy + equ_cy
49
+
50
+ persp = cv2.remap(
51
+ img,
52
+ lon.astype(np.float32),
53
+ lat.astype(np.float32),
54
+ cv2.INTER_CUBIC,
55
+ borderMode=cv2.BORDER_WRAP,
56
+ )
57
+ return persp