assile commited on
Commit
b21f0f2
·
verified ·
1 Parent(s): b8be4c1

Delete roop

Browse files
roop/FaceSet.py DELETED
@@ -1,20 +0,0 @@
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 DELETED
@@ -1,7 +0,0 @@
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 DELETED
@@ -1,911 +0,0 @@
1
- import os
2
- import cv2
3
- import numpy as np
4
- import psutil
5
-
6
- from roop.ProcessOptions import ProcessOptions
7
-
8
- from roop.face_util import get_first_face, get_all_faces, rotate_anticlockwise, rotate_clockwise, clamp_cut_values
9
- from roop.utilities import compute_cosine_distance, get_device, str_to_class, shuffle_array
10
- import roop.vr_util as vr
11
-
12
- from typing import Any, List, Callable
13
- from roop.typing import Frame, Face
14
- from concurrent.futures import ThreadPoolExecutor, as_completed
15
- from threading import Thread, Lock
16
- from queue import Queue
17
- from tqdm import tqdm
18
- from roop.ffmpeg_writer import FFMPEG_VideoWriter
19
- from roop.StreamWriter import StreamWriter
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
- USE_LAST_SWAPPED = 4
31
-
32
-
33
-
34
- def create_queue(temp_frame_paths: List[str]) -> Queue[str]:
35
- queue: Queue[str] = Queue()
36
- for frame_path in temp_frame_paths:
37
- queue.put(frame_path)
38
- return queue
39
-
40
-
41
- def pick_queue(queue: Queue[str], queue_per_future: int) -> List[str]:
42
- queues = []
43
- for _ in range(queue_per_future):
44
- if not queue.empty():
45
- queues.append(queue.get())
46
- return queues
47
-
48
-
49
-
50
- class ProcessMgr():
51
- input_face_datas = []
52
- target_face_datas = []
53
-
54
- imagemask = None
55
-
56
- processors = []
57
- options : ProcessOptions = None
58
-
59
- num_threads = 1
60
- current_index = 0
61
- processing_threads = 1
62
- buffer_wait_time = 0.1
63
-
64
- lock = Lock()
65
-
66
- frames_queue = None
67
- processed_queue = None
68
-
69
- videowriter= None
70
- streamwriter = None
71
-
72
- progress_gradio = None
73
- total_frames = 0
74
-
75
- num_frames_no_face = 0
76
- last_swapped_frame = None
77
-
78
- output_to_file = None
79
- output_to_cam = None
80
-
81
-
82
- plugins = {
83
- 'faceswap' : 'FaceSwapInsightFace',
84
- 'mask_clip2seg' : 'Mask_Clip2Seg',
85
- 'mask_xseg' : 'Mask_XSeg',
86
- 'codeformer' : 'Enhance_CodeFormer',
87
- 'gfpgan' : 'Enhance_GFPGAN',
88
- 'dmdnet' : 'Enhance_DMDNet',
89
- 'gpen' : 'Enhance_GPEN',
90
- 'restoreformer++' : 'Enhance_RestoreFormerPPlus',
91
- 'colorizer' : 'Frame_Colorizer',
92
- 'filter_generic' : 'Frame_Filter',
93
- 'removebg' : 'Frame_Masking',
94
- 'upscale' : 'Frame_Upscale'
95
- }
96
-
97
- def __init__(self, progress):
98
- if progress is not None:
99
- self.progress_gradio = progress
100
-
101
- def reuseOldProcessor(self, name:str):
102
- for p in self.processors:
103
- if p.processorname == name:
104
- return p
105
-
106
- return None
107
-
108
-
109
- def initialize(self, input_faces, target_faces, options):
110
- self.input_face_datas = input_faces
111
- self.target_face_datas = target_faces
112
- self.num_frames_no_face = 0
113
- self.last_swapped_frame = None
114
- self.options = options
115
- devicename = get_device()
116
-
117
- roop.globals.g_desired_face_analysis=["landmark_3d_68", "landmark_2d_106","detection","recognition"]
118
- if options.swap_mode == "all_female" or options.swap_mode == "all_male":
119
- roop.globals.g_desired_face_analysis.append("genderage")
120
- elif options.swap_mode == "all_random":
121
- # don't modify original list
122
- self.input_face_datas = input_faces.copy()
123
- shuffle_array(self.input_face_datas)
124
-
125
-
126
- for p in self.processors:
127
- newp = next((x for x in options.processors.keys() if x == p.processorname), None)
128
- if newp is None:
129
- p.Release()
130
- del p
131
-
132
- newprocessors = []
133
- for key, extoption in options.processors.items():
134
- p = self.reuseOldProcessor(key)
135
- if p is None:
136
- classname = self.plugins[key]
137
- module = 'roop.processors.' + classname
138
- p = str_to_class(module, classname)
139
- if p is not None:
140
- extoption.update({"devicename": devicename})
141
- if p.type == "swap":
142
- if self.options.swap_modelname == "InSwapper 128":
143
- extoption.update({"modelname": "inswapper_128.onnx"})
144
- elif self.options.swap_modelname == "ReSwapper 128":
145
- extoption.update({"modelname": "reswapper_128.onnx"})
146
- elif self.options.swap_modelname == "ReSwapper 256":
147
- extoption.update({"modelname": "reswapper_256.onnx"})
148
-
149
- p.Initialize(extoption)
150
- newprocessors.append(p)
151
- else:
152
- print(f"Not using {module}")
153
- self.processors = newprocessors
154
-
155
-
156
-
157
- if isinstance(self.options.imagemask, dict) and self.options.imagemask.get("layers") and len(self.options.imagemask["layers"]) > 0:
158
- self.options.imagemask = self.options.imagemask.get("layers")[0]
159
- # Get rid of alpha
160
- self.options.imagemask = cv2.cvtColor(self.options.imagemask, cv2.COLOR_RGBA2GRAY)
161
- if np.any(self.options.imagemask):
162
- mo = self.input_face_datas[0].faces[0].mask_offsets
163
- self.options.imagemask = self.blur_area(self.options.imagemask, mo[4], mo[5])
164
- self.options.imagemask = self.options.imagemask.astype(np.float32) / 255
165
- self.options.imagemask = cv2.cvtColor(self.options.imagemask, cv2.COLOR_GRAY2RGB)
166
- else:
167
- self.options.imagemask = None
168
-
169
- self.options.frame_processing = False
170
- for p in self.processors:
171
- if p.type.startswith("frame_"):
172
- self.options.frame_processing = True
173
-
174
-
175
-
176
-
177
-
178
-
179
- def run_batch(self, source_files, target_files, threads:int = 1):
180
- progress_bar_format = '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]'
181
- self.total_frames = len(source_files)
182
- self.num_threads = threads
183
- with tqdm(total=self.total_frames, desc='Processing', unit='frame', dynamic_ncols=True, bar_format=progress_bar_format) as progress:
184
- with ThreadPoolExecutor(max_workers=threads) as executor:
185
- futures = []
186
- queue = create_queue(source_files)
187
- queue_per_future = max(len(source_files) // threads, 1)
188
- while not queue.empty():
189
- future = executor.submit(self.process_frames, source_files, target_files, pick_queue(queue, queue_per_future), lambda: self.update_progress(progress))
190
- futures.append(future)
191
- for future in as_completed(futures):
192
- future.result()
193
-
194
-
195
- def process_frames(self, source_files: List[str], target_files: List[str], current_files, update: Callable[[], None]) -> None:
196
- for f in current_files:
197
- if not roop.globals.processing:
198
- return
199
-
200
- # Decode the byte array into an OpenCV image
201
- temp_frame = cv2.imdecode(np.fromfile(f, dtype=np.uint8), cv2.IMREAD_COLOR)
202
- if temp_frame is not None:
203
- if self.options.frame_processing:
204
- for p in self.processors:
205
- frame = p.Run(temp_frame)
206
- resimg = frame
207
- else:
208
- resimg = self.process_frame(temp_frame)
209
- if resimg is not None:
210
- i = source_files.index(f)
211
- # Also let numpy write the file to support utf-8/16 filenames
212
- cv2.imencode(f'.{roop.globals.CFG.output_image_format}',resimg)[1].tofile(target_files[i])
213
- if update:
214
- update()
215
-
216
-
217
-
218
- def read_frames_thread(self, cap, frame_start, frame_end, num_threads):
219
- num_frame = 0
220
- total_num = frame_end - frame_start
221
- if frame_start > 0:
222
- cap.set(cv2.CAP_PROP_POS_FRAMES,frame_start)
223
-
224
- while True and roop.globals.processing:
225
- ret, frame = cap.read()
226
- if not ret:
227
- break
228
-
229
- self.frames_queue[num_frame % num_threads].put(frame, block=True)
230
- num_frame += 1
231
- if num_frame == total_num:
232
- break
233
-
234
- for i in range(num_threads):
235
- self.frames_queue[i].put(None)
236
-
237
-
238
-
239
- def process_videoframes(self, threadindex, progress) -> None:
240
- while True:
241
- frame = self.frames_queue[threadindex].get()
242
- if frame is None:
243
- self.processing_threads -= 1
244
- self.processed_queue[threadindex].put((False, None))
245
- return
246
- else:
247
- if self.options.frame_processing:
248
- for p in self.processors:
249
- frame = p.Run(frame)
250
- resimg = frame
251
- else:
252
- resimg = self.process_frame(frame)
253
- self.processed_queue[threadindex].put((True, resimg))
254
- del frame
255
- progress()
256
-
257
-
258
- def write_frames_thread(self):
259
- nextindex = 0
260
- num_producers = self.num_threads
261
-
262
- while True:
263
- process, frame = self.processed_queue[nextindex % self.num_threads].get()
264
- nextindex += 1
265
- if frame is not None:
266
- if self.output_to_file:
267
- self.videowriter.write_frame(frame)
268
- if self.output_to_cam:
269
- self.streamwriter.WriteToStream(frame)
270
- del frame
271
- elif process == False:
272
- num_producers -= 1
273
- if num_producers < 1:
274
- return
275
-
276
-
277
-
278
- def run_batch_inmem(self, output_method, source_video, target_video, frame_start, frame_end, fps, threads:int = 1):
279
- if len(self.processors) < 1:
280
- print("No processor defined!")
281
- return
282
-
283
- cap = cv2.VideoCapture(source_video)
284
- # frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
285
- frame_count = (frame_end - frame_start) + 1
286
- width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
287
- height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
288
-
289
- processed_resolution = None
290
- for p in self.processors:
291
- if hasattr(p, 'getProcessedResolution'):
292
- processed_resolution = p.getProcessedResolution(width, height)
293
- print(f"Processed resolution: {processed_resolution}")
294
- if processed_resolution is not None:
295
- width = processed_resolution[0]
296
- height = processed_resolution[1]
297
-
298
-
299
- self.total_frames = frame_count
300
- self.num_threads = threads
301
-
302
- self.processing_threads = self.num_threads
303
- self.frames_queue = []
304
- self.processed_queue = []
305
- for _ in range(threads):
306
- self.frames_queue.append(Queue(1))
307
- self.processed_queue.append(Queue(1))
308
-
309
- self.output_to_file = output_method != "Virtual Camera"
310
- self.output_to_cam = output_method == "Virtual Camera" or output_method == "Both"
311
-
312
- if self.output_to_file:
313
- self.videowriter = FFMPEG_VideoWriter(target_video, (width, height), fps, codec=roop.globals.video_encoder, crf=roop.globals.video_quality, audiofile=None)
314
- if self.output_to_cam:
315
- self.streamwriter = StreamWriter((width, height), int(fps))
316
-
317
- readthread = Thread(target=self.read_frames_thread, args=(cap, frame_start, frame_end, threads))
318
- readthread.start()
319
-
320
- writethread = Thread(target=self.write_frames_thread)
321
- writethread.start()
322
-
323
- progress_bar_format = '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]'
324
- with tqdm(total=self.total_frames, desc='Processing', unit='frames', dynamic_ncols=True, bar_format=progress_bar_format) as progress:
325
- with ThreadPoolExecutor(thread_name_prefix='swap_proc', max_workers=self.num_threads) as executor:
326
- futures = []
327
-
328
- for threadindex in range(threads):
329
- future = executor.submit(self.process_videoframes, threadindex, lambda: self.update_progress(progress))
330
- futures.append(future)
331
-
332
- for future in as_completed(futures):
333
- future.result()
334
- # wait for the task to complete
335
- readthread.join()
336
- writethread.join()
337
- cap.release()
338
- if self.output_to_file:
339
- self.videowriter.close()
340
- if self.output_to_cam:
341
- self.streamwriter.Close()
342
-
343
- self.frames_queue.clear()
344
- self.processed_queue.clear()
345
-
346
-
347
-
348
-
349
- def update_progress(self, progress: Any = None) -> None:
350
- process = psutil.Process(os.getpid())
351
- memory_usage = process.memory_info().rss / 1024 / 1024 / 1024
352
- progress.set_postfix({
353
- 'memory_usage': '{:.2f}'.format(memory_usage).zfill(5) + 'GB',
354
- 'execution_threads': self.num_threads
355
- })
356
- progress.update(1)
357
- if self.progress_gradio is not None:
358
- self.progress_gradio((progress.n, self.total_frames), desc='Processing', total=self.total_frames, unit='frames')
359
-
360
-
361
-
362
- def process_frame(self, frame:Frame):
363
- if len(self.input_face_datas) < 1 and not self.options.show_face_masking:
364
- return frame
365
- temp_frame = frame.copy()
366
- num_swapped, temp_frame = self.swap_faces(frame, temp_frame)
367
- if num_swapped > 0:
368
- if roop.globals.no_face_action == eNoFaceAction.SKIP_FRAME_IF_DISSIMILAR:
369
- if len(self.input_face_datas) > num_swapped:
370
- return None
371
- self.num_frames_no_face = 0
372
- self.last_swapped_frame = temp_frame.copy()
373
- return temp_frame
374
- if roop.globals.no_face_action == eNoFaceAction.USE_LAST_SWAPPED:
375
- if self.last_swapped_frame is not None and self.num_frames_no_face < self.options.max_num_reuse_frame:
376
- self.num_frames_no_face += 1
377
- return self.last_swapped_frame.copy()
378
- return frame
379
-
380
- elif roop.globals.no_face_action == eNoFaceAction.USE_ORIGINAL_FRAME:
381
- return frame
382
- if roop.globals.no_face_action == eNoFaceAction.SKIP_FRAME:
383
- #This only works with in-mem processing, as it simply skips the frame.
384
- #For 'extract frames' it simply leaves the unprocessed frame unprocessed and it gets used in the final output by ffmpeg.
385
- #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?????
386
- #alternatively, it could mark all the necessary frames for deletion, delete them at the end, then rename the remaining frames that might work?
387
- return None
388
- else:
389
- return self.retry_rotated(frame)
390
-
391
- def retry_rotated(self, frame):
392
- copyframe = frame.copy()
393
- copyframe = rotate_clockwise(copyframe)
394
- temp_frame = copyframe.copy()
395
- num_swapped, temp_frame = self.swap_faces(copyframe, temp_frame)
396
- if num_swapped > 0:
397
- return rotate_anticlockwise(temp_frame)
398
-
399
- copyframe = frame.copy()
400
- copyframe = rotate_anticlockwise(copyframe)
401
- temp_frame = copyframe.copy()
402
- num_swapped, temp_frame = self.swap_faces(copyframe, temp_frame)
403
- if num_swapped > 0:
404
- return rotate_clockwise(temp_frame)
405
- del copyframe
406
- return frame
407
-
408
-
409
-
410
- def swap_faces(self, frame, temp_frame):
411
- num_faces_found = 0
412
-
413
- if self.options.swap_mode == "first":
414
- face = get_first_face(frame)
415
-
416
- if face is None:
417
- return num_faces_found, frame
418
-
419
- num_faces_found += 1
420
- temp_frame = self.process_face(self.options.selected_index, face, temp_frame)
421
- del face
422
-
423
- else:
424
- faces = get_all_faces(frame)
425
- if faces is None:
426
- return num_faces_found, frame
427
-
428
- if self.options.swap_mode == "all":
429
- for face in faces:
430
- num_faces_found += 1
431
- temp_frame = self.process_face(self.options.selected_index, face, temp_frame)
432
-
433
- elif self.options.swap_mode == "all_input" or self.options.swap_mode == "all_random":
434
- for i,face in enumerate(faces):
435
- num_faces_found += 1
436
- if i < len(self.input_face_datas):
437
- temp_frame = self.process_face(i, face, temp_frame)
438
- else:
439
- break
440
-
441
- elif self.options.swap_mode == "selected":
442
- num_targetfaces = len(self.target_face_datas)
443
- use_index = num_targetfaces == 1
444
- for i,tf in enumerate(self.target_face_datas):
445
- for face in faces:
446
- if compute_cosine_distance(tf.embedding, face.embedding) <= self.options.face_distance_threshold:
447
- if i < len(self.input_face_datas):
448
- if use_index:
449
- temp_frame = self.process_face(self.options.selected_index, face, temp_frame)
450
- else:
451
- temp_frame = self.process_face(i, face, temp_frame)
452
- num_faces_found += 1
453
- if not roop.globals.vr_mode and num_faces_found == num_targetfaces:
454
- break
455
- elif self.options.swap_mode == "all_female" or self.options.swap_mode == "all_male":
456
- gender = 'F' if self.options.swap_mode == "all_female" else 'M'
457
- for face in faces:
458
- if face.sex == gender:
459
- num_faces_found += 1
460
- temp_frame = self.process_face(self.options.selected_index, face, temp_frame)
461
-
462
- # might be slower but way more clean to release everything here
463
- for face in faces:
464
- del face
465
- faces.clear()
466
-
467
-
468
-
469
- if roop.globals.vr_mode and num_faces_found % 2 > 0:
470
- # stereo image, there has to be an even number of faces
471
- num_faces_found = 0
472
- return num_faces_found, frame
473
- if num_faces_found == 0:
474
- return num_faces_found, frame
475
-
476
- #maskprocessor = next((x for x in self.processors if x.type == 'mask'), None)
477
-
478
- if self.options.imagemask is not None and self.options.imagemask.shape == frame.shape:
479
- temp_frame = self.simple_blend_with_mask(temp_frame, frame, self.options.imagemask)
480
- return num_faces_found, temp_frame
481
-
482
-
483
- def rotation_action(self, original_face:Face, frame:Frame):
484
- (height, width) = frame.shape[:2]
485
-
486
- bounding_box_width = original_face.bbox[2] - original_face.bbox[0]
487
- bounding_box_height = original_face.bbox[3] - original_face.bbox[1]
488
- horizontal_face = bounding_box_width > bounding_box_height
489
-
490
- center_x = width // 2.0
491
- start_x = original_face.bbox[0]
492
- end_x = original_face.bbox[2]
493
- bbox_center_x = start_x + (bounding_box_width // 2.0)
494
-
495
- # need to leverage the array of landmarks as decribed here:
496
- # https://github.com/deepinsight/insightface/tree/master/alignment/coordinate_reg
497
- # basically, we should be able to check for the relative position of eyes and nose
498
- # then use that to determine which way the face is actually facing when in a horizontal position
499
- # and use that to determine the correct rotation_action
500
-
501
- forehead_x = original_face.landmark_2d_106[72][0]
502
- chin_x = original_face.landmark_2d_106[0][0]
503
-
504
- if horizontal_face:
505
- if chin_x < forehead_x:
506
- # this is someone lying down with their face like this (:
507
- return "rotate_anticlockwise"
508
- elif forehead_x < chin_x:
509
- # this is someone lying down with their face like this :)
510
- return "rotate_clockwise"
511
- if bbox_center_x >= center_x:
512
- # this is someone lying down with their face in the right hand side of the frame
513
- return "rotate_anticlockwise"
514
- if bbox_center_x < center_x:
515
- # this is someone lying down with their face in the left hand side of the frame
516
- return "rotate_clockwise"
517
-
518
- return None
519
-
520
-
521
- def auto_rotate_frame(self, original_face, frame:Frame):
522
- target_face = original_face
523
- original_frame = frame
524
-
525
- rotation_action = self.rotation_action(original_face, frame)
526
-
527
- if rotation_action == "rotate_anticlockwise":
528
- #face is horizontal, rotating frame anti-clockwise and getting face bounding box from rotated frame
529
- frame = rotate_anticlockwise(frame)
530
- elif rotation_action == "rotate_clockwise":
531
- #face is horizontal, rotating frame clockwise and getting face bounding box from rotated frame
532
- frame = rotate_clockwise(frame)
533
-
534
- return target_face, frame, rotation_action
535
-
536
-
537
- def auto_unrotate_frame(self, frame:Frame, rotation_action):
538
- if rotation_action == "rotate_anticlockwise":
539
- return rotate_clockwise(frame)
540
- elif rotation_action == "rotate_clockwise":
541
- return rotate_anticlockwise(frame)
542
-
543
- return frame
544
-
545
-
546
-
547
- def process_face(self,face_index, target_face:Face, frame:Frame):
548
- from roop.face_util import align_crop
549
-
550
- enhanced_frame = None
551
- if(len(self.input_face_datas) > 0):
552
- inputface = self.input_face_datas[face_index].faces[0]
553
- else:
554
- inputface = None
555
-
556
- rotation_action = None
557
- if roop.globals.autorotate_faces:
558
- # check for sideways rotation of face
559
- rotation_action = self.rotation_action(target_face, frame)
560
- if rotation_action is not None:
561
- (startX, startY, endX, endY) = target_face["bbox"].astype("int")
562
- width = endX - startX
563
- height = endY - startY
564
- offs = int(max(width,height) * 0.25)
565
- rotcutframe,startX, startY, endX, endY = self.cutout(frame, startX - offs, startY - offs, endX + offs, endY + offs)
566
- if rotation_action == "rotate_anticlockwise":
567
- rotcutframe = rotate_anticlockwise(rotcutframe)
568
- elif rotation_action == "rotate_clockwise":
569
- rotcutframe = rotate_clockwise(rotcutframe)
570
- # rotate image and re-detect face to correct wonky landmarks
571
- rotface = get_first_face(rotcutframe)
572
- if rotface is None:
573
- rotation_action = None
574
- else:
575
- saved_frame = frame.copy()
576
- frame = rotcutframe
577
- target_face = rotface
578
-
579
-
580
-
581
- # if roop.globals.vr_mode:
582
- # bbox = target_face.bbox
583
- # [orig_width, orig_height, _] = frame.shape
584
-
585
- # # Convert bounding box to ints
586
- # x1, y1, x2, y2 = map(int, bbox)
587
-
588
- # # Determine the center of the bounding box
589
- # x_center = (x1 + x2) / 2
590
- # y_center = (y1 + y2) / 2
591
-
592
- # # Normalize coordinates to range [-1, 1]
593
- # x_center_normalized = x_center / (orig_width / 2) - 1
594
- # y_center_normalized = y_center / (orig_width / 2) - 1
595
-
596
- # # Convert normalized coordinates to spherical (theta, phi)
597
- # theta = x_center_normalized * 180 # Theta ranges from -180 to 180 degrees
598
- # phi = -y_center_normalized * 90 # Phi ranges from -90 to 90 degrees
599
-
600
- # img = vr.GetPerspective(frame, 90, theta, phi, 1280, 1280) # Generate perspective image
601
-
602
-
603
- """ Code ported/adapted from Facefusion which borrowed the idea from Rope:
604
- Kind of subsampling the cutout and aligned face image and faceswapping slices of it up to
605
- the desired output resolution. This works around the current resolution limitations without using enhancers.
606
- """
607
- model_output_size = self.options.swap_output_size
608
- subsample_size = max(self.options.subsample_size, model_output_size)
609
- subsample_total = subsample_size // model_output_size
610
- aligned_img, M = align_crop(frame, target_face.kps, subsample_size)
611
-
612
- fake_frame = aligned_img
613
- target_face.matrix = M
614
-
615
- for p in self.processors:
616
- if p.type == 'swap':
617
- swap_result_frames = []
618
- subsample_frames = self.implode_pixel_boost(aligned_img, model_output_size, subsample_total)
619
- for sliced_frame in subsample_frames:
620
- for _ in range(0,self.options.num_swap_steps):
621
- sliced_frame = self.prepare_crop_frame(sliced_frame)
622
- sliced_frame = p.Run(inputface, target_face, sliced_frame)
623
- sliced_frame = self.normalize_swap_frame(sliced_frame)
624
- swap_result_frames.append(sliced_frame)
625
- fake_frame = self.explode_pixel_boost(swap_result_frames, model_output_size, subsample_total, subsample_size)
626
- fake_frame = fake_frame.astype(np.uint8)
627
- scale_factor = 0.0
628
- elif p.type == 'mask':
629
- fake_frame = self.process_mask(p, aligned_img, fake_frame)
630
- else:
631
- enhanced_frame, scale_factor = p.Run(self.input_face_datas[face_index], target_face, fake_frame)
632
-
633
- upscale = 512
634
- orig_width = fake_frame.shape[1]
635
- if orig_width != upscale:
636
- fake_frame = cv2.resize(fake_frame, (upscale, upscale), cv2.INTER_CUBIC)
637
- mask_offsets = (0,0,0,0,1,20) if inputface is None else inputface.mask_offsets
638
-
639
-
640
- if enhanced_frame is None:
641
- scale_factor = int(upscale / orig_width)
642
- result = self.paste_upscale(fake_frame, fake_frame, target_face.matrix, frame, scale_factor, mask_offsets)
643
- else:
644
- result = self.paste_upscale(fake_frame, enhanced_frame, target_face.matrix, frame, scale_factor, mask_offsets)
645
-
646
- # Restore mouth before unrotating
647
- if self.options.restore_original_mouth:
648
- mouth_cutout, mouth_bb = self.create_mouth_mask(target_face, frame)
649
- result = self.apply_mouth_area(result, mouth_cutout, mouth_bb)
650
-
651
- if rotation_action is not None:
652
- fake_frame = self.auto_unrotate_frame(result, rotation_action)
653
- result = self.paste_simple(fake_frame, saved_frame, startX, startY)
654
-
655
- return result
656
-
657
-
658
-
659
-
660
- def cutout(self, frame:Frame, start_x, start_y, end_x, end_y):
661
- if start_x < 0:
662
- start_x = 0
663
- if start_y < 0:
664
- start_y = 0
665
- if end_x > frame.shape[1]:
666
- end_x = frame.shape[1]
667
- if end_y > frame.shape[0]:
668
- end_y = frame.shape[0]
669
- return frame[start_y:end_y, start_x:end_x], start_x, start_y, end_x, end_y
670
-
671
- def paste_simple(self, src:Frame, dest:Frame, start_x, start_y):
672
- end_x = start_x + src.shape[1]
673
- end_y = start_y + src.shape[0]
674
-
675
- start_x, end_x, start_y, end_y = clamp_cut_values(start_x, end_x, start_y, end_y, dest)
676
- dest[start_y:end_y, start_x:end_x] = src
677
- return dest
678
-
679
- def simple_blend_with_mask(self, image1, image2, mask):
680
- # Blend the images
681
- blended_image = image1.astype(np.float32) * (1.0 - mask) + image2.astype(np.float32) * mask
682
- return blended_image.astype(np.uint8)
683
-
684
-
685
- def paste_upscale(self, fake_face, upsk_face, M, target_img, scale_factor, mask_offsets):
686
- M_scale = M * scale_factor
687
- IM = cv2.invertAffineTransform(M_scale)
688
-
689
- face_matte = np.full((target_img.shape[0],target_img.shape[1]), 255, dtype=np.uint8)
690
- # Generate white square sized as a upsk_face
691
- img_matte = np.zeros((upsk_face.shape[0],upsk_face.shape[1]), dtype=np.uint8)
692
-
693
- w = img_matte.shape[1]
694
- h = img_matte.shape[0]
695
-
696
- top = int(mask_offsets[0] * h)
697
- bottom = int(h - (mask_offsets[1] * h))
698
- left = int(mask_offsets[2] * w)
699
- right = int(w - (mask_offsets[3] * w))
700
- img_matte[top:bottom,left:right] = 255
701
-
702
- # Transform white square back to target_img
703
- img_matte = cv2.warpAffine(img_matte, IM, (target_img.shape[1], target_img.shape[0]), flags=cv2.INTER_NEAREST, borderValue=0.0)
704
- ##Blacken the edges of face_matte by 1 pixels (so the mask in not expanded on the image edges)
705
- img_matte[:1,:] = img_matte[-1:,:] = img_matte[:,:1] = img_matte[:,-1:] = 0
706
-
707
- img_matte = self.blur_area(img_matte, mask_offsets[4], mask_offsets[5])
708
- #Normalize images to float values and reshape
709
- img_matte = img_matte.astype(np.float32)/255
710
- face_matte = face_matte.astype(np.float32)/255
711
- img_matte = np.minimum(face_matte, img_matte)
712
- if self.options.show_face_area_overlay:
713
- # Additional steps for green overlay
714
- green_overlay = np.zeros_like(target_img)
715
- green_color = [0, 255, 0] # RGB for green
716
- for i in range(3): # Apply green color where img_matte is not zero
717
- green_overlay[:, :, i] = np.where(img_matte > 0, green_color[i], 0) ##Transform upcaled face back to target_img
718
- img_matte = np.reshape(img_matte, [img_matte.shape[0],img_matte.shape[1],1])
719
- paste_face = cv2.warpAffine(upsk_face, IM, (target_img.shape[1], target_img.shape[0]), borderMode=cv2.BORDER_REPLICATE)
720
- if upsk_face is not fake_face:
721
- fake_face = cv2.warpAffine(fake_face, IM, (target_img.shape[1], target_img.shape[0]), borderMode=cv2.BORDER_REPLICATE)
722
- paste_face = cv2.addWeighted(paste_face, self.options.blend_ratio, fake_face, 1.0 - self.options.blend_ratio, 0)
723
-
724
- # Re-assemble image
725
- paste_face = img_matte * paste_face
726
- paste_face = paste_face + (1-img_matte) * target_img.astype(np.float32)
727
- if self.options.show_face_area_overlay:
728
- # Overlay the green overlay on the final image
729
- paste_face = cv2.addWeighted(paste_face.astype(np.uint8), 1 - 0.5, green_overlay, 0.5, 0)
730
- return paste_face.astype(np.uint8)
731
-
732
-
733
- def blur_area(self, img_matte, num_erosion_iterations, blur_amount):
734
- # Detect the affine transformed white area
735
- mask_h_inds, mask_w_inds = np.where(img_matte==255)
736
- # Calculate the size (and diagonal size) of transformed white area width and height boundaries
737
- mask_h = np.max(mask_h_inds) - np.min(mask_h_inds)
738
- mask_w = np.max(mask_w_inds) - np.min(mask_w_inds)
739
- mask_size = int(np.sqrt(mask_h*mask_w))
740
- # Calculate the kernel size for eroding img_matte by kernel (insightface empirical guess for best size was max(mask_size//10,10))
741
- # k = max(mask_size//12, 8)
742
- k = max(mask_size//(blur_amount // 2) , blur_amount // 2)
743
- kernel = np.ones((k,k),np.uint8)
744
- img_matte = cv2.erode(img_matte,kernel,iterations = num_erosion_iterations)
745
- #Calculate the kernel size for blurring img_matte by blur_size (insightface empirical guess for best size was max(mask_size//20, 5))
746
- # k = max(mask_size//24, 4)
747
- k = max(mask_size//blur_amount, blur_amount//5)
748
- kernel_size = (k, k)
749
- blur_size = tuple(2*i+1 for i in kernel_size)
750
- return cv2.GaussianBlur(img_matte, blur_size, 0)
751
-
752
-
753
- def prepare_crop_frame(self, swap_frame):
754
- model_type = 'inswapper'
755
- model_mean = [0.0, 0.0, 0.0]
756
- model_standard_deviation = [1.0, 1.0, 1.0]
757
-
758
- if model_type == 'ghost':
759
- swap_frame = swap_frame[:, :, ::-1] / 127.5 - 1
760
- else:
761
- swap_frame = swap_frame[:, :, ::-1] / 255.0
762
- swap_frame = (swap_frame - model_mean) / model_standard_deviation
763
- swap_frame = swap_frame.transpose(2, 0, 1)
764
- swap_frame = np.expand_dims(swap_frame, axis = 0).astype(np.float32)
765
- return swap_frame
766
-
767
-
768
- def normalize_swap_frame(self, swap_frame):
769
- model_type = 'inswapper'
770
- swap_frame = swap_frame.transpose(1, 2, 0)
771
-
772
- if model_type == 'ghost':
773
- swap_frame = (swap_frame * 127.5 + 127.5).round()
774
- else:
775
- swap_frame = (swap_frame * 255.0).round()
776
- swap_frame = swap_frame[:, :, ::-1]
777
- return swap_frame
778
-
779
- def implode_pixel_boost(self, aligned_face_frame, model_size, pixel_boost_total : int):
780
- subsample_frame = aligned_face_frame.reshape(model_size, pixel_boost_total, model_size, pixel_boost_total, 3)
781
- subsample_frame = subsample_frame.transpose(1, 3, 0, 2, 4).reshape(pixel_boost_total ** 2, model_size, model_size, 3)
782
- return subsample_frame
783
-
784
-
785
- def explode_pixel_boost(self, subsample_frame, model_size, pixel_boost_total, pixel_boost_size):
786
- final_frame = np.stack(subsample_frame, axis = 0).reshape(pixel_boost_total, pixel_boost_total, model_size, model_size, 3)
787
- final_frame = final_frame.transpose(2, 0, 3, 1, 4).reshape(pixel_boost_size, pixel_boost_size, 3)
788
- return final_frame
789
-
790
- def process_mask(self, processor, frame:Frame, target:Frame):
791
- img_mask = processor.Run(frame, self.options.masking_text)
792
- img_mask = cv2.resize(img_mask, (target.shape[1], target.shape[0]))
793
- img_mask = np.reshape(img_mask, [img_mask.shape[0],img_mask.shape[1],1])
794
-
795
- if self.options.show_face_masking:
796
- result = (1 - img_mask) * frame.astype(np.float32)
797
- return np.uint8(result)
798
-
799
-
800
- target = target.astype(np.float32)
801
- result = (1-img_mask) * target
802
- result += img_mask * frame.astype(np.float32)
803
- return np.uint8(result)
804
-
805
-
806
- # Code for mouth restoration adapted from https://github.com/iVideoGameBoss/iRoopDeepFaceCam
807
-
808
- def create_mouth_mask(self, face: Face, frame: Frame):
809
- mouth_cutout = None
810
-
811
- landmarks = face.landmark_2d_106
812
- if landmarks is not None:
813
- # Get mouth landmarks (indices 52 to 71 typically represent the outer mouth)
814
- mouth_points = landmarks[52:71].astype(np.int32)
815
-
816
- # Add padding to mouth area
817
- min_x, min_y = np.min(mouth_points, axis=0)
818
- max_x, max_y = np.max(mouth_points, axis=0)
819
- min_x = max(0, min_x - (15*6))
820
- min_y = max(0, min_y - 22)
821
- max_x = min(frame.shape[1], max_x + (15*6))
822
- max_y = min(frame.shape[0], max_y + (90*6))
823
-
824
- # Extract the mouth area from the frame using the calculated bounding box
825
- mouth_cutout = frame[min_y:max_y, min_x:max_x].copy()
826
-
827
- return mouth_cutout, (min_x, min_y, max_x, max_y)
828
-
829
-
830
-
831
- def create_feathered_mask(self, shape, feather_amount=30):
832
- mask = np.zeros(shape[:2], dtype=np.float32)
833
- center = (shape[1] // 2, shape[0] // 2)
834
- cv2.ellipse(mask, center, (shape[1] // 2 - feather_amount, shape[0] // 2 - feather_amount),
835
- 0, 0, 360, 1, -1)
836
- mask = cv2.GaussianBlur(mask, (feather_amount*2+1, feather_amount*2+1), 0)
837
- return mask / np.max(mask)
838
-
839
- def apply_mouth_area(self, frame: np.ndarray, mouth_cutout: np.ndarray, mouth_box: tuple) -> np.ndarray:
840
- min_x, min_y, max_x, max_y = mouth_box
841
- box_width = max_x - min_x
842
- box_height = max_y - min_y
843
-
844
-
845
- # Resize the mouth cutout to match the mouth box size
846
- if mouth_cutout is None or box_width is None or box_height is None:
847
- return frame
848
- try:
849
- resized_mouth_cutout = cv2.resize(mouth_cutout, (box_width, box_height))
850
-
851
- # Extract the region of interest (ROI) from the target frame
852
- roi = frame[min_y:max_y, min_x:max_x]
853
-
854
- # Ensure the ROI and resized_mouth_cutout have the same shape
855
- if roi.shape != resized_mouth_cutout.shape:
856
- resized_mouth_cutout = cv2.resize(resized_mouth_cutout, (roi.shape[1], roi.shape[0]))
857
-
858
- # Apply color transfer from ROI to mouth cutout
859
- color_corrected_mouth = self.apply_color_transfer(resized_mouth_cutout, roi)
860
-
861
- # Create a feathered mask with increased feather amount
862
- feather_amount = min(30, box_width // 15, box_height // 15)
863
- mask = self.create_feathered_mask(resized_mouth_cutout.shape, feather_amount)
864
-
865
- # Blend the color-corrected mouth cutout with the ROI using the feathered mask
866
- mask = mask[:,:,np.newaxis] # Add channel dimension to mask
867
- blended = (color_corrected_mouth * mask + roi * (1 - mask)).astype(np.uint8)
868
-
869
- # Place the blended result back into the frame
870
- frame[min_y:max_y, min_x:max_x] = blended
871
- except Exception as e:
872
- print(f'Error {e}')
873
- pass
874
-
875
- return frame
876
-
877
- def apply_color_transfer(self, source, target):
878
- """
879
- Apply color transfer from target to source image
880
- """
881
- source = cv2.cvtColor(source, cv2.COLOR_BGR2LAB).astype("float32")
882
- target = cv2.cvtColor(target, cv2.COLOR_BGR2LAB).astype("float32")
883
-
884
- source_mean, source_std = cv2.meanStdDev(source)
885
- target_mean, target_std = cv2.meanStdDev(target)
886
-
887
- # Reshape mean and std to be broadcastable
888
- source_mean = source_mean.reshape(1, 1, 3)
889
- source_std = source_std.reshape(1, 1, 3)
890
- target_mean = target_mean.reshape(1, 1, 3)
891
- target_std = target_std.reshape(1, 1, 3)
892
-
893
- # Perform the color transfer
894
- source = (source - source_mean) * (target_std / source_std) + target_mean
895
- return cv2.cvtColor(np.clip(source, 0, 255).astype("uint8"), cv2.COLOR_LAB2BGR)
896
-
897
-
898
-
899
- def unload_models():
900
- pass
901
-
902
-
903
- def release_resources(self):
904
- for p in self.processors:
905
- p.Release()
906
- self.processors.clear()
907
- if self.videowriter is not None:
908
- self.videowriter.close()
909
- if self.streamwriter is not None:
910
- self.streamwriter.Close()
911
-
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
roop/ProcessOptions.py DELETED
@@ -1,16 +0,0 @@
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, restore_original_mouth, 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
15
- self.restore_original_mouth = restore_original_mouth
16
- self.max_num_reuse_frame = 15
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
roop/StreamWriter.py DELETED
@@ -1,60 +0,0 @@
1
- import threading
2
- import time
3
- import pyvirtualcam
4
-
5
-
6
- class StreamWriter():
7
- FPS = 30
8
- VCam = None
9
- Active = False
10
- THREAD_LOCK_STREAM = threading.Lock()
11
- time_last_process = None
12
- timespan_min = 0.0
13
-
14
- def __enter__(self):
15
- return self
16
-
17
- def __exit__(self, exc_type, exc_value, traceback):
18
- self.Close()
19
-
20
- def __init__(self, size, fps):
21
- self.time_last_process = time.perf_counter()
22
- self.FPS = fps
23
- self.timespan_min = 1.0 / fps
24
- print('Detecting virtual cam devices')
25
- self.VCam = pyvirtualcam.Camera(width=size[0], height=size[1], fps=fps, fmt=pyvirtualcam.PixelFormat.BGR, print_fps=False)
26
- if self.VCam is None:
27
- print("No virtual camera found!")
28
- return
29
- print(f'Using virtual camera: {self.VCam.device}')
30
- print(f'Using {self.VCam.native_fmt}')
31
- self.Active = True
32
-
33
-
34
- def LimitFrames(self):
35
- while True:
36
- current_time = time.perf_counter()
37
- time_passed = current_time - self.time_last_process
38
- if time_passed >= self.timespan_min:
39
- break
40
-
41
- # First version used a queue and threading. Surprisingly this
42
- # totally simple, blocking version is 10 times faster!
43
- def WriteToStream(self, frame):
44
- if self.VCam is None:
45
- return
46
- with self.THREAD_LOCK_STREAM:
47
- self.LimitFrames()
48
- self.VCam.send(frame)
49
- self.time_last_process = time.perf_counter()
50
-
51
-
52
- def Close(self):
53
- self.Active = False
54
- if self.VCam is None:
55
- self.VCam.close()
56
- self.VCam = None
57
-
58
-
59
-
60
-
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
roop/__init__.py DELETED
File without changes
roop/capturer.py DELETED
@@ -1,46 +0,0 @@
1
- from typing import Optional
2
- import cv2
3
- import numpy as np
4
-
5
- from roop.typing import Frame
6
-
7
- current_video_path = None
8
- current_frame_total = 0
9
- current_capture = None
10
-
11
- def get_image_frame(filename: str):
12
- try:
13
- return cv2.imdecode(np.fromfile(filename, dtype=np.uint8), cv2.IMREAD_COLOR)
14
- except:
15
- print(f"Exception reading {filename}")
16
- return None
17
-
18
-
19
- def get_video_frame(video_path: str, frame_number: int = 0) -> Optional[Frame]:
20
- global current_video_path, current_capture, current_frame_total
21
-
22
- if video_path != current_video_path:
23
- release_video()
24
- current_capture = cv2.VideoCapture(video_path)
25
- current_video_path = video_path
26
- current_frame_total = current_capture.get(cv2.CAP_PROP_FRAME_COUNT)
27
-
28
- current_capture.set(cv2.CAP_PROP_POS_FRAMES, min(current_frame_total, frame_number - 1))
29
- has_frame, frame = current_capture.read()
30
- if has_frame:
31
- return frame
32
- return None
33
-
34
- def release_video():
35
- global current_capture
36
-
37
- if current_capture is not None:
38
- current_capture.release()
39
- current_capture = None
40
-
41
-
42
- def get_video_frame_total(video_path: str) -> int:
43
- capture = cv2.VideoCapture(video_path)
44
- video_frame_total = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
45
- capture.release()
46
- return video_frame_total
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
roop/core.py DELETED
@@ -1,404 +0,0 @@
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
- import argparse
18
-
19
- from time import time
20
-
21
- import roop.globals
22
- import roop.metadata
23
- import roop.utilities as util
24
- import roop.util_ffmpeg as ffmpeg
25
- import ui.main as main
26
- from settings import Settings
27
- from roop.face_util import extract_face_images
28
- from roop.ProcessEntry import ProcessEntry
29
- from roop.ProcessMgr import ProcessMgr
30
- from roop.ProcessOptions import ProcessOptions
31
- from roop.capturer import get_video_frame_total, release_video
32
-
33
-
34
- clip_text = None
35
-
36
- call_display_ui = None
37
-
38
- process_mgr = None
39
-
40
-
41
- if 'ROCMExecutionProvider' in roop.globals.execution_providers:
42
- del torch
43
-
44
- warnings.filterwarnings('ignore', category=FutureWarning, module='insightface')
45
- warnings.filterwarnings('ignore', category=UserWarning, module='torchvision')
46
-
47
-
48
- def parse_args() -> None:
49
- signal.signal(signal.SIGINT, lambda signal_number, frame: destroy())
50
- roop.globals.headless = False
51
-
52
- program = argparse.ArgumentParser(formatter_class=lambda prog: argparse.HelpFormatter(prog, max_help_position=100))
53
- program.add_argument('--server_share', help='Public server', dest='server_share', action='store_true', default=False)
54
- program.add_argument('--cuda_device_id', help='Index of the cuda gpu to use', dest='cuda_device_id', type=int, default=0)
55
- roop.globals.startup_args = program.parse_args()
56
- # Always enable all processors when using GUI
57
- roop.globals.frame_processors = ['face_swapper', 'face_enhancer']
58
-
59
-
60
- def encode_execution_providers(execution_providers: List[str]) -> List[str]:
61
- return [execution_provider.replace('ExecutionProvider', '').lower() for execution_provider in execution_providers]
62
-
63
-
64
- def decode_execution_providers(execution_providers: List[str]) -> List[str]:
65
- list_providers = [provider for provider, encoded_execution_provider in zip(onnxruntime.get_available_providers(), encode_execution_providers(onnxruntime.get_available_providers()))
66
- if any(execution_provider in encoded_execution_provider for execution_provider in execution_providers)]
67
-
68
- try:
69
- for i in range(len(list_providers)):
70
- if list_providers[i] == 'CUDAExecutionProvider':
71
- list_providers[i] = ('CUDAExecutionProvider', {'device_id': roop.globals.cuda_device_id})
72
- torch.cuda.set_device(roop.globals.cuda_device_id)
73
- break
74
- except:
75
- pass
76
-
77
- return list_providers
78
-
79
-
80
-
81
- def suggest_max_memory() -> int:
82
- if platform.system().lower() == 'darwin':
83
- return 4
84
- return 16
85
-
86
-
87
- def suggest_execution_providers() -> List[str]:
88
- return encode_execution_providers(onnxruntime.get_available_providers())
89
-
90
-
91
- def suggest_execution_threads() -> int:
92
- if 'DmlExecutionProvider' in roop.globals.execution_providers:
93
- return 1
94
- if 'ROCMExecutionProvider' in roop.globals.execution_providers:
95
- return 1
96
- return 8
97
-
98
-
99
- def limit_resources() -> None:
100
- # limit memory usage
101
- if roop.globals.max_memory:
102
- memory = roop.globals.max_memory * 1024 ** 3
103
- if platform.system().lower() == 'darwin':
104
- memory = roop.globals.max_memory * 1024 ** 6
105
- if platform.system().lower() == 'windows':
106
- import ctypes
107
- kernel32 = ctypes.windll.kernel32 # type: ignore[attr-defined]
108
- kernel32.SetProcessWorkingSetSize(-1, ctypes.c_size_t(memory), ctypes.c_size_t(memory))
109
- else:
110
- import resource
111
- resource.setrlimit(resource.RLIMIT_DATA, (memory, memory))
112
-
113
-
114
-
115
- def release_resources() -> None:
116
- import gc
117
- global process_mgr
118
-
119
- if process_mgr is not None:
120
- process_mgr.release_resources()
121
- process_mgr = None
122
-
123
- gc.collect()
124
- # if 'CUDAExecutionProvider' in roop.globals.execution_providers and torch.cuda.is_available():
125
- # with torch.cuda.device('cuda'):
126
- # torch.cuda.empty_cache()
127
- # torch.cuda.ipc_collect()
128
-
129
-
130
- def pre_check() -> bool:
131
- if sys.version_info < (3, 9):
132
- update_status('Python version is not supported - please upgrade to 3.9 or higher.')
133
- return False
134
-
135
- download_directory_path = util.resolve_relative_path('../models')
136
- util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/inswapper_128.onnx'])
137
- util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/GFPGANv1.4.onnx'])
138
- util.conditional_download(download_directory_path, ['https://github.com/csxmli2016/DMDNet/releases/download/v1/DMDNet.pth'])
139
- util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/GPEN-BFR-512.onnx'])
140
- util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/restoreformer_plus_plus.onnx'])
141
- util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/xseg.onnx'])
142
- download_directory_path = util.resolve_relative_path('../models/CLIP')
143
- util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/rd64-uni-refined.pth'])
144
- download_directory_path = util.resolve_relative_path('../models/CodeFormer')
145
- util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/CodeFormerv0.1.onnx'])
146
- download_directory_path = util.resolve_relative_path('../models/Frame')
147
- util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/deoldify_artistic.onnx'])
148
- util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/deoldify_stable.onnx'])
149
- util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/isnet-general-use.onnx'])
150
- util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/real_esrgan_x4.onnx'])
151
- util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/real_esrgan_x2.onnx'])
152
- util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/lsdir_x4.onnx'])
153
-
154
- if not shutil.which('ffmpeg'):
155
- update_status('ffmpeg is not installed.')
156
- return True
157
-
158
- def set_display_ui(function):
159
- global call_display_ui
160
-
161
- call_display_ui = function
162
-
163
-
164
- def update_status(message: str) -> None:
165
- global call_display_ui
166
-
167
- print(message)
168
- if call_display_ui is not None:
169
- call_display_ui(message)
170
-
171
-
172
-
173
-
174
- def start() -> None:
175
- if roop.globals.headless:
176
- print('Headless mode currently unsupported - starting UI!')
177
- # faces = extract_face_images(roop.globals.source_path, (False, 0))
178
- # roop.globals.INPUT_FACES.append(faces[roop.globals.source_face_index])
179
- # faces = extract_face_images(roop.globals.target_path, (False, util.has_image_extension(roop.globals.target_path)))
180
- # roop.globals.TARGET_FACES.append(faces[roop.globals.target_face_index])
181
- # if 'face_enhancer' in roop.globals.frame_processors:
182
- # roop.globals.selected_enhancer = 'GFPGAN'
183
-
184
- batch_process_regular(None, False, None)
185
-
186
-
187
- def get_processing_plugins(masking_engine):
188
- processors = { "faceswap": {}}
189
- if masking_engine is not None:
190
- processors.update({masking_engine: {}})
191
-
192
- if roop.globals.selected_enhancer == 'GFPGAN':
193
- processors.update({"gfpgan": {}})
194
- elif roop.globals.selected_enhancer == 'Codeformer':
195
- processors.update({"codeformer": {}})
196
- elif roop.globals.selected_enhancer == 'DMDNet':
197
- processors.update({"dmdnet": {}})
198
- elif roop.globals.selected_enhancer == 'GPEN':
199
- processors.update({"gpen": {}})
200
- elif roop.globals.selected_enhancer == 'Restoreformer++':
201
- processors.update({"restoreformer++": {}})
202
- return processors
203
-
204
-
205
- def live_swap(frame, options):
206
- global process_mgr
207
-
208
- if frame is None:
209
- return frame
210
-
211
- if process_mgr is None:
212
- process_mgr = ProcessMgr(None)
213
-
214
- # if len(roop.globals.INPUT_FACESETS) <= selected_index:
215
- # selected_index = 0
216
- process_mgr.initialize(roop.globals.INPUT_FACESETS, roop.globals.TARGET_FACES, options)
217
- newframe = process_mgr.process_frame(frame)
218
- if newframe is None:
219
- return frame
220
- return newframe
221
-
222
-
223
- def batch_process_regular(output_method, files:list[ProcessEntry], masking_engine:str, new_clip_text:str, use_new_method, imagemask, restore_original_mouth, num_swap_steps, progress, selected_index = 0) -> None:
224
- global clip_text, process_mgr
225
-
226
- release_resources()
227
- limit_resources()
228
- if process_mgr is None:
229
- process_mgr = ProcessMgr(progress)
230
- mask = imagemask["layers"][0] if imagemask is not None else None
231
- if len(roop.globals.INPUT_FACESETS) <= selected_index:
232
- selected_index = 0
233
- options = ProcessOptions(get_processing_plugins(masking_engine), roop.globals.distance_threshold, roop.globals.blend_ratio,
234
- roop.globals.face_swap_mode, selected_index, new_clip_text, mask, num_swap_steps,
235
- roop.globals.subsample_size, False, restore_original_mouth)
236
- process_mgr.initialize(roop.globals.INPUT_FACESETS, roop.globals.TARGET_FACES, options)
237
- batch_process(output_method, files, use_new_method)
238
- return
239
-
240
- def batch_process_with_options(files:list[ProcessEntry], options, progress):
241
- global clip_text, process_mgr
242
-
243
- release_resources()
244
- limit_resources()
245
- if process_mgr is None:
246
- process_mgr = ProcessMgr(progress)
247
- process_mgr.initialize(roop.globals.INPUT_FACESETS, roop.globals.TARGET_FACES, options)
248
- roop.globals.keep_frames = False
249
- roop.globals.wait_after_extraction = False
250
- roop.globals.skip_audio = False
251
- batch_process("Files", files, True)
252
-
253
-
254
-
255
- def batch_process(output_method, files:list[ProcessEntry], use_new_method) -> None:
256
- global clip_text, process_mgr
257
-
258
- roop.globals.processing = True
259
-
260
- # limit threads for some providers
261
- max_threads = suggest_execution_threads()
262
- if max_threads == 1:
263
- roop.globals.execution_threads = 1
264
-
265
- imagefiles:list[ProcessEntry] = []
266
- videofiles:list[ProcessEntry] = []
267
-
268
- update_status('Sorting videos/images')
269
-
270
-
271
- for index, f in enumerate(files):
272
- fullname = f.filename
273
- if util.has_image_extension(fullname):
274
- destination = util.get_destfilename_from_path(fullname, roop.globals.output_path, f'.{roop.globals.CFG.output_image_format}')
275
- destination = util.replace_template(destination, index=index)
276
- pathlib.Path(os.path.dirname(destination)).mkdir(parents=True, exist_ok=True)
277
- f.finalname = destination
278
- imagefiles.append(f)
279
-
280
- elif util.is_video(fullname) or util.has_extension(fullname, ['gif']):
281
- destination = util.get_destfilename_from_path(fullname, roop.globals.output_path, f'__temp.{roop.globals.CFG.output_video_format}')
282
- f.finalname = destination
283
- videofiles.append(f)
284
-
285
-
286
-
287
- if(len(imagefiles) > 0):
288
- update_status('Processing image(s)')
289
- origimages = []
290
- fakeimages = []
291
- for f in imagefiles:
292
- origimages.append(f.filename)
293
- fakeimages.append(f.finalname)
294
-
295
- process_mgr.run_batch(origimages, fakeimages, roop.globals.execution_threads)
296
- origimages.clear()
297
- fakeimages.clear()
298
-
299
- if(len(videofiles) > 0):
300
- for index,v in enumerate(videofiles):
301
- if not roop.globals.processing:
302
- end_processing('Processing stopped!')
303
- return
304
- fps = v.fps if v.fps > 0 else util.detect_fps(v.filename)
305
- if v.endframe == 0:
306
- v.endframe = get_video_frame_total(v.filename)
307
-
308
- is_streaming_only = output_method == "Virtual Camera"
309
- if is_streaming_only == False:
310
- update_status(f'Creating {os.path.basename(v.finalname)} with {fps} FPS...')
311
-
312
- start_processing = time()
313
- if is_streaming_only == False and roop.globals.keep_frames or not use_new_method:
314
- util.create_temp(v.filename)
315
- update_status('Extracting frames...')
316
- ffmpeg.extract_frames(v.filename,v.startframe,v.endframe, fps)
317
- if not roop.globals.processing:
318
- end_processing('Processing stopped!')
319
- return
320
-
321
- temp_frame_paths = util.get_temp_frame_paths(v.filename)
322
- process_mgr.run_batch(temp_frame_paths, temp_frame_paths, roop.globals.execution_threads)
323
- if not roop.globals.processing:
324
- end_processing('Processing stopped!')
325
- return
326
- if roop.globals.wait_after_extraction:
327
- extract_path = os.path.dirname(temp_frame_paths[0])
328
- util.open_folder(extract_path)
329
- input("Press any key to continue...")
330
- print("Resorting frames to create video")
331
- util.sort_rename_frames(extract_path)
332
-
333
- ffmpeg.create_video(v.filename, v.finalname, fps)
334
- if not roop.globals.keep_frames:
335
- util.delete_temp_frames(temp_frame_paths[0])
336
- else:
337
- if util.has_extension(v.filename, ['gif']):
338
- skip_audio = True
339
- else:
340
- skip_audio = roop.globals.skip_audio
341
- process_mgr.run_batch_inmem(output_method, v.filename, v.finalname, v.startframe, v.endframe, fps,roop.globals.execution_threads)
342
-
343
- if not roop.globals.processing:
344
- end_processing('Processing stopped!')
345
- return
346
-
347
- video_file_name = v.finalname
348
- if os.path.isfile(video_file_name):
349
- destination = ''
350
- if util.has_extension(v.filename, ['gif']):
351
- gifname = util.get_destfilename_from_path(v.filename, roop.globals.output_path, '.gif')
352
- destination = util.replace_template(gifname, index=index)
353
- pathlib.Path(os.path.dirname(destination)).mkdir(parents=True, exist_ok=True)
354
-
355
- update_status('Creating final GIF')
356
- ffmpeg.create_gif_from_video(video_file_name, destination)
357
- if os.path.isfile(destination):
358
- os.remove(video_file_name)
359
- else:
360
- skip_audio = roop.globals.skip_audio
361
- destination = util.replace_template(video_file_name, index=index)
362
- pathlib.Path(os.path.dirname(destination)).mkdir(parents=True, exist_ok=True)
363
-
364
- if not skip_audio:
365
- ffmpeg.restore_audio(video_file_name, v.filename, v.startframe, v.endframe, destination)
366
- if os.path.isfile(destination):
367
- os.remove(video_file_name)
368
- else:
369
- shutil.move(video_file_name, destination)
370
-
371
- elif is_streaming_only == False:
372
- update_status(f'Failed processing {os.path.basename(v.finalname)}!')
373
- elapsed_time = time() - start_processing
374
- average_fps = (v.endframe - v.startframe) / elapsed_time
375
- update_status(f'\nProcessing {os.path.basename(destination)} took {elapsed_time:.2f} secs, {average_fps:.2f} frames/s')
376
- end_processing('Finished')
377
-
378
-
379
- def end_processing(msg:str):
380
- update_status(msg)
381
- roop.globals.target_folder_path = None
382
- release_resources()
383
-
384
-
385
- def destroy() -> None:
386
- if roop.globals.target_path:
387
- util.clean_temp(roop.globals.target_path)
388
- release_resources()
389
- sys.exit()
390
-
391
-
392
- def run() -> None:
393
- parse_args()
394
- if not pre_check():
395
- return
396
- roop.globals.CFG = Settings('config.yaml')
397
- roop.globals.cuda_device_id = roop.globals.startup_args.cuda_device_id
398
- roop.globals.execution_threads = roop.globals.CFG.max_threads
399
- roop.globals.video_encoder = roop.globals.CFG.output_video_codec
400
- roop.globals.video_quality = roop.globals.CFG.video_quality
401
- roop.globals.max_memory = roop.globals.CFG.memory_limit if roop.globals.CFG.memory_limit > 0 else None
402
- if roop.globals.startup_args.server_share:
403
- roop.globals.CFG.server_share = True
404
- main.run()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
roop/face_util.py DELETED
@@ -1,309 +0,0 @@
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_thread_semaphore
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 conditional_thread_semaphore():
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 DELETED
@@ -1,218 +0,0 @@
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 DELETED
@@ -1,56 +0,0 @@
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
- startup_args = None
9
-
10
- cuda_device_id = 0
11
- frame_processors: List[str] = []
12
- keep_fps = None
13
- keep_frames = None
14
- autorotate_faces = None
15
- vr_mode = None
16
- skip_audio = None
17
- wait_after_extraction = None
18
- many_faces = None
19
- use_batch = None
20
- source_face_index = 0
21
- target_face_index = 0
22
- face_position = None
23
- video_encoder = None
24
- video_quality = None
25
- max_memory = None
26
- execution_providers: List[str] = []
27
- execution_threads = None
28
- headless = None
29
- log_level = 'error'
30
- selected_enhancer = None
31
- subsample_size = 128
32
- face_swap_mode = None
33
- blend_ratio = 0.5
34
- distance_threshold = 0.65
35
- default_det_size = True
36
-
37
- no_face_action = 0
38
-
39
- processing = False
40
-
41
- g_current_face_analysis = None
42
- g_desired_face_analysis = None
43
-
44
- FACE_ENHANCER = None
45
-
46
- INPUT_FACESETS = []
47
- TARGET_FACES = []
48
-
49
-
50
- IMAGE_CHAIN_PROCESSOR = None
51
- VIDEO_CHAIN_PROCESSOR = None
52
- BATCH_IMAGE_CHAIN_PROCESSOR = None
53
-
54
- CFG: Settings = None
55
-
56
-
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
roop/metadata.py DELETED
@@ -1,2 +0,0 @@
1
- name = 'roop unleashed'
2
- version = '4.3.3'
 
 
 
roop/processors/Enhance_CodeFormer.py DELETED
@@ -1,71 +0,0 @@
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_CodeFormer():
11
- model_codeformer = None
12
-
13
- plugin_options:dict = None
14
-
15
- processorname = 'codeformer'
16
- type = 'enhance'
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_codeformer is None:
26
- # replace Mac mps with cpu for the moment
27
- self.devicename = self.plugin_options["devicename"].replace('mps', 'cpu')
28
- model_path = resolve_relative_path('../models/CodeFormer/CodeFormerv0.1.onnx')
29
- self.model_codeformer = onnxruntime.InferenceSession(model_path, None, providers=roop.globals.execution_providers)
30
- self.model_inputs = self.model_codeformer.get_inputs()
31
- model_outputs = self.model_codeformer.get_outputs()
32
- self.io_binding = self.model_codeformer.io_binding()
33
- self.io_binding.bind_cpu_input(self.model_inputs[1].name, np.array([0.5]))
34
- self.io_binding.bind_output(model_outputs[0].name, self.devicename)
35
-
36
-
37
- def Run(self, source_faceset: FaceSet, target_face: Face, temp_frame: Frame) -> Frame:
38
- input_size = temp_frame.shape[1]
39
- # preprocess
40
- temp_frame = cv2.resize(temp_frame, (512, 512), cv2.INTER_CUBIC)
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
- self.io_binding.bind_cpu_input(self.model_inputs[0].name, temp_frame.astype(np.float32))
47
- self.model_codeformer.run_with_iobinding(self.io_binding)
48
- ort_outs = self.io_binding.copy_outputs_to_cpu()
49
- result = ort_outs[0][0]
50
- del ort_outs
51
-
52
- # post-process
53
- result = result.transpose((1, 2, 0))
54
-
55
- un_min = -1.0
56
- un_max = 1.0
57
- result = np.clip(result, un_min, un_max)
58
- result = (result - un_min) / (un_max - un_min)
59
-
60
- result = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
61
- result = (result * 255.0).round()
62
- scale_factor = int(result.shape[1] / input_size)
63
- return result.astype(np.uint8), scale_factor
64
-
65
-
66
- def Release(self):
67
- del self.model_codeformer
68
- self.model_codeformer = None
69
- del self.io_binding
70
- self.io_binding = None
71
-
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
roop/processors/Enhance_DMDNet.py DELETED
@@ -1,898 +0,0 @@
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 DELETED
@@ -1,73 +0,0 @@
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_GFPGAN():
11
- plugin_options:dict = None
12
-
13
- model_gfpgan = None
14
- name = None
15
- devicename = None
16
-
17
- processorname = 'gfpgan'
18
- type = 'enhance'
19
-
20
-
21
- def Initialize(self, plugin_options:dict):
22
- if self.plugin_options is not None:
23
- if self.plugin_options["devicename"] != plugin_options["devicename"]:
24
- self.Release()
25
-
26
- self.plugin_options = plugin_options
27
- if self.model_gfpgan is None:
28
- model_path = resolve_relative_path('../models/GFPGANv1.4.onnx')
29
- self.model_gfpgan = onnxruntime.InferenceSession(model_path, None, providers=roop.globals.execution_providers)
30
- # replace Mac mps with cpu for the moment
31
- self.devicename = self.plugin_options["devicename"].replace('mps', 'cpu')
32
-
33
- self.name = self.model_gfpgan.get_inputs()[0].name
34
-
35
- def Run(self, source_faceset: FaceSet, target_face: Face, temp_frame: Frame) -> Frame:
36
- # preprocess
37
- input_size = temp_frame.shape[1]
38
- temp_frame = cv2.resize(temp_frame, (512, 512), cv2.INTER_CUBIC)
39
-
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
- io_binding = self.model_gfpgan.io_binding()
46
- io_binding.bind_cpu_input("input", temp_frame)
47
- io_binding.bind_output("1288", self.devicename)
48
- self.model_gfpgan.run_with_iobinding(io_binding)
49
- ort_outs = io_binding.copy_outputs_to_cpu()
50
- result = ort_outs[0][0]
51
-
52
- # post-process
53
- result = np.clip(result, -1, 1)
54
- result = (result + 1) / 2
55
- result = result.transpose(1, 2, 0) * 255.0
56
- result = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
57
- scale_factor = int(result.shape[1] / input_size)
58
- return result.astype(np.uint8), scale_factor
59
-
60
-
61
- def Release(self):
62
- self.model_gfpgan = None
63
-
64
-
65
-
66
-
67
-
68
-
69
-
70
-
71
-
72
-
73
-
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
roop/processors/Enhance_GPEN.py DELETED
@@ -1,63 +0,0 @@
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 DELETED
@@ -1,64 +0,0 @@
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 DELETED
@@ -1,61 +0,0 @@
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 DELETED
@@ -1,70 +0,0 @@
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 DELETED
@@ -1,105 +0,0 @@
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 DELETED
@@ -1,71 +0,0 @@
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 DELETED
@@ -1,129 +0,0 @@
1
- import cv2
2
- import numpy as np
3
- import onnxruntime
4
- import roop.globals
5
-
6
- from roop.utilities import resolve_relative_path, conditional_thread_semaphore
7
- from roop.typing import Frame
8
-
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
-
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.prev_type is not None and self.prev_type != self.plugin_options["subtype"]:
27
- self.Release()
28
- self.prev_type = self.plugin_options["subtype"]
29
- if self.model_upscale is None:
30
- # replace Mac mps with cpu for the moment
31
- self.devicename = self.plugin_options["devicename"].replace('mps', 'cpu')
32
- if self.prev_type == "esrganx4":
33
- model_path = resolve_relative_path('../models/Frame/real_esrgan_x4.onnx')
34
- self.scale = 4
35
- elif self.prev_type == "esrganx2":
36
- model_path = resolve_relative_path('../models/Frame/real_esrgan_x2.onnx')
37
- self.scale = 2
38
- elif self.prev_type == "lsdirx4":
39
- model_path = resolve_relative_path('../models/Frame/lsdir_x4.onnx')
40
- self.scale = 4
41
- onnxruntime.set_default_logger_severity(3)
42
- self.model_upscale = onnxruntime.InferenceSession(model_path, None, providers=roop.globals.execution_providers)
43
- self.model_inputs = self.model_upscale.get_inputs()
44
- model_outputs = self.model_upscale.get_outputs()
45
- self.io_binding = self.model_upscale.io_binding()
46
- self.io_binding.bind_output(model_outputs[0].name, self.devicename)
47
-
48
- def getProcessedResolution(self, width, height):
49
- return (width * self.scale, height * self.scale)
50
-
51
- # borrowed from facefusion -> https://github.com/facefusion/facefusion
52
- def prepare_tile_frame(self, tile_frame : Frame) -> Frame:
53
- tile_frame = np.expand_dims(tile_frame[:, :, ::-1], axis = 0)
54
- tile_frame = tile_frame.transpose(0, 3, 1, 2)
55
- tile_frame = tile_frame.astype(np.float32) / 255
56
- return tile_frame
57
-
58
-
59
- def normalize_tile_frame(self, tile_frame : Frame) -> Frame:
60
- tile_frame = tile_frame.transpose(0, 2, 3, 1).squeeze(0) * 255
61
- tile_frame = tile_frame.clip(0, 255).astype(np.uint8)[:, :, ::-1]
62
- return tile_frame
63
-
64
- def create_tile_frames(self, input_frame : Frame, size):
65
- input_frame = np.pad(input_frame, ((size[1], size[1]), (size[1], size[1]), (0, 0)))
66
- tile_width = size[0] - 2 * size[2]
67
- pad_size_bottom = size[2] + tile_width - input_frame.shape[0] % tile_width
68
- pad_size_right = size[2] + tile_width - input_frame.shape[1] % tile_width
69
- pad_vision_frame = np.pad(input_frame, ((size[2], pad_size_bottom), (size[2], pad_size_right), (0, 0)))
70
- pad_height, pad_width = pad_vision_frame.shape[:2]
71
- row_range = range(size[2], pad_height - size[2], tile_width)
72
- col_range = range(size[2], pad_width - size[2], tile_width)
73
- tile_frames = []
74
-
75
- for row_frame in row_range:
76
- top = row_frame - size[2]
77
- bottom = row_frame + size[2] + tile_width
78
- for column_vision_frame in col_range:
79
- left = column_vision_frame - size[2]
80
- right = column_vision_frame + size[2] + tile_width
81
- tile_frames.append(pad_vision_frame[top:bottom, left:right, :])
82
- return tile_frames, pad_width, pad_height
83
-
84
-
85
- def merge_tile_frames(self, tile_frames, temp_width : int, temp_height : int, pad_width : int, pad_height : int, size) -> Frame:
86
- merge_frame = np.zeros((pad_height, pad_width, 3)).astype(np.uint8)
87
- tile_width = tile_frames[0].shape[1] - 2 * size[2]
88
- tiles_per_row = min(pad_width // tile_width, len(tile_frames))
89
-
90
- for index, tile_frame in enumerate(tile_frames):
91
- tile_frame = tile_frame[size[2]:-size[2], size[2]:-size[2]]
92
- row_index = index // tiles_per_row
93
- col_index = index % tiles_per_row
94
- top = row_index * tile_frame.shape[0]
95
- bottom = top + tile_frame.shape[0]
96
- left = col_index * tile_frame.shape[1]
97
- right = left + tile_frame.shape[1]
98
- merge_frame[top:bottom, left:right, :] = tile_frame
99
- merge_frame = merge_frame[size[1] : size[1] + temp_height, size[1]: size[1] + temp_width, :]
100
- return merge_frame
101
-
102
-
103
- def Run(self, temp_frame: Frame) -> Frame:
104
- size = (128, 8, 2)
105
- temp_height, temp_width = temp_frame.shape[:2]
106
- upscale_tile_frames, pad_width, pad_height = self.create_tile_frames(temp_frame, size)
107
-
108
- for index, tile_frame in enumerate(upscale_tile_frames):
109
- tile_frame = self.prepare_tile_frame(tile_frame)
110
- with conditional_thread_semaphore():
111
- self.io_binding.bind_cpu_input(self.model_inputs[0].name, tile_frame)
112
- self.model_upscale.run_with_iobinding(self.io_binding)
113
- ort_outs = self.io_binding.copy_outputs_to_cpu()
114
- result = ort_outs[0]
115
- upscale_tile_frames[index] = self.normalize_tile_frame(result)
116
- final_frame = self.merge_tile_frames(upscale_tile_frames, temp_width * self.scale
117
- , temp_height * self.scale
118
- , pad_width * self.scale, pad_height * self.scale
119
- , (size[0] * self.scale, size[1] * self.scale, size[2] * self.scale))
120
- return final_frame.astype(np.uint8)
121
-
122
-
123
-
124
- def Release(self):
125
- del self.model_upscale
126
- self.model_upscale = None
127
- del self.io_binding
128
- self.io_binding = None
129
-
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
roop/processors/Mask_Clip2Seg.py DELETED
@@ -1,94 +0,0 @@
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 DELETED
@@ -1,58 +0,0 @@
1
- import numpy as np
2
- import cv2
3
- import onnxruntime
4
- import roop.globals
5
-
6
- from roop.typing import Frame
7
- from roop.utilities import resolve_relative_path, conditional_thread_semaphore
8
-
9
-
10
-
11
- class Mask_XSeg():
12
- plugin_options:dict = None
13
-
14
- model_xseg = None
15
-
16
- processorname = 'mask_xseg'
17
- type = 'mask'
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_xseg is None:
27
- model_path = resolve_relative_path('../models/xseg.onnx')
28
- onnxruntime.set_default_logger_severity(3)
29
- self.model_xseg = onnxruntime.InferenceSession(model_path, None, providers=roop.globals.execution_providers)
30
- self.model_inputs = self.model_xseg.get_inputs()
31
- self.model_outputs = self.model_xseg.get_outputs()
32
-
33
- # replace Mac mps with cpu for the moment
34
- self.devicename = self.plugin_options["devicename"].replace('mps', 'cpu')
35
-
36
-
37
- def Run(self, img1, keywords:str) -> Frame:
38
- temp_frame = cv2.resize(img1, (256, 256), cv2.INTER_CUBIC)
39
- temp_frame = temp_frame.astype('float32') / 255.0
40
- temp_frame = temp_frame[None, ...]
41
- io_binding = self.model_xseg.io_binding()
42
- io_binding.bind_cpu_input(self.model_inputs[0].name, temp_frame)
43
- io_binding.bind_output(self.model_outputs[0].name, self.devicename)
44
- self.model_xseg.run_with_iobinding(io_binding)
45
- ort_outs = io_binding.copy_outputs_to_cpu()
46
- result = ort_outs[0][0]
47
- result = np.clip(result, 0, 1.0)
48
- result[result < 0.1] = 0
49
- # invert values to mask areas to keep
50
- result = 1.0 - result
51
- return result
52
-
53
-
54
- def Release(self):
55
- del self.model_xseg
56
- self.model_xseg = None
57
-
58
-
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
roop/template_parser.py DELETED
@@ -1,23 +0,0 @@
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 DELETED
@@ -1,9 +0,0 @@
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 DELETED
@@ -1,130 +0,0 @@
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, release_video
77
-
78
- fps = util.detect_fps(video_path)
79
- frame = get_video_frame(video_path)
80
- release_video()
81
-
82
- scalex = frame.shape[0]
83
- scaley = frame.shape[1]
84
-
85
- if scalex >= scaley:
86
- scaley = -1
87
- else:
88
- scalex = -1
89
-
90
- run_ffmpeg(['-i', video_path, '-vf', f'fps={fps},scale={int(scalex)}:{int(scaley)}:flags=lanczos,split[s0][s1];[s0]palettegen[p];[s1][p]paletteuse', '-loop', '0', gif_path])
91
-
92
-
93
-
94
- def create_video_from_gif(gif_path: str, output_path):
95
- fps = util.detect_fps(gif_path)
96
- filter = """scale='trunc(in_w/2)*2':'trunc(in_h/2)*2',format=yuv420p,fps=10"""
97
- run_ffmpeg(['-i', gif_path, '-vf', f'"{filter}"', '-movflags', '+faststart', '-shortest', output_path])
98
-
99
-
100
- def repair_video(original_video: str, final_video : str):
101
- run_ffmpeg(['-i', original_video, '-movflags', 'faststart', '-acodec', 'copy', '-vcodec', 'copy', final_video])
102
-
103
-
104
- def restore_audio(intermediate_video: str, original_video: str, trim_frame_start, trim_frame_end, final_video : str) -> None:
105
- fps = util.detect_fps(original_video)
106
- commands = [ '-i', intermediate_video ]
107
- if trim_frame_start is None and trim_frame_end is None:
108
- commands.extend([ '-c:a', 'copy' ])
109
- else:
110
- # if trim_frame_start is not None:
111
- # start_time = trim_frame_start / fps
112
- # commands.extend([ '-ss', format(start_time, ".2f")])
113
- # else:
114
- # commands.extend([ '-ss', '0' ])
115
- # if trim_frame_end is not None:
116
- # end_time = trim_frame_end / fps
117
- # commands.extend([ '-to', format(end_time, ".2f")])
118
- # commands.extend([ '-c:a', 'aac' ])
119
- if trim_frame_start is not None:
120
- start_time = trim_frame_start / fps
121
- commands.extend([ '-ss', format(start_time, ".2f")])
122
- else:
123
- commands.extend([ '-ss', '0' ])
124
- if trim_frame_end is not None:
125
- end_time = trim_frame_end / fps
126
- commands.extend([ '-to', format(end_time, ".2f")])
127
- commands.extend([ '-i', original_video, "-c", "copy" ])
128
-
129
- commands.extend([ '-map', '0:v:0', '-map', '1:a:0?', '-shortest', final_video ])
130
- run_ffmpeg(commands)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
roop/utilities.py DELETED
@@ -1,378 +0,0 @@
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
- import threading
17
- import threading
18
-
19
- from typing import Union, Any
20
- from contextlib import nullcontext
21
-
22
- from pathlib import Path
23
- from typing import List, Any
24
- from tqdm import tqdm
25
- from scipy.spatial import distance
26
-
27
- import roop.template_parser as template_parser
28
-
29
- import roop.globals
30
-
31
- TEMP_FILE = "temp.mp4"
32
- TEMP_DIRECTORY = "temp"
33
-
34
- THREAD_SEMAPHORE = threading.Semaphore()
35
- NULL_CONTEXT = nullcontext()
36
-
37
-
38
- # monkey patch ssl for mac
39
- if platform.system().lower() == "darwin":
40
- ssl._create_default_https_context = ssl._create_unverified_context
41
-
42
-
43
- # https://github.com/facefusion/facefusion/blob/master/facefusion
44
- def detect_fps(target_path: str) -> float:
45
- fps = 24.0
46
- cap = cv2.VideoCapture(target_path)
47
- if cap.isOpened():
48
- fps = cap.get(cv2.CAP_PROP_FPS)
49
- cap.release()
50
- return fps
51
-
52
-
53
- # Gradio wants Images in RGB
54
- def convert_to_gradio(image):
55
- if image is None:
56
- return None
57
- return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
58
-
59
-
60
- def sort_filenames_ignore_path(filenames):
61
- """Sorts a list of filenames containing a complete path by their filename,
62
- while retaining their original path.
63
-
64
- Args:
65
- filenames: A list of filenames containing a complete path.
66
-
67
- Returns:
68
- A sorted list of filenames containing a complete path.
69
- """
70
- filename_path_tuples = [
71
- (os.path.split(filename)[1], filename) for filename in filenames
72
- ]
73
- sorted_filename_path_tuples = sorted(filename_path_tuples, key=lambda x: x[0])
74
- return [
75
- filename_path_tuple[1] for filename_path_tuple in sorted_filename_path_tuples
76
- ]
77
-
78
-
79
- def sort_rename_frames(path: str):
80
- filenames = os.listdir(path)
81
- filenames.sort()
82
- for i in range(len(filenames)):
83
- of = os.path.join(path, filenames[i])
84
- newidx = i + 1
85
- new_filename = os.path.join(
86
- path, f"{newidx:06d}." + roop.globals.CFG.output_image_format
87
- )
88
- os.rename(of, new_filename)
89
-
90
-
91
- def get_temp_frame_paths(target_path: str) -> List[str]:
92
- temp_directory_path = get_temp_directory_path(target_path)
93
- return glob.glob(
94
- (
95
- os.path.join(
96
- glob.escape(temp_directory_path),
97
- f"*.{roop.globals.CFG.output_image_format}",
98
- )
99
- )
100
- )
101
-
102
-
103
- def get_temp_directory_path(target_path: str) -> str:
104
- target_name, _ = os.path.splitext(os.path.basename(target_path))
105
- target_directory_path = os.path.dirname(target_path)
106
- return os.path.join(target_directory_path, TEMP_DIRECTORY, target_name)
107
-
108
-
109
- def get_temp_output_path(target_path: str) -> str:
110
- temp_directory_path = get_temp_directory_path(target_path)
111
- return os.path.join(temp_directory_path, TEMP_FILE)
112
-
113
-
114
- def normalize_output_path(source_path: str, target_path: str, output_path: str) -> Any:
115
- if source_path and target_path:
116
- source_name, _ = os.path.splitext(os.path.basename(source_path))
117
- target_name, target_extension = os.path.splitext(os.path.basename(target_path))
118
- if os.path.isdir(output_path):
119
- return os.path.join(
120
- output_path, source_name + "-" + target_name + target_extension
121
- )
122
- return output_path
123
-
124
-
125
- def get_destfilename_from_path(
126
- srcfilepath: str, destfilepath: str, extension: str
127
- ) -> str:
128
- fn, ext = os.path.splitext(os.path.basename(srcfilepath))
129
- if "." in extension:
130
- return os.path.join(destfilepath, f"{fn}{extension}")
131
- return os.path.join(destfilepath, f"{fn}{extension}{ext}")
132
-
133
-
134
- def replace_template(file_path: str, index: int = 0) -> str:
135
- fn, ext = os.path.splitext(os.path.basename(file_path))
136
-
137
- # Remove the "__temp" placeholder that was used as a temporary filename
138
- fn = fn.replace("__temp", "")
139
-
140
- template = roop.globals.CFG.output_template
141
- replaced_filename = template_parser.parse(
142
- template, {"index": str(index), "file": fn}
143
- )
144
-
145
- return os.path.join(roop.globals.output_path, f"{replaced_filename}{ext}")
146
-
147
-
148
- def create_temp(target_path: str) -> None:
149
- temp_directory_path = get_temp_directory_path(target_path)
150
- Path(temp_directory_path).mkdir(parents=True, exist_ok=True)
151
-
152
-
153
- def move_temp(target_path: str, output_path: str) -> None:
154
- temp_output_path = get_temp_output_path(target_path)
155
- if os.path.isfile(temp_output_path):
156
- if os.path.isfile(output_path):
157
- os.remove(output_path)
158
- shutil.move(temp_output_path, output_path)
159
-
160
-
161
- def clean_temp(target_path: str) -> None:
162
- temp_directory_path = get_temp_directory_path(target_path)
163
- parent_directory_path = os.path.dirname(temp_directory_path)
164
- if not roop.globals.keep_frames and os.path.isdir(temp_directory_path):
165
- shutil.rmtree(temp_directory_path)
166
- if os.path.exists(parent_directory_path) and not os.listdir(parent_directory_path):
167
- os.rmdir(parent_directory_path)
168
-
169
-
170
- def delete_temp_frames(filename: str) -> None:
171
- dir = os.path.dirname(os.path.dirname(filename))
172
- shutil.rmtree(dir)
173
-
174
-
175
- def has_image_extension(image_path: str) -> bool:
176
- return image_path.lower().endswith(("png", "jpg", "jpeg", "webp"))
177
-
178
-
179
- def has_extension(filepath: str, extensions: List[str]) -> bool:
180
- return filepath.lower().endswith(tuple(extensions))
181
-
182
-
183
- def is_image(image_path: str) -> bool:
184
- if image_path and os.path.isfile(image_path):
185
- if image_path.endswith(".webp"):
186
- return True
187
- mimetype, _ = mimetypes.guess_type(image_path)
188
- return bool(mimetype and mimetype.startswith("image/"))
189
- return False
190
-
191
-
192
- def is_video(video_path: str) -> bool:
193
- if video_path and os.path.isfile(video_path):
194
- mimetype, _ = mimetypes.guess_type(video_path)
195
- return bool(mimetype and mimetype.startswith("video/"))
196
- return False
197
-
198
-
199
- def conditional_download(download_directory_path: str, urls: List[str]) -> None:
200
- if not os.path.exists(download_directory_path):
201
- os.makedirs(download_directory_path)
202
- for url in urls:
203
- download_file_path = os.path.join(
204
- download_directory_path, os.path.basename(url)
205
- )
206
- if not os.path.exists(download_file_path):
207
- request = urllib.request.urlopen(url) # type: ignore[attr-defined]
208
- total = int(request.headers.get("Content-Length", 0))
209
- with tqdm(
210
- total=total,
211
- desc=f"Downloading {url}",
212
- unit="B",
213
- unit_scale=True,
214
- unit_divisor=1024,
215
- ) as progress:
216
- urllib.request.urlretrieve(url, download_file_path, reporthook=lambda count, block_size, total_size: progress.update(block_size)) # type: ignore[attr-defined]
217
-
218
-
219
- def get_local_files_from_folder(folder: str) -> List[str]:
220
- if not os.path.exists(folder) or not os.path.isdir(folder):
221
- return None
222
- files = [
223
- os.path.join(folder, f)
224
- for f in os.listdir(folder)
225
- if os.path.isfile(os.path.join(folder, f))
226
- ]
227
- return files
228
-
229
-
230
- def resolve_relative_path(path: str) -> str:
231
- return os.path.abspath(os.path.join(os.path.dirname(__file__), path))
232
-
233
-
234
- def get_device() -> str:
235
- if len(roop.globals.execution_providers) < 1:
236
- roop.globals.execution_providers = ["CPUExecutionProvider"]
237
-
238
- prov = roop.globals.execution_providers[0]
239
- if "CoreMLExecutionProvider" in prov:
240
- return "mps"
241
- if "CUDAExecutionProvider" in prov or "ROCMExecutionProvider" in prov:
242
- return "cuda"
243
- if "OpenVINOExecutionProvider" in prov:
244
- return "mkl"
245
- return "cpu"
246
-
247
-
248
- def str_to_class(module_name, class_name) -> Any:
249
- from importlib import import_module
250
-
251
- class_ = None
252
- try:
253
- module_ = import_module(module_name)
254
- try:
255
- class_ = getattr(module_, class_name)()
256
- except AttributeError:
257
- print(f"Class {class_name} does not exist")
258
- except ImportError:
259
- print(f"Module {module_name} does not exist")
260
- return class_
261
-
262
- def is_installed(name:str) -> bool:
263
- return shutil.which(name);
264
-
265
- # Taken from https://stackoverflow.com/a/68842705
266
- def get_platform() -> str:
267
- if sys.platform == "linux":
268
- try:
269
- proc_version = open("/proc/version").read()
270
- if "Microsoft" in proc_version:
271
- return "wsl"
272
- except:
273
- pass
274
- return sys.platform
275
-
276
- def open_with_default_app(filename:str):
277
- if filename == None:
278
- return
279
- platform = get_platform()
280
- if platform == "darwin":
281
- subprocess.call(("open", filename))
282
- elif platform in ["win64", "win32"]: os.startfile(filename.replace("/", "\\"))
283
- elif platform == "wsl":
284
- subprocess.call("cmd.exe /C start".split() + [filename])
285
- else: # linux variants
286
- subprocess.call("xdg-open", filename)
287
-
288
-
289
- def prepare_for_batch(target_files) -> str:
290
- print("Preparing temp files")
291
- tempfolder = os.path.join(tempfile.gettempdir(), "rooptmp")
292
- if os.path.exists(tempfolder):
293
- shutil.rmtree(tempfolder)
294
- Path(tempfolder).mkdir(parents=True, exist_ok=True)
295
- for f in target_files:
296
- newname = os.path.basename(f.name)
297
- shutil.move(f.name, os.path.join(tempfolder, newname))
298
- return tempfolder
299
-
300
-
301
- def zip(files, zipname):
302
- with zipfile.ZipFile(zipname, "w") as zip_file:
303
- for f in files:
304
- zip_file.write(f, os.path.basename(f))
305
-
306
-
307
- def unzip(zipfilename: str, target_path: str):
308
- with zipfile.ZipFile(zipfilename, "r") as zip_file:
309
- zip_file.extractall(target_path)
310
-
311
-
312
- def mkdir_with_umask(directory):
313
- oldmask = os.umask(0)
314
- # mode needs octal
315
- os.makedirs(directory, mode=0o775, exist_ok=True)
316
- os.umask(oldmask)
317
-
318
-
319
- def open_folder(path: str):
320
- platform = get_platform()
321
- try:
322
- if platform == "darwin":
323
- subprocess.call(("open", path))
324
- elif platform in ["win64", "win32"]:
325
- open_with_default_app(path)
326
- elif platform == "wsl":
327
- subprocess.call("cmd.exe /C start".split() + [path])
328
- else: # linux variants
329
- subprocess.Popen(["xdg-open", path])
330
- except Exception as e:
331
- traceback.print_exc()
332
- pass
333
- # import webbrowser
334
- # webbrowser.open(url)
335
-
336
-
337
- def create_version_html() -> str:
338
- python_version = ".".join([str(x) for x in sys.version_info[0:3]])
339
- versions_html = f"""
340
- python: <span title="{sys.version}">{python_version}</span>
341
-
342
- torch: {getattr(torch, '__long_version__',torch.__version__)}
343
-
344
- gradio: {gradio.__version__}
345
- """
346
- return versions_html
347
-
348
-
349
- def compute_cosine_distance(emb1, emb2) -> float:
350
- return distance.cosine(emb1, emb2)
351
-
352
- def has_cuda_device():
353
- return torch.cuda is not None and torch.cuda.is_available()
354
-
355
-
356
- def print_cuda_info():
357
- try:
358
- print(f'Number of CUDA devices: {torch.cuda.device_count()} Currently used Id: {torch.cuda.current_device()} Device Name: {torch.cuda.get_device_name(torch.cuda.current_device())}')
359
- except:
360
- print('No CUDA device found!')
361
-
362
- def clean_dir(path: str):
363
- contents = os.listdir(path)
364
- for item in contents:
365
- item_path = os.path.join(path, item)
366
- try:
367
- if os.path.isfile(item_path):
368
- os.remove(item_path)
369
- elif os.path.isdir(item_path):
370
- shutil.rmtree(item_path)
371
- except Exception as e:
372
- print(e)
373
-
374
-
375
- def conditional_thread_semaphore() -> Union[Any, Any]:
376
- if 'DmlExecutionProvider' in roop.globals.execution_providers or 'ROCMExecutionProvider' in roop.globals.execution_providers:
377
- return THREAD_SEMAPHORE
378
- return NULL_CONTEXT
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
roop/virtualcam.py DELETED
@@ -1,88 +0,0 @@
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, use_xseg, use_mouthrestore, 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
-
50
- options = ProcessOptions(get_processing_plugins("mask_xseg" if use_xseg else None), roop.globals.distance_threshold, roop.globals.blend_ratio,
51
- "all", 0, None, None, 1, subsample_size, False, use_mouthrestore)
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, use_xseg, use_mouthrestore, 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, use_xseg, use_mouthrestore, 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 DELETED
@@ -1,57 +0,0 @@
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