Spaces:
Running
Running
File size: 4,262 Bytes
0bc8459 fb94566 0bc8459 fb94566 0bc8459 fb94566 0bc8459 fb94566 0bc8459 fb94566 0bc8459 fb94566 0bc8459 fb94566 0bc8459 fb94566 0bc8459 fb94566 0bc8459 fb94566 0bc8459 fb94566 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 |
from typing import Any, List, Callable
import cv2
import threading
import gfpgan
import os
import logging
import roop.globals
import roop.processors.frame.core
from roop.core import update_status
from roop.face_analyser import get_one_face
from roop.typing import Frame, Face
from roop.utilities import conditional_download, resolve_relative_path, is_image, is_video
FACE_ENHANCER = None
THREAD_SEMAPHORE = threading.Semaphore()
THREAD_LOCK = threading.Lock()
NAME = 'ROOP.FACE-ENHANCER'
# Configure logging
logging.basicConfig(level=logging.INFO)
def get_face_enhancer() -> Any:
global FACE_ENHANCER
with THREAD_LOCK:
if FACE_ENHANCER is None:
model_path = resolve_relative_path('../models/GFPGANv1.4.pth')
try:
FACE_ENHANCER = gfpgan.GFPGANer(model_path=model_path, upscale=5) # type: ignore[attr-defined]
logging.info(f"Loaded face enhancer model from {model_path}")
except Exception as e:
logging.error(f"Failed to load face enhancer model: {e}")
FACE_ENHANCER = None
return FACE_ENHANCER
def pre_check() -> bool:
download_directory_path = resolve_relative_path('../models')
try:
conditional_download(download_directory_path, ['https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/GFPGANv1.4.pth'])
logging.info("Pre-check completed successfully.")
return True
except Exception as e:
logging.error(f"Pre-check failed: {e}")
return False
def pre_start() -> bool:
try:
if not is_image(roop.globals.target_path) and not is_video(roop.globals.target_path):
update_status('Select an image or video for target path.', NAME)
return False
logging.info("Pre-start checks passed.")
return True
except Exception as e:
logging.error(f"Pre-start check failed: {e}")
return False
def post_process() -> None:
global FACE_ENHANCER
FACE_ENHANCER = None
logging.info("Post-process cleanup done.")
def enhance_face(temp_frame: Frame) -> Frame:
try:
with THREAD_SEMAPHORE:
_, _, temp_frame = get_face_enhancer().enhance(
temp_frame,
paste_back=True
)
return temp_frame
except Exception as e:
logging.error(f"Error enhancing face: {e}")
return temp_frame # Return the unmodified frame in case of error
def process_frame(source_face: Face, temp_frame: Frame) -> Frame:
try:
target_face = get_one_face(temp_frame)
if target_face:
temp_frame = enhance_face(temp_frame)
return temp_frame
except Exception as e:
logging.error(f"Error processing frame: {e}")
return temp_frame # Return the unmodified frame in case of error
def process_frames(source_path: str, temp_frame_paths: List[str], update: Callable[[], None]) -> None:
try:
for temp_frame_path in temp_frame_paths:
temp_frame = cv2.imread(temp_frame_path)
if temp_frame is None:
raise ValueError(f"Failed to read frame from path: {temp_frame_path}")
result = process_frame(None, temp_frame)
cv2.imwrite(temp_frame_path, result)
if update:
update()
logging.info("Frames processed successfully.")
except Exception as e:
logging.error(f"Error processing frames: {e}")
def process_image(source_path: str, target_path: str, output_path: str) -> None:
try:
target_frame = cv2.imread(target_path)
if target_frame is None:
raise ValueError("Failed to read target frame.")
result = process_frame(None, target_frame)
cv2.imwrite(output_path, result)
logging.info(f"Image processed and saved to {output_path}.")
except Exception as e:
logging.error(f"Error processing image: {e}")
def process_video(source_path: str, temp_frame_paths: List[str]) -> None:
try:
roop.processors.frame.core.process_video(None, temp_frame_paths, process_frames)
logging.info("Video processing completed.")
except Exception as e:
logging.error(f"Error processing video: {e}")
|