|
import torch
|
from typing import List
|
from insightface.app.common import Face
|
|
from gfpgan.utils import GFPGANer
|
from .typing import Frame # 你已有定义
|
from auto_media_publisher.config.conf_base import FACE_ENHANCER_MODEL_PATH
|
|
SIMILAR_FACE_DISTANCE: float = 0.45 # < 0.4 同一人 0.4 -0.6 大概是同一人 >0.6 不是同一个
|
|
|
class FaceEnhancer:
|
def __init__(self):
|
self.device = self._select_device()
|
self._enhancer = GFPGANer(model_path=FACE_ENHANCER_MODEL_PATH, upscale=1, device=self.device)
|
|
def _select_device(self):
|
if torch.cuda.is_available():
|
return 'cuda'
|
elif torch.backends.mps.is_available():
|
return 'mps'
|
else:
|
return 'cpu'
|
|
def _enhance_face(self,target_face: Face, temp_frame: Frame) -> Frame:
|
start_x, start_y, end_x, end_y = map(int, target_face['bbox'])
|
padding_x = int((end_x - start_x) * 0.5)
|
padding_y = int((end_y - start_y) * 0.5)
|
start_x = max(0, start_x - padding_x)
|
start_y = max(0, start_y - padding_y)
|
end_x = max(0, end_x + padding_x)
|
end_y = max(0, end_y + padding_y)
|
temp_face = temp_frame[start_y:end_y, start_x:end_x]
|
if temp_face.size:
|
_, _, temp_face = self._enhancer.enhance(
|
temp_face,
|
paste_back=True
|
)
|
temp_frame[start_y:end_y, start_x:end_x] = temp_face
|
return temp_frame
|
|
|
def process_frame(self,temp_frame: Frame,many_faces:List) -> Frame:
|
#many_faces = self._get_many_faces(temp_frame) check face_analyser.py
|
if many_faces:
|
for target_face in many_faces:
|
temp_frame = self._enhance_face(target_face, temp_frame)
|
return temp_frame
|
