michaelj commited on
Commit
ddd4a8d
1 Parent(s): 1914f99

change enhance

Browse files
roop/processors/frame/face_enhancer.py CHANGED
@@ -17,88 +17,88 @@ NAME = 'ROOP.FACE-ENHANCER'
17
 
18
 
19
  def get_face_enhancer() -> Any:
20
- global FACE_ENHANCER
21
 
22
- with THREAD_LOCK:
23
- if FACE_ENHANCER is None:
24
- model_path = resolve_relative_path('../models/GFPGANv1.4.pth')
25
- # todo: set models path -> https://github.com/TencentARC/GFPGAN/issues/399
26
- FACE_ENHANCER = GFPGANer(model_path=model_path, upscale=1, device=get_device())
27
- return FACE_ENHANCER
28
 
29
 
30
  def get_device() -> str:
31
- if 'CUDAExecutionProvider' in roop.globals.execution_providers:
32
- return 'cuda'
33
- if 'CoreMLExecutionProvider' in roop.globals.execution_providers:
34
- return 'mps'
35
- return 'cpu'
36
 
37
 
38
  def clear_face_enhancer() -> None:
39
- global FACE_ENHANCER
40
 
41
- FACE_ENHANCER = None
42
 
43
 
44
  def pre_check() -> bool:
45
- download_directory_path = resolve_relative_path('../models')
46
- conditional_download(download_directory_path, ['https://github.com/TencentARC/GFPGAN/releases/download/v1.3.8/GFPGANv1.4.pth'])
47
- return True
48
 
49
 
50
  def pre_start() -> bool:
51
- if not is_image(roop.globals.target_path) and not is_video(roop.globals.target_path):
52
- update_status('Select an image or video for target path.', NAME)
53
- return False
54
- return True
55
 
56
 
57
  def post_process() -> None:
58
- clear_face_enhancer()
59
 
60
 
61
  def enhance_face(target_face: Face, temp_frame: Frame) -> Frame:
62
- start_x, start_y, end_x, end_y = map(int, target_face['bbox'])
63
- padding_x = int((end_x - start_x) * 0.5)
64
- padding_y = int((end_y - start_y) * 0.5)
65
- start_x = max(0, start_x - padding_x)
66
- start_y = max(0, start_y - padding_y)
67
- end_x = max(0, end_x + padding_x)
68
- end_y = max(0, end_y + padding_y)
69
- temp_face = temp_frame[start_y:end_y, start_x:end_x]
70
- if temp_face.size:
71
- with THREAD_SEMAPHORE:
72
- _, _, temp_face = get_face_enhancer().enhance(
73
- temp_face,
74
- paste_back=True
75
- )
76
- temp_frame[start_y:end_y, start_x:end_x] = temp_face
77
- return temp_frame
78
 
79
 
80
  def process_frame(source_face: Face, reference_face: Face, temp_frame: Frame) -> Frame:
81
- many_faces = get_many_faces(temp_frame)
82
- if many_faces:
83
- for target_face in many_faces:
84
- temp_frame = enhance_face(target_face, temp_frame)
85
- return temp_frame
86
 
87
 
88
  def process_frames(source_path: str, temp_frame_paths: List[str], update: Callable[[], None]) -> None:
89
- for temp_frame_path in temp_frame_paths:
90
- temp_frame = cv2.imread(temp_frame_path)
91
- result = process_frame(None, None, temp_frame)
92
- cv2.imwrite(temp_frame_path, result)
93
- if update:
94
- update()
95
 
96
 
97
  def process_image(source_path: str, target_path: str, output_path: str) -> None:
98
- target_frame = cv2.imread(target_path)
99
- result = process_frame(None, None, target_frame)
100
- cv2.imwrite(output_path, result)
101
 
102
 
103
  def process_video(source_path: str, temp_frame_paths: List[str]) -> None:
104
- roop.processors.frame.core.process_video(None, temp_frame_paths, process_frames)
 
17
 
18
 
19
  def get_face_enhancer() -> Any:
20
+ global FACE_ENHANCER
21
 
22
+ with THREAD_LOCK:
23
+ if FACE_ENHANCER is None:
24
+ model_path = resolve_relative_path('../models/GFPGANv1.4.pth')
25
+ # todo: set models path -> https://github.com/TencentARC/GFPGAN/issues/399
26
+ FACE_ENHANCER = GFPGANer(model_path=model_path, upscale=1, device=get_device())
27
+ return FACE_ENHANCER
28
 
29
 
30
  def get_device() -> str:
31
+ if 'CUDAExecutionProvider' in roop.globals.execution_providers:
32
+ return 'cuda'
33
+ if 'CoreMLExecutionProvider' in roop.globals.execution_providers:
34
+ return 'mps'
35
+ return 'cpu'
36
 
37
 
38
  def clear_face_enhancer() -> None:
39
+ global FACE_ENHANCER
40
 
41
+ FACE_ENHANCER = None
42
 
43
 
44
  def pre_check() -> bool:
45
+ download_directory_path = resolve_relative_path('../models')
46
+ conditional_download(download_directory_path, ['https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/GFPGANv1.4.pth'])
47
+ return True
48
 
49
 
50
  def pre_start() -> bool:
51
+ if not is_image(roop.globals.target_path) and not is_video(roop.globals.target_path):
52
+ update_status('Select an image or video for target path.', NAME)
53
+ return False
54
+ return True
55
 
56
 
57
  def post_process() -> None:
58
+ clear_face_enhancer()
59
 
60
 
61
  def enhance_face(target_face: Face, temp_frame: Frame) -> Frame:
62
+ start_x, start_y, end_x, end_y = map(int, target_face['bbox'])
63
+ padding_x = int((end_x - start_x) * 0.5)
64
+ padding_y = int((end_y - start_y) * 0.5)
65
+ start_x = max(0, start_x - padding_x)
66
+ start_y = max(0, start_y - padding_y)
67
+ end_x = max(0, end_x + padding_x)
68
+ end_y = max(0, end_y + padding_y)
69
+ temp_face = temp_frame[start_y:end_y, start_x:end_x]
70
+ if temp_face.size:
71
+ with THREAD_SEMAPHORE:
72
+ _, _, temp_face = get_face_enhancer().enhance(
73
+ temp_face,
74
+ paste_back=True
75
+ )
76
+ temp_frame[start_y:end_y, start_x:end_x] = temp_face
77
+ return temp_frame
78
 
79
 
80
  def process_frame(source_face: Face, reference_face: Face, temp_frame: Frame) -> Frame:
81
+ many_faces = get_many_faces(temp_frame)
82
+ if many_faces:
83
+ for target_face in many_faces:
84
+ temp_frame = enhance_face(target_face, temp_frame)
85
+ return temp_frame
86
 
87
 
88
  def process_frames(source_path: str, temp_frame_paths: List[str], update: Callable[[], None]) -> None:
89
+ for temp_frame_path in temp_frame_paths:
90
+ temp_frame = cv2.imread(temp_frame_path)
91
+ result = process_frame(None, None, temp_frame)
92
+ cv2.imwrite(temp_frame_path, result)
93
+ if update:
94
+ update()
95
 
96
 
97
  def process_image(source_path: str, target_path: str, output_path: str) -> None:
98
+ target_frame = cv2.imread(target_path)
99
+ result = process_frame(None, None, target_frame)
100
+ cv2.imwrite(output_path, result)
101
 
102
 
103
  def process_video(source_path: str, temp_frame_paths: List[str]) -> None:
104
+ roop.processors.frame.core.process_video(None, temp_frame_paths, process_frames)