关于提取脸部速度问题的解决

day270010678 · 发表于 5 天前

星级打分

1
2
3
4
5

平均分:NAN 参与人数:0 我的评分:未评

我发个我自己修改的用的版本吧。

因为我的是python3.10版本以及相关库的环境，所以如果出错，你们自己解决。。

首先是FANExtractor.py

import os
import traceback
from pathlib import Path
import cv2
import numpy as np
from numpy import linalg as npla
from facelib import FaceType, LandmarksProcessor
from core.leras import nn
"""
ported from https://github.com/1adrianb/face-alignment
"""
class FANExtractor(object):
def __init__(self, landmarks_3D=False, place_model_on_cpu=False):
model_path = Path(__file__).parent / ("2DFAN.npy" if not landmarks_3D else "3DFAN.npy")
if not model_path.exists():
raise Exception("Unable to load FANExtractor model")
nn.initialize(data_format="NHWC")
tf = nn.tf
class ConvBlock(nn.ModelBase):
def on_build(self, in_planes, out_planes):
self.in_planes = in_planes
self.out_planes = out_planes
self.bn1 = nn.BatchNorm2D(in_planes)
self.conv1 = nn.Conv2D(in_planes, out_planes // 2, kernel_size=3, strides=1, padding='SAME', use_bias=False)
self.bn2 = nn.BatchNorm2D(out_planes // 2)
self.conv2 = nn.Conv2D(out_planes // 2, out_planes // 4, kernel_size=3, strides=1, padding='SAME', use_bias=False)
self.bn3 = nn.BatchNorm2D(out_planes // 4)
self.conv3 = nn.Conv2D(out_planes // 4, out_planes // 4, kernel_size=3, strides=1, padding='SAME', use_bias=False)
if self.in_planes != self.out_planes:
self.down_bn1 = nn.BatchNorm2D(in_planes)
self.down_conv1 = nn.Conv2D(in_planes, out_planes, kernel_size=1, strides=1, padding='VALID', use_bias=False)
else:
self.down_bn1 = None
self.down_conv1 = None
def forward(self, input):
x = input
x = self.bn1(x)
x = tf.nn.relu(x)
x = out1 = self.conv1(x)
x = self.bn2(x)
x = tf.nn.relu(x)
x = out2 = self.conv2(x)
x = self.bn3(x)
x = tf.nn.relu(x)
x = out3 = self.conv3(x)
x = tf.concat([out1, out2, out3], axis=-1)
if self.in_planes != self.out_planes:
downsample = self.down_bn1(input)
downsample = tf.nn.relu(downsample)
downsample = self.down_conv1(downsample)
x = x + downsample
else:
x = x + input
return x
class HourGlass(nn.ModelBase):
def on_build(self, in_planes, depth):
self.b1 = ConvBlock(in_planes, 256)
self.b2 = ConvBlock(in_planes, 256)
if depth > 1:
self.b2_plus = HourGlass(256, depth - 1)
else:
self.b2_plus = ConvBlock(256, 256)
self.b3 = ConvBlock(256, 256)
def forward(self, input):
up1 = self.b1(input)
low1 = tf.nn.avg_pool(input, [1, 2, 2, 1], [1, 2, 2, 1], 'VALID')
low1 = self.b2(low1)
low2 = self.b2_plus(low1)
low3 = self.b3(low2)
up2 = nn.upsample2d(low3)
return up1 + up2
class FAN(nn.ModelBase):
def __init__(self):
super().__init__(name='FAN')
def on_build(self):
self.conv1 = nn.Conv2D(3, 64, kernel_size=7, strides=2, padding='SAME')
self.bn1 = nn.BatchNorm2D(64)
self.conv2 = ConvBlock(64, 128)
self.conv3 = ConvBlock(128, 128)
self.conv4 = ConvBlock(128, 256)
self.m = []
self.top_m = []
self.conv_last = []
self.bn_end = []
self.l = []
self.bl = []
self.al = []
for i in range(4):
self.m += [HourGlass(256, 4)]
self.top_m += [ConvBlock(256, 256)]
self.conv_last += [nn.Conv2D(256, 256, kernel_size=1, strides=1, padding='VALID')]
self.bn_end += [nn.BatchNorm2D(256)]
self.l += [nn.Conv2D(256, 68, kernel_size=1, strides=1, padding='VALID')]
if i < 4 - 1:
self.bl += [nn.Conv2D(256, 256, kernel_size=1, strides=1, padding='VALID')]
self.al += [nn.Conv2D(68, 256, kernel_size=1, strides=1, padding='VALID')]
def forward(self, inp):
x, = inp
x = self.conv1(x)
x = self.bn1(x)
x = tf.nn.relu(x)
x = self.conv2(x)
x = tf.nn.avg_pool(x, [1, 2, 2, 1], [1, 2, 2, 1], 'VALID')
x = self.conv3(x)
x = self.conv4(x)
outputs = []
previous = x
for i in range(4):
ll = self.m[i](previous)
ll = self.top_m[i](ll)
ll = self.conv_last[i](ll)
ll = self.bn_end[i](ll)
ll = tf.nn.relu(ll)
tmp_out = self.l[i](ll)
outputs.append(tmp_out)
if i < 4 - 1:
ll = self.bl[i](ll)
previous = previous + ll + self.al[i](tmp_out)
x = outputs[-1]
x = tf.transpose(x, (0, 3, 1, 2))
return x
e = None
if place_model_on_cpu:
e = tf.device("/CPU:0")
if e is not None:
e.__enter__()
try:
self.model = FAN()
self.model.load_weights(str(model_path))
finally:
if e is not None:
e.__exit__(None, None, None)
self.model.build_for_run([(tf.float32, (None, 256, 256, 3))])
self.gpu_memory_limit = self._get_gpu_memory_limit()
self.current_memory_usage = 0
def _get_gpu_memory_limit(self):
import tensorflow as tf
try:
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
details = tf.config.experimental.get_device_details(gpus[0])
if 'device_limitations' in details and 'total_memory' in details['device_limitations']:
return details['device_limitations']['total_memory'] // (1024 ** 2)
else:
return 8192
except:
return 8192
return 8192
def _estimate_memory_usage(self, batch_size, avg_image_size):
memory_per_image = (avg_image_size * 3 * 4) / (1024 ** 2) # 图像数据 (float32)
model_overhead = 150 # FAN模型激活值和中间层的更大开销
return batch_size * (memory_per_image + model_overhead)
def _get_current_gpu_memory_usage(self):
try:
import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
try:
import pynvml
pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
info = pynvml.nvmlDeviceGetMemoryInfo(handle)
return info.used // (1024 ** 2) # 转换为MB
except ImportError:
return self.current_memory_usage
return 0
except:
return self.current_memory_usage
def extract_batch(self, input_images, all_rects, second_pass_extractor=None,
is_bgr=True, multi_sample=False, initial_batch_size=4, memory_safety_margin=0.7):
if not input_images or not all_rects or len(input_images) != len(all_rects):
return [[] for _ in range(len(input_images))]
all_results = []
avg_image_size = np.mean([img.shape[0] * img.shape[1] for img in input_images])
current_batch_size = initial_batch_size
max_retries = 3
retry_count = 0
i = 0
while i < len(input_images):
current_batch_size = self._adjust_batch_size(
current_batch_size, avg_image_size, memory_safety_margin
)
end_idx = min(i + current_batch_size, len(input_images))
batch_images = input_images[i:end_idx]
batch_rects = all_rects[i:end_idx]
batch_landmarks = []
success = True
try:
for img, rects in zip(batch_images, batch_rects):
landmarks = self._extract_single_optimized(img, rects, second_pass_extractor, is_bgr, multi_sample)
batch_landmarks.append(landmarks)
self.current_memory_usage += 100 # 估算每次处理约100MB
all_results.extend(batch_landmarks)
i = end_idx
retry_count = 0 # 重置重试计数
except Exception as e:
if "out of memory" in str(e).lower() or "memory" in str(e).lower():
if retry_count < max_retries:
current_batch_size = max(1, current_batch_size // 2)
retry_count += 1
print(f"FANExtractor: 内存不足，减少批大小至 {current_batch_size}")
continue
else:
print(f"FANExtractor: 无法处理图像，跳过。错误: {e}")
for _ in range(len(batch_images)):
all_results.append(None)
i = end_idx
retry_count = 0
else:
print(f"FANExtractor: 处理图像时发生错误: {e}")
for _ in range(len(batch_images)):
all_results.append(None)
i = end_idx
retry_count = 0
return all_results
def _adjust_batch_size(self, current_batch_size, avg_image_size, safety_margin):
if self.gpu_memory_limit == 0: # CPU模式
return min(current_batch_size, 2)
current_usage = self._get_current_gpu_memory_usage()
available_memory = self.gpu_memory_limit * safety_margin - current_usage
required_memory = self._estimate_memory_usage(current_batch_size, avg_image_size)
if required_memory > available_memory and current_batch_size > 1:
adjusted_batch_size = int(current_batch_size * (available_memory / required_memory))
adjusted_batch_size = max(1, adjusted_batch_size)
return min(current_batch_size, adjusted_batch_size)
return current_batch_size
def _process_batch(self, input_images, all_rects, second_pass_extractor=None, is_bgr=True, multi_sample=False):
batch_results = []
for img, rects in zip(input_images, all_rects):
landmarks = self._extract_single_optimized(img, rects, second_pass_extractor, is_bgr, multi_sample)
batch_results.append(landmarks)
return batch_results
def _extract_single_optimized(self, input_image, rects, second_pass_extractor=None, is_bgr=True, multi_sample=False):
if len(rects) == 0:
return []
if is_bgr:
input_image = input_image[:, :, ::-1]
is_bgr = False
(h, w, ch) = input_image.shape
landmarks = []
for (left, top, right, bottom) in rects:
scale = (right - left + bottom - top) / 195.0
center = np.array([(left + right) / 2.0, (top + bottom) / 2.0])
centers = [center]
if multi_sample:
centers += [center + [-1, -1],
center + [1, -1],
center + [1, 1],
center + [-1, 1],
]
images = []
ptss = []
try:
for c in centers:
images += [self.crop(input_image, c, scale)]
images = np.stack(images)
images = images.astype(np.float32) / 255.0
predicted = []
for i in range(len(images)):
pred_result = self.model.run([images[i][None, ...]])
predicted.append(pred_result[0])
for i, pred in enumerate(predicted):
ptss += [self.get_pts_from_predict(pred, centers[i], scale)]
pts_img = np.mean(np.array(ptss), 0)
landmarks.append(pts_img)
except Exception as e:
print(f"Error in FANExtractor.extract: {str(e)}")
landmarks.append(None)
if second_pass_extractor is not None and any(lmrks is not None for lmrks in landmarks):
for i, lmrks in enumerate(landmarks):
if lmrks is not None:
try:
image_to_face_mat = LandmarksProcessor.get_transform_mat(lmrks, 256, FaceType.FULL)
face_image = cv2.warpAffine(input_image, image_to_face_mat, (256, 256), cv2.INTER_CUBIC)
rects2 = second_pass_extractor.extract(face_image, is_bgr=is_bgr)
if len(rects2) == 1: # 只有在检测到1个面部时才进行二次提取
lmrks2 = self.extract(face_image, [rects2[0]], is_bgr=is_bgr, multi_sample=True)[0]
landmarks[i] = LandmarksProcessor.transform_points(lmrks2, image_to_face_mat, True)
except Exception as e:
print(f"Error in second pass extraction: {str(e)}")
pass
return landmarks
def extract(self, input_image, rects, second_pass_extractor=None, is_bgr=True, multi_sample=False):
return self._extract_single_optimized(input_image, rects, second_pass_extractor, is_bgr, multi_sample)
def transform(self, point, center, scale, resolution):
pt = np.array([point[0], point[1], 1.0])
h = 200.0 * scale
m = np.eye(3)
m[0, 0] = resolution / h
m[1, 1] = resolution / h
m[0, 2] = resolution * (-center[0] / h + 0.5)
m[1, 2] = resolution * (-center[1] / h + 0.5)
m = np.linalg.inv(m)
return np.matmul(m, pt)[0:2]
def crop(self, image, center, scale, resolution=256.0):
ul = self.transform([1, 1], center, scale, resolution).astype(np.int32)
br = self.transform([resolution, resolution], center, scale, resolution).astype(np.int32)
ht, wd = image.shape[0], image.shape[1]
newX = np.array([max(1, -ul[0] + 1), min(br[0], wd) - ul[0]], dtype=np.int32)
newY = np.array([max(1, -ul[1] + 1), min(br[1], ht) - ul[1]], dtype=np.int32)
oldX = np.array([max(1, ul[0] + 1), min(br[0], wd)], dtype=np.int32)
oldY = np.array([max(1, ul[1] + 1), min(br[1], ht)], dtype=np.int32)
if image.ndim > 2:
newDim = np.array([br[1] - ul[1], br[0] - ul[0], image.shape[2]], dtype=np.int32)
newImg = np.zeros(newDim, dtype=np.uint8)
else:
newDim = np.array([br[1] - ul[1], br[0] - ul[0]], dtype=np.int32)
newImg = np.zeros(newDim, dtype=np.uint8)
newImg[newY[0] - 1:newY[1], newX[0] - 1:newX[1]] = image[oldY[0] - 1:oldY[1], oldX[0] - 1:oldX[1], :]
newImg = cv2.resize(newImg, dsize=(int(resolution), int(resolution)), interpolation=cv2.INTER_LINEAR)
return newImg
def get_pts_from_predict(self, a, center, scale):
a_ch, a_h, a_w = a.shape
b = a.reshape((a_ch, a_h * a_w))
c = b.argmax(1).reshape((a_ch, 1)).repeat(2, axis=1).astype(np.float32)
c[:, 0] %= a_w
c[:, 1] = np.floor(c[:, 1] / a_w)
for i in range(a_ch):
pX, pY = int(c[i, 0]), int(c[i, 1])
if 0 < pX < a_w - 1 and 0 < pY < a_h - 1:
diff = np.array([a[i, pY, pX + 1] - a[i, pY, pX - 1],
a[i, pY + 1, pX] - a[i, pY - 1, pX]])
c[i] += np.sign(diff) * 0.25
c += 0.5
return np.array([self.transform(c[i], center, scale, a_w) for i in range(a_ch)])

复制代码

第二个文件 S3FDExtractor.py

import operator
from pathlib import Path
import cv2
import numpy as np
from core.leras import nn
class S3FDExtractor(object):
def __init__(self, place_model_on_cpu=False):
nn.initialize(data_format="NHWC")
tf = nn.tf
model_path = Path(__file__).parent / "S3FD.npy"
if not model_path.exists():
raise Exception("Unable to load S3FD.npy")
class L2Norm(nn.LayerBase):
def __init__(self, n_channels, **kwargs):
self.n_channels = n_channels
super().__init__(**kwargs)
def build_weights(self):
self.weight = tf.get_variable("weight", (1, 1, 1, self.n_channels), dtype=nn.floatx, initializer=tf.initializers.ones)
def get_weights(self):
return [self.weight]
def __call__(self, inputs):
x = inputs
x = x / (tf.sqrt(tf.reduce_sum(tf.pow(x, 2), axis=-1, keepdims=True)) + 1e-10) * self.weight
return x
class S3FD(nn.ModelBase):
def __init__(self):
super().__init__(name='S3FD')
def on_build(self):
self.minus = tf.constant([104, 117, 123], dtype=nn.floatx)
self.conv1_1 = nn.Conv2D(3, 64, kernel_size=3, strides=1, padding='SAME')
self.conv1_2 = nn.Conv2D(64, 64, kernel_size=3, strides=1, padding='SAME')
self.conv2_1 = nn.Conv2D(64, 128, kernel_size=3, strides=1, padding='SAME')
self.conv2_2 = nn.Conv2D(128, 128, kernel_size=3, strides=1, padding='SAME')
self.conv3_1 = nn.Conv2D(128, 256, kernel_size=3, strides=1, padding='SAME')
self.conv3_2 = nn.Conv2D(256, 256, kernel_size=3, strides=1, padding='SAME')
self.conv3_3 = nn.Conv2D(256, 256, kernel_size=3, strides=1, padding='SAME')
self.conv4_1 = nn.Conv2D(256, 512, kernel_size=3, strides=1, padding='SAME')
self.conv4_2 = nn.Conv2D(512, 512, kernel_size=3, strides=1, padding='SAME')
self.conv4_3 = nn.Conv2D(512, 512, kernel_size=3, strides=1, padding='SAME')
self.conv5_1 = nn.Conv2D(512, 512, kernel_size=3, strides=1, padding='SAME')
self.conv5_2 = nn.Conv2D(512, 512, kernel_size=3, strides=1, padding='SAME')
self.conv5_3 = nn.Conv2D(512, 512, kernel_size=3, strides=1, padding='SAME')
self.fc6 = nn.Conv2D(512, 1024, kernel_size=3, strides=1, padding=3)
self.fc7 = nn.Conv2D(1024, 1024, kernel_size=1, strides=1, padding='SAME')
self.conv6_1 = nn.Conv2D(1024, 256, kernel_size=1, strides=1, padding='SAME')
self.conv6_2 = nn.Conv2D(256, 512, kernel_size=3, strides=2, padding='SAME')
self.conv7_1 = nn.Conv2D(512, 128, kernel_size=1, strides=1, padding='SAME')
self.conv7_2 = nn.Conv2D(128, 256, kernel_size=3, strides=2, padding='SAME')
self.conv3_3_norm = L2Norm(256)
self.conv4_3_norm = L2Norm(512)
self.conv5_3_norm = L2Norm(512)
self.conv3_3_norm_mbox_conf = nn.Conv2D(256, 4, kernel_size=3, strides=1, padding='SAME')
self.conv3_3_norm_mbox_loc = nn.Conv2D(256, 4, kernel_size=3, strides=1, padding='SAME')
self.conv4_3_norm_mbox_conf = nn.Conv2D(512, 2, kernel_size=3, strides=1, padding='SAME')
self.conv4_3_norm_mbox_loc = nn.Conv2D(512, 4, kernel_size=3, strides=1, padding='SAME')
self.conv5_3_norm_mbox_conf = nn.Conv2D(512, 2, kernel_size=3, strides=1, padding='SAME')
self.conv5_3_norm_mbox_loc = nn.Conv2D(512, 4, kernel_size=3, strides=1, padding='SAME')
self.fc7_mbox_conf = nn.Conv2D(1024, 2, kernel_size=3, strides=1, padding='SAME')
self.fc7_mbox_loc = nn.Conv2D(1024, 4, kernel_size=3, strides=1, padding='SAME')
self.conv6_2_mbox_conf = nn.Conv2D(512, 2, kernel_size=3, strides=1, padding='SAME')
self.conv6_2_mbox_loc = nn.Conv2D(512, 4, kernel_size=3, strides=1, padding='SAME')
self.conv7_2_mbox_conf = nn.Conv2D(256, 2, kernel_size=3, strides=1, padding='SAME')
self.conv7_2_mbox_loc = nn.Conv2D(256, 4, kernel_size=3, strides=1, padding='SAME')
def forward(self, inp):
x, = inp
x = x - self.minus
x = tf.nn.relu(self.conv1_1(x))
x = tf.nn.relu(self.conv1_2(x))
x = tf.nn.max_pool(x, [1, 2, 2, 1], [1, 2, 2, 1], "VALID")
x = tf.nn.relu(self.conv2_1(x))
x = tf.nn.relu(self.conv2_2(x))
x = tf.nn.max_pool(x, [1, 2, 2, 1], [1, 2, 2, 1], "VALID")
x = tf.nn.relu(self.conv3_1(x))
x = tf.nn.relu(self.conv3_2(x))
x = tf.nn.relu(self.conv3_3(x))
f3_3 = x
x = tf.nn.max_pool(x, [1, 2, 2, 1], [1, 2, 2, 1], "VALID")
x = tf.nn.relu(self.conv4_1(x))
x = tf.nn.relu(self.conv4_2(x))
x = tf.nn.relu(self.conv4_3(x))
f4_3 = x
x = tf.nn.max_pool(x, [1, 2, 2, 1], [1, 2, 2, 1], "VALID")
x = tf.nn.relu(self.conv5_1(x))
x = tf.nn.relu(self.conv5_2(x))
x = tf.nn.relu(self.conv5_3(x))
f5_3 = x
x = tf.nn.max_pool(x, [1, 2, 2, 1], [1, 2, 2, 1], "VALID")
x = tf.nn.relu(self.fc6(x))
x = tf.nn.relu(self.fc7(x))
ffc7 = x
x = tf.nn.relu(self.conv6_1(x))
x = tf.nn.relu(self.conv6_2(x))
f6_2 = x
x = tf.nn.relu(self.conv7_1(x))
x = tf.nn.relu(self.conv7_2(x))
f7_2 = x
f3_3 = self.conv3_3_norm(f3_3)
f4_3 = self.conv4_3_norm(f4_3)
f5_3 = self.conv5_3_norm(f5_3)
cls1 = self.conv3_3_norm_mbox_conf(f3_3)
reg1 = self.conv3_3_norm_mbox_loc(f3_3)
cls2 = tf.nn.softmax(self.conv4_3_norm_mbox_conf(f4_3))
reg2 = self.conv4_3_norm_mbox_loc(f4_3)
cls3 = tf.nn.softmax(self.conv5_3_norm_mbox_conf(f5_3))
reg3 = self.conv5_3_norm_mbox_loc(f5_3)
cls4 = tf.nn.softmax(self.fc7_mbox_conf(ffc7))
reg4 = self.fc7_mbox_loc(ffc7)
cls5 = tf.nn.softmax(self.conv6_2_mbox_conf(f6_2))
reg5 = self.conv6_2_mbox_loc(f6_2)
cls6 = tf.nn.softmax(self.conv7_2_mbox_conf(f7_2))
reg6 = self.conv7_2_mbox_loc(f7_2)
# max-out background label
bmax = tf.maximum(tf.maximum(cls1[:, :, :, 0:1], cls1[:, :, :, 1:2]), cls1[:, :, :, 2:3])
cls1 = tf.concat([bmax, cls1[:, :, :, 3:4]], axis=-1)
cls1 = tf.nn.softmax(cls1)
return [cls1, reg1, cls2, reg2, cls3, reg3, cls4, reg4, cls5, reg5, cls6, reg6]
e = None
if place_model_on_cpu:
e = tf.device("/CPU:0")
if e is not None:
e.__enter__()
try:
self.model = S3FD()
self.model.load_weights(model_path)
finally:
if e is not None:
e.__exit__(None, None, None)
self.model.build_for_run([(tf.float32, nn.get4Dshape(None, None, 3))])
self.gpu_memory_limit = self._get_gpu_memory_limit()
self.current_memory_usage = 0
def _get_gpu_memory_limit(self):
import tensorflow as tf
try:
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
details = tf.config.experimental.get_device_details(gpus[0])
if 'device_limitations' in details and 'total_memory' in details['device_limitations']:
return details['device_limitations']['total_memory'] // (1024 ** 2)
else:
return 8192
except:
return 8192
return 8192
def _estimate_memory_usage(self, batch_size, avg_image_size):
memory_per_image = (avg_image_size * 3 * 4) / (1024 ** 2) # 图像数据 (float32)
model_overhead = 50 # 模型激活值和中间层的开销
return batch_size * (memory_per_image + model_overhead)
def _get_current_gpu_memory_usage(self):
try:
import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
try:
import pynvml
pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
info = pynvml.nvmlDeviceGetMemoryInfo(handle)
return info.used // (1024 ** 2) # 转换为MB
except ImportError:
return self.current_memory_usage
return 0
except:
return self.current_memory_usage
def extract_batch(self, input_images, is_bgr=True, is_remove_intersects=False,
initial_batch_size=8, memory_safety_margin=0.8):
if not input_images:
return [[] for _ in range(len(input_images))]
all_results = []
avg_image_size = np.mean([img.shape[0] * img.shape[1] for img in input_images])
current_batch_size = initial_batch_size
max_retries = 3
retry_count = 0
i = 0
while i < len(input_images):
current_batch_size = self._adjust_batch_size(
current_batch_size, avg_image_size, memory_safety_margin
)
end_idx = min(i + current_batch_size, len(input_images))
batch = input_images[i:end_idx]
try:
processed_batch = []
scales = []
for img in batch:
if is_bgr:
img = img[:, :, ::-1]
h, w, ch = img.shape
d = max(w, h)
scale_to = 640 if d >= 1280 else d / 2
scale_to = max(64, scale_to)
input_scale = d / scale_to
resized_img = cv2.resize(img, (int(w / input_scale), int(h / input_scale)),
interpolation=cv2.INTER_LINEAR)
processed_batch.append(resized_img)
scales.append(input_scale)
batch_tensor = np.stack(processed_batch).astype(np.float32)
self.current_memory_usage += len(batch) * 50 # 估算每张图约50MB
olist_batch = self.model.run([batch_tensor])
for idx in range(len(batch)):
single_olist = [tensor[idx:idx + 1] for tensor in olist_batch]
detected_faces = []
for ltrb in self.refine(single_olist):
l, t, r, b = [x * scales[idx] for x in ltrb]
bt = b - t
if min(r - l, bt) < 40:
continue
b += bt * 0.1
detected_faces.append([int(x) for x in (l, t, r, b)])
detected_faces = [[(l, t, r, b), (r - l) * (b - t)] for (l, t, r, b) in detected_faces]
detected_faces = sorted(detected_faces, key=operator.itemgetter(1), reverse=True)
detected_faces = [x[0] for x in detected_faces]
if is_remove_intersects:
detected_faces = self._remove_intersecting_faces_optimized(detected_faces)
all_results.append(detected_faces)
i = end_idx
retry_count = 0 # 重置重试计数
except Exception as e:
if "out of memory" in str(e).lower() or "memory" in str(e).lower():
if retry_count < max_retries:
current_batch_size = max(1, current_batch_size // 2)
retry_count += 1
print(f"S3FDExtractor: 内存不足，减少批大小至 {current_batch_size}")
continue
else:
print(f"S3FDExtractor: 无法处理图像，跳过。错误: {e}")
for _ in range(len(batch)):
all_results.append([])
i = end_idx
retry_count = 0
else:
print(f"S3FDExtractor: 处理图像时发生错误: {e}")
for _ in range(len(batch)):
all_results.append([])
i = end_idx
retry_count = 0
return all_results
def _adjust_batch_size(self, current_batch_size, avg_image_size, safety_margin):
if self.gpu_memory_limit == 0: # CPU模式
return min(current_batch_size, 4)
current_usage = self._get_current_gpu_memory_usage()
available_memory = self.gpu_memory_limit * safety_margin - current_usage
required_memory = self._estimate_memory_usage(current_batch_size, avg_image_size)
if required_memory > available_memory and current_batch_size > 1:
adjusted_batch_size = int(current_batch_size * (available_memory / required_memory))
adjusted_batch_size = max(1, adjusted_batch_size)
return min(current_batch_size, adjusted_batch_size)
return current_batch_size
def _process_batch(self, input_images, is_bgr=True, is_remove_intersects=False):
batch_results = []
for img in input_images:
result = self.extract_single(img, is_bgr, is_remove_intersects)
batch_results.append(result)
return batch_results
def extract_single(self, input_image, is_bgr=True, is_remove_intersects=False):
if is_bgr:
input_image = input_image[:, :, ::-1]
is_bgr = False
(h, w, ch) = input_image.shape
d = max(w, h)
scale_to = 640 if d >= 1280 else d / 2
scale_to = max(64, scale_to)
input_scale = d / scale_to
input_image = cv2.resize(input_image, (int(w / input_scale), int(h / input_scale)), interpolation=cv2.INTER_LINEAR)
olist = self.model.run([input_image[None, ...]])
detected_faces = []
for ltrb in self.refine(olist):
l, t, r, b = [x * input_scale for x in ltrb]
bt = b - t
if min(r - l, bt) < 40: # filtering faces < 40pix by any side
continue
b += bt * 0.1 # enlarging bottom line a bit for 2DFAN-4
detected_faces.append([int(x) for x in (l, t, r, b)])
# sort by largest area first
detected_faces = [[(l, t, r, b), (r - l) * (b - t)] for (l, t, r, b) in detected_faces]
detected_faces = sorted(detected_faces, key=operator.itemgetter(1), reverse=True)
detected_faces = [x[0] for x in detected_faces]
if is_remove_intersects:
detected_faces = self._remove_intersecting_faces_optimized(detected_faces)
return detected_faces
def _remove_intersecting_faces_optimized(self, detected_faces):
if len(detected_faces) <= 1:
return detected_faces
bboxes = np.array(detected_faces)
areas = (bboxes[:, 2] - bboxes[:, 0]) * (bboxes[:, 3] - bboxes[:, 1])
x1, y1, x2, y2 = bboxes[:, 0], bboxes[:, 1], bboxes[:, 2], bboxes[:, 3]
xx1 = np.maximum(x1[:, None], x1[None, :])
yy1 = np.maximum(y1[:, None], y1[None, :])
xx2 = np.minimum(x2[:, None], x2[None, :])
yy2 = np.minimum(y2[:, None], y2[None, :])
w = np.maximum(0, xx2 - xx1)
h = np.maximum(0, yy2 - yy1)
inter = w * h
union = areas[:, None] + areas[None, :] - inter
iou = inter / union
to_remove = set()
for i in range(len(bboxes)):
if i in to_remove:
continue
overlap_indices = np.where((iou[i, :] > 0.3) & (np.arange(len(bboxes)) != i))[0]
for j in overlap_indices:
if j in to_remove:
continue
if areas[i] < areas[j]:
to_remove.add(i)
break
else:
to_remove.add(j)
for idx in sorted(to_remove, reverse=True):
if idx < len(detected_faces):
detected_faces.pop(idx)
return detected_faces
def extract(self, input_image, is_bgr=True, is_remove_intersects=False):
return self.extract_single(input_image, is_bgr, is_remove_intersects)
def refine(self, olist):
bboxlist = []
for i, ((ocls,), (oreg,)) in enumerate(zip(olist[::2], olist[1::2])):
stride = 2 ** (i + 2) # 4,8,16,32,64,128
s_d2 = stride / 2
s_m4 = stride * 4
scores = ocls[..., 1]
high_score_indices = np.where(scores > 0.05)
for hindex, windex in zip(*high_score_indices):
score = scores[hindex, windex]
loc = oreg[hindex, windex, :]
priors = np.array([windex * stride + s_d2, hindex * stride + s_d2, s_m4, s_m4])
priors_2p = priors[2:]
box = np.concatenate((priors[:2] + loc[:2] * 0.1 * priors_2p,
priors_2p * np.exp(loc[2:] * 0.2)))
box[:2] -= box[2:] / 2
box[2:] += box[:2]
bboxlist.append([*box, score])
bboxlist = np.array(bboxlist)
if len(bboxlist) == 0:
bboxlist = np.zeros((1, 5))
bboxlist = bboxlist[self.refine_nms(bboxlist, 0.3), :]
bboxlist = [x[:-1].astype(np.int32) for x in bboxlist if x[-1] >= 0.5]
return bboxlist
def refine_nms(self, dets, thresh):
if len(dets) == 0:
return []
dets = np.asarray(dets)
x1 = dets[:, 0]
y1 = dets[:, 1]
x2 = dets[:, 2]
y2 = dets[:, 3]
scores = dets[:, 4]
areas = (x2 - x1 + 1) * (y2 - y1 + 1)
order = scores.argsort()[::-1]
keep = []
while order.size > 0:
i = order[0]
keep.append(i)
xx1 = np.maximum(x1[i], x1[order[1:]])
yy1 = np.maximum(y1[i], y1[order[1:]])
xx2 = np.minimum(x2[i], x2[order[1:]])
yy2 = np.minimum(y2[i], y2[order[1:]])
w = np.maximum(0.0, xx2 - xx1 + 1)
h = np.maximum(0.0, yy2 - yy1 + 1)
inter = w * h
ovr = inter / (areas[i] + areas[order[1:]] - inter)
inds = np.where(ovr <= thresh)[0]
order = order[inds + 1]
return keep

复制代码

day270010678 · 发表于 5 天前

最后就是主调用文件Extractor.py

import traceback
import math
import multiprocessing
import operator
import os
import shutil
import sys
import time
from pathlib import Path
import cv2
import numpy as np
from numpy import linalg as npla
import facelib
from core import imagelib
from core import mathlib
from facelib import FaceType, LandmarksProcessor
from core.interact import interact as io
from core.joblib import Subprocessor
from core.leras import nn
from core import pathex
from core.cv2ex import *
from DFLIMG import *
DEBUG = False
class ExtractSubprocessor(Subprocessor):
class Data(object):
def __init__(self, filepath=None, rects=None, landmarks=None, landmarks_accurate=True, manual=False, force_output_path=None, final_output_files=None):
self.filepath = filepath
self.rects = rects or []
self.rects_rotation = 0
self.landmarks_accurate = landmarks_accurate
self.manual = manual
self.landmarks = landmarks or []
self.force_output_path = force_output_path
self.final_output_files = final_output_files or []
self.faces_detected = 0
class Cli(Subprocessor.Cli):
# override
def on_initialize(self, client_dict):
self.type = client_dict['type']
self.image_size = client_dict['image_size']
self.jpeg_quality = client_dict['jpeg_quality']
self.face_type = client_dict['face_type']
self.max_faces_from_image = client_dict['max_faces_from_image']
self.device_idx = client_dict['device_idx']
self.cpu_only = client_dict['device_type'] == 'CPU'
self.final_output_path = client_dict['final_output_path']
self.output_debug_path = client_dict['output_debug_path']
# transfer and set stdin in order to work code.interact in debug subprocess
stdin_fd = client_dict['stdin_fd']
if stdin_fd is not None and DEBUG:
sys.stdin = os.fdopen(stdin_fd)
if self.cpu_only:
device_config = nn.DeviceConfig.CPU()
place_model_on_cpu = True
else:
device_config = nn.DeviceConfig.GPUIndexes([self.device_idx])
place_model_on_cpu = device_config.devices[0].total_mem_gb < 4
if self.type == 'all' or 'rects' in self.type or 'landmarks' in self.type:
nn.initialize(device_config)
self.log_info(f"Running on {client_dict['device_name']}")
if self.type == 'all' or self.type == 'rects-s3fd' or 'landmarks' in self.type:
self.rects_extractor = facelib.S3FDExtractor(place_model_on_cpu=place_model_on_cpu)
if self.type == 'all' or 'landmarks' in self.type:
# for head type, extract "3D landmarks"
self.landmarks_extractor = facelib.FANExtractor(landmarks_3D=self.face_type >= FaceType.HEAD,
place_model_on_cpu=place_model_on_cpu)
# 缓存图像，减少重复读取
self.cached_image = (None, None)
# override
def process_data(self, data):
if 'landmarks' in self.type and len(data.rects) == 0:
return data
filepath = data.filepath
cached_filepath, image = self.cached_image
if cached_filepath != filepath:
image = cv2_imread(filepath)
if image is None:
self.log_err(f'Failed to open {filepath}, reason: cv2_imread() fail.')
return data
image = imagelib.normalize_channels(image, 3)
image = imagelib.cut_odd_image(image)
self.cached_image = (filepath, image)
h, w, c = image.shape
if 'rects' in self.type or self.type == 'all':
data = ExtractSubprocessor.Cli.rects_stage(data=data,
image=image,
max_faces_from_image=self.max_faces_from_image,
rects_extractor=self.rects_extractor,
)
if 'landmarks' in self.type or self.type == 'all':
data = ExtractSubprocessor.Cli.landmarks_stage(data=data,
image=image,
landmarks_extractor=self.landmarks_extractor,
rects_extractor=self.rects_extractor,
)
if self.type == 'final' or self.type == 'all':
data = ExtractSubprocessor.Cli.final_stage(data=data,
image=image,
face_type=self.face_type,
image_size=self.image_size,
jpeg_quality=self.jpeg_quality,
output_debug_path=self.output_debug_path,
final_output_path=self.final_output_path,
)
return data
@staticmethod
def rects_stage(data,
image,
max_faces_from_image,
rects_extractor,
):
"""
优化的人脸矩形检测阶段
"""
h, w, c = image.shape
if min(h, w) < 128:
# Image is too small
data.rects = []
else:
# 使用旋转优化的检测
rotation_angles = [0, 90, 270, 180]
for rot in rotation_angles:
if rot == 0:
rotated_image = image
elif rot == 90:
rotated_image = image.swapaxes(0, 1)[:, ::-1, :]
elif rot == 180:
rotated_image = image[::-1, ::-1, :]
elif rot == 270:
rotated_image = image.swapaxes(0, 1)[::-1, :, :]
rects = data.rects = rects_extractor.extract(rotated_image, is_bgr=True)
if len(rects) != 0:
data.rects_rotation = rot
break
if max_faces_from_image is not None and \
max_faces_from_image > 0 and \
len(data.rects) > 0:
data.rects = data.rects[0:max_faces_from_image]
return data
@staticmethod
def landmarks_stage(data,
image,
landmarks_extractor,
rects_extractor,
):
"""
优化的关键点提取阶段
"""
h, w, ch = image.shape
if data.rects_rotation == 0:
rotated_image = image
elif data.rects_rotation == 90:
rotated_image = image.swapaxes(0, 1)[:, ::-1, :]
elif data.rects_rotation == 180:
rotated_image = image[::-1, ::-1, :]
elif data.rects_rotation == 270:
rotated_image = image.swapaxes(0, 1)[::-1, :, :]
data.landmarks = landmarks_extractor.extract(rotated_image, data.rects,
rects_extractor if (data.landmarks_accurate) else None,
is_bgr=True)
# 优化旋转后的关键点调整
if data.rects_rotation != 0:
for i, (rect, lmrks) in enumerate(zip(data.rects, data.landmarks)):
new_rect, new_lmrks = rect, lmrks
(l, t, r, b) = rect
if data.rects_rotation == 90:
new_rect = (t, h - l, b, h - r)
if lmrks is not None:
new_lmrks = lmrks[:, ::-1].copy()
new_lmrks[:, 1] = h - new_lmrks[:, 1]
elif data.rects_rotation == 180:
if lmrks is not None:
new_rect = (w - l, h - t, w - r, h - b)
new_lmrks = lmrks.copy()
new_lmrks[:, 0] = w - new_lmrks[:, 0]
new_lmrks[:, 1] = h - new_lmrks[:, 1]
elif data.rects_rotation == 270:
new_rect = (w - b, l, w - t, r)
if lmrks is not None:
new_lmrks = lmrks[:, ::-1].copy()
new_lmrks[:, 0] = w - new_lmrks[:, 0]
data.rects[i], data.landmarks[i] = new_rect, new_lmrks
return data
@staticmethod
def final_stage(data,
image,
face_type,
image_size,
jpeg_quality,
output_debug_path=None,
final_output_path=None,
):
"""
优化的最终处理阶段
"""
data.final_output_files = []
filepath = data.filepath
rects = data.rects
landmarks = data.landmarks
if output_debug_path is not None:
debug_image = image.copy()
face_idx = 0
for rect, image_landmarks in zip(rects, landmarks):
if image_landmarks is None:
continue
rect = np.array(rect)
if face_type == FaceType.MARK_ONLY:
image_to_face_mat = None
face_image = image
face_image_landmarks = image_landmarks
else:
image_to_face_mat = LandmarksProcessor.get_transform_mat(image_landmarks, image_size, face_type)
face_image = cv2.warpAffine(image, image_to_face_mat, (image_size, image_size), cv2.INTER_LANCZOS4)
face_image_landmarks = LandmarksProcessor.transform_points(image_landmarks, image_to_face_mat)
landmarks_bbox = LandmarksProcessor.transform_points([(0, 0), (0, image_size - 1), (image_size - 1, image_size - 1), (image_size - 1, 0)],
image_to_face_mat, True)
rect_area = mathlib.polygon_area(np.array(rect[[0, 2, 2, 0]]).astype(np.float32), np.array(rect[[1, 1, 3, 3]]).astype(np.float32))
landmarks_area = mathlib.polygon_area(landmarks_bbox[:, 0].astype(np.float32), landmarks_bbox[:, 1].astype(np.float32))
if not data.manual and face_type <= FaceType.FULL_NO_ALIGN and landmarks_area > 4 * rect_area: # get rid of faces which umeyama-landmark-area > 4*detector-rect-area
continue
if output_debug_path is not None:
LandmarksProcessor.draw_rect_landmarks(debug_image, rect, image_landmarks, face_type, image_size, transparent_mask=True)
output_path = final_output_path
if data.force_output_path is not None:
output_path = data.force_output_path
output_filepath = output_path / f"{filepath.stem}_{face_idx}.jpg"
cv2_imwrite(output_filepath, face_image, [int(cv2.IMWRITE_JPEG_QUALITY), jpeg_quality])
dflimg = DFLJPG.load(output_filepath)
dflimg.set_face_type(FaceType.toString(face_type))
dflimg.set_landmarks(face_image_landmarks.tolist())
dflimg.set_source_filename(filepath.name)
dflimg.set_source_rect(rect)
dflimg.set_source_landmarks(image_landmarks.tolist())
dflimg.set_image_to_face_mat(image_to_face_mat)
dflimg.save()
data.final_output_files.append(output_filepath)
face_idx += 1
data.faces_detected = face_idx
if output_debug_path is not None:
cv2_imwrite(output_debug_path / (filepath.stem + '.jpg'), debug_image, [int(cv2.IMWRITE_JPEG_QUALITY), 50])
return data
# overridable
def get_data_name(self, data):
# return string identificator of your data
return data.filepath
@staticmethod
def get_devices_for_config(type, device_config):
devices = device_config.devices
cpu_only = len(devices) == 0
if 'rects' in type or \
'landmarks' in type or \
'all' in type:
if not cpu_only:
if type == 'landmarks-manual':
devices = [devices.get_best_device()]
result = []
for device in devices:
count = 1
if count == 1:
result += [(device.index, 'GPU', device.name, device.total_mem_gb)]
else:
for i in range(count):
result += [(device.index, 'GPU', f"{device.name} #{i}", device.total_mem_gb)]
return result
else:
if type == 'landmarks-manual':
return [(0, 'CPU', 'CPU', 0)]
else:
return [(i, 'CPU', 'CPU%d' % (i), 0) for i in range(min(8, multiprocessing.cpu_count() // 2))]
elif type == 'final':
return [(i, 'CPU', 'CPU%d' % (i), 0) for i in (range(min(8, multiprocessing.cpu_count())) if not DEBUG else [0])]
def __init__(self, input_data, type, image_size=None, jpeg_quality=None, face_type=None, output_debug_path=None, manual_window_size=0, max_faces_from_image=0,
final_output_path=None, device_config=None):
if type == 'landmarks-manual':
for x in input_data:
x.manual = True
self.input_data = input_data
self.type = type
self.image_size = image_size
self.jpeg_quality = jpeg_quality
self.face_type = face_type
self.output_debug_path = output_debug_path
self.final_output_path = final_output_path
self.manual_window_size = manual_window_size
self.max_faces_from_image = max_faces_from_image
self.result = []
self.devices = ExtractSubprocessor.get_devices_for_config(self.type, device_config)
super().__init__('Extractor', ExtractSubprocessor.Cli,
999999 if type == 'landmarks-manual' or DEBUG else 120)
# override
def on_clients_initialized(self):
if self.type == 'landmarks-manual':
self.wnd_name = 'Manual pass'
io.named_window(self.wnd_name)
io.capture_mouse(self.wnd_name)
io.capture_keys(self.wnd_name)
self.cache_original_image = (None, None)
self.cache_image = (None, None)
self.cache_text_lines_img = (None, None)
self.hide_help = False
self.landmarks_accurate = True
self.force_landmarks = False
self.landmarks = None
self.x = 0
self.y = 0
self.rect_size = 100
self.rect_locked = False
self.extract_needed = True
self.image = None
self.image_filepath = None
io.progress_bar(None, len(self.input_data))
# override
def on_clients_finalized(self):
if self.type == 'landmarks-manual':
io.destroy_all_windows()
io.progress_bar_close()
# override
def process_info_generator(self):
base_dict = {'type': self.type,
'image_size': self.image_size,
'jpeg_quality': self.jpeg_quality,
'face_type': self.face_type,
'max_faces_from_image': self.max_faces_from_image,
'output_debug_path': self.output_debug_path,
'final_output_path': self.final_output_path,
'stdin_fd': sys.stdin.fileno()}
for (device_idx, device_type, device_name, device_total_vram_gb) in self.devices:
client_dict = base_dict.copy()
client_dict['device_idx'] = device_idx
client_dict['device_name'] = device_name
client_dict['device_type'] = device_type
yield client_dict['device_name'], {}, client_dict
# override
def get_data(self, host_dict):
if self.type == 'landmarks-manual':
need_remark_face = False
while len(self.input_data) > 0:
data = self.input_data[0]
filepath, data_rects, data_landmarks = data.filepath, data.rects, data.landmarks
is_frame_done = False
if self.image_filepath != filepath:
self.image_filepath = filepath
if self.cache_original_image[0] == filepath:
self.original_image = self.cache_original_image[1]
else:
self.original_image = imagelib.normalize_channels(cv2_imread(filepath), 3)
self.cache_original_image = (filepath, self.original_image)
(h, w, c) = self.original_image.shape
self.view_scale = 1.0 if self.manual_window_size == 0 else self.manual_window_size / (h * (16.0 / 9.0))
if self.cache_image[0] == (h, w, c) + (self.view_scale, filepath):
self.image = self.cache_image[1]
else:
self.image = cv2.resize(self.original_image, (int(w * self.view_scale), int(h * self.view_scale)), interpolation=cv2.INTER_LINEAR)
self.cache_image = ((h, w, c) + (self.view_scale, filepath), self.image)
(h, w, c) = self.image.shape
sh = (0, 0, w, min(100, h))
if self.cache_text_lines_img[0] == sh:
self.text_lines_img = self.cache_text_lines_img[1]
else:
self.text_lines_img = (imagelib.get_draw_text_lines(self.image, sh,
['[L Mouse click] - lock/unlock selection. [Mouse wheel] - change rect',
'[R Mouse Click] - manual face rectangle',
'[Enter] / [Space] - confirm / skip frame',
'[,] [.]- prev frame, next frame. [Q] - skip remaining frames',
'[a] - accuracy on/off (more fps)',
'[h] - hide this help'
], (1, 1, 1)) * 255).astype(np.uint8)
self.cache_text_lines_img = (sh, self.text_lines_img)
if need_remark_face: # need remark image from input data that already has a marked face?
need_remark_face = False
if len(data_rects) != 0: # If there was already a face then lock the rectangle to it until the mouse is clicked
self.rect = data_rects.pop()
self.landmarks = data_landmarks.pop()
data_rects.clear()
data_landmarks.clear()
self.rect_locked = True
self.rect_size = (self.rect[2] - self.rect[0]) / 2
self.x = (self.rect[0] + self.rect[2]) / 2
self.y = (self.rect[1] + self.rect[3]) / 2
self.redraw()
if len(data_rects) == 0:
(h, w, c) = self.image.shape
while True:
io.process_messages(0.0001)
if not self.force_landmarks:
new_x = self.x
new_y = self.y
new_rect_size = self.rect_size
mouse_events = io.get_mouse_events(self.wnd_name)
for ev in mouse_events:
(x, y, ev, flags) = ev
if ev == io.EVENT_MOUSEWHEEL and not self.rect_locked:
mod = 1 if flags > 0 else -1
diff = 1 if new_rect_size <= 40 else np.clip(new_rect_size / 10, 1, 10)
new_rect_size = max(5, new_rect_size + diff * mod)
elif ev == io.EVENT_LBUTTONDOWN:
if self.force_landmarks:
self.x = new_x
self.y = new_y
self.force_landmarks = False
self.rect_locked = True
self.redraw()
else:
self.rect_locked = not self.rect_locked
self.extract_needed = True
elif ev == io.EVENT_RBUTTONDOWN:
self.force_landmarks = not self.force_landmarks
if self.force_landmarks:
self.rect_locked = False
elif not self.rect_locked:
new_x = np.clip(x, 0, w - 1) / self.view_scale
new_y = np.clip(y, 0, h - 1) / self.view_scale
key_events = io.get_key_events(self.wnd_name)
key, chr_key, ctrl_pressed, alt_pressed, shift_pressed = key_events[-1] if len(key_events) > 0 else (0, 0, False, False, False)
if key == ord('\r') or key == ord('\n'):
# confirm frame
is_frame_done = True
data_rects.append(self.rect)
data_landmarks.append(self.landmarks)
break
elif key == ord(' '):
# confirm skip frame
is_frame_done = True
break
elif key == ord(',') and len(self.result) > 0:
# go prev frame
if self.rect_locked:
self.rect_locked = False
# Only save the face if the rect is still locked
data_rects.append(self.rect)
data_landmarks.append(self.landmarks)
self.input_data.insert(0, self.result.pop())
io.progress_bar_inc(-1)
need_remark_face = True
break
elif key == ord('.'):
# go next frame
if self.rect_locked:
self.rect_locked = False
# Only save the face if the rect is still locked
data_rects.append(self.rect)
data_landmarks.append(self.landmarks)
need_remark_face = True
is_frame_done = True
break
elif key == ord('q'):
# skip remaining
if self.rect_locked:
self.rect_locked = False
data.rects.append(self.rect)
data.landmarks.append(self.landmarks)
while len(self.input_data) > 0:
self.result.append(self.input_data.pop(0))
io.progress_bar_inc(1)
break
elif key == ord('h'):
self.hide_help = not self.hide_help
break
elif key == ord('a'):
self.landmarks_accurate = not self.landmarks_accurate
break
if self.force_landmarks:
pt2 = np.float32([new_x, new_y])
pt1 = np.float32([self.x, self.y])
pt_vec_len = npla.norm(pt2 - pt1)
pt_vec = pt2 - pt1
if pt_vec_len != 0:
pt_vec /= pt_vec_len
self.rect_size = pt_vec_len
self.rect = (int(self.x - self.rect_size),
int(self.y - self.rect_size),
int(self.x + self.rect_size),
int(self.y + self.rect_size))
if pt_vec_len > 0:
lmrks = np.concatenate((np.zeros((17, 2), np.float32), LandmarksProcessor.landmarks_2D), axis=0)
lmrks -= lmrks[30:31, :]
mat = cv2.getRotationMatrix2D((0, 0), -np.arctan2(pt_vec[1], pt_vec[0]) * 180 / math.pi, pt_vec_len)
mat[:, 2] += (self.x, self.y)
self.landmarks = LandmarksProcessor.transform_points(lmrks, mat)
self.redraw()
elif self.x != new_x or \
self.y != new_y or \
self.rect_size != new_rect_size or \
self.extract_needed:
self.x = new_x
self.y = new_y
self.rect_size = new_rect_size
self.rect = (int(self.x - self.rect_size),
int(self.y - self.rect_size),
int(self.x + self.rect_size),
int(self.y + self.rect_size))
return ExtractSubprocessor.Data(filepath, rects=[self.rect], landmarks_accurate=self.landmarks_accurate)
else:
is_frame_done = True
if is_frame_done:
self.result.append(data)
self.input_data.pop(0)
io.progress_bar_inc(1)
self.extract_needed = True
self.rect_locked = False
else:
if len(self.input_data) > 0:
return self.input_data.pop(0)
return None
# override
def on_data_return(self, host_dict, data):
if not self.type != 'landmarks-manual':
self.input_data.insert(0, data)
def redraw(self):
(h, w, c) = self.image.shape
if not self.hide_help:
image = cv2.addWeighted(self.image, 1.0, self.text_lines_img, 1.0, 0)
else:
image = self.image.copy()
view_rect = (np.array(self.rect) * self.view_scale).astype(np.int32).tolist()
view_landmarks = (np.array(self.landmarks) * self.view_scale).astype(np.int32).tolist()
if self.rect_size <= 40:
scaled_rect_size = h // 3 if h < w else w // 3
p1 = (self.x - self.rect_size, self.y - self.rect_size)
p2 = (self.x + self.rect_size, self.y - self.rect_size)
p3 = (self.x - self.rect_size, self.y + self.rect_size)
wh = h if h < w else w
np1 = (w / 2 - wh / 4, h / 2 - wh / 4)
np2 = (w / 2 + wh / 4, h / 2 - wh / 4)
np3 = (w / 2 - wh / 4, h / 2 + wh / 4)
mat = cv2.getAffineTransform(np.float32([p1, p2, p3]) * self.view_scale, np.float32([np1, np2, np3]))
image = cv2.warpAffine(image, mat, (w, h))
view_landmarks = LandmarksProcessor.transform_points(view_landmarks, mat)
landmarks_color = (255, 255, 0) if self.rect_locked else (0, 255, 0)
LandmarksProcessor.draw_rect_landmarks(image, view_rect, view_landmarks, self.face_type, self.image_size, landmarks_color=landmarks_color)
self.extract_needed = False
io.show_image(self.wnd_name, image)
# override
def on_result(self, host_dict, data, result):
if self.type == 'landmarks-manual':
filepath, landmarks = result.filepath, result.landmarks
if len(landmarks) != 0 and landmarks[0] is not None:
self.landmarks = landmarks[0]
self.redraw()
else:
self.result.append(result)
io.progress_bar_inc(1)
# override
def get_result(self):
return self.result
class DeletedFilesSearcherSubprocessor(Subprocessor):
class Cli(Subprocessor.Cli):
# override
def on_initialize(self, client_dict):
self.debug_paths_stems = client_dict['debug_paths_stems']
return None
# override
def process_data(self, data):
input_path_stem = Path(data[0]).stem
return any([input_path_stem == d_stem for d_stem in self.debug_paths_stems])
# override
def get_data_name(self, data):
# return string identificator of your data
return data[0]
# override
def __init__(self, input_paths, debug_paths):
self.input_paths = input_paths
self.debug_paths_stems = [Path(d).stem for d in debug_paths]
self.result = []
super().__init__('DeletedFilesSearcherSubprocessor', DeletedFilesSearcherSubprocessor.Cli, 60)
# override
def process_info_generator(self):
for i in range(min(multiprocessing.cpu_count(), 8)):
yield 'CPU%d' % (i), {}, {'debug_paths_stems': self.debug_paths_stems}
# override
def on_clients_initialized(self):
io.progress_bar("Searching deleted files", len(self.input_paths))
# override
def on_clients_finalized(self):
io.progress_bar_close()
# override
def get_data(self, host_dict):
if len(self.input_paths) > 0:
return [self.input_paths.pop(0)]
return None
# override
def on_data_return(self, host_dict, data):
self.input_paths.insert(0, data[0])
# override
def on_result(self, host_dict, data, result):
if result == False:
self.result.append(data[0])
io.progress_bar_inc(1)
# override
def get_result(self):
return self.result
def main(detector=None,
input_path=None,
output_path=None,
output_debug=None,
manual_fix=False,
manual_output_debug_fix=False,
manual_window_size=1368,
face_type='full_face',
max_faces_from_image=None,
image_size=None,
jpeg_quality=None,
cpu_only=False,
force_gpu_idxs=None,
):
if not input_path.exists():
io.log_err('Input directory not found. Please ensure it exists.')
return
if not output_path.exists():
output_path.mkdir(parents=True, exist_ok=True)
if face_type is not None:
face_type = FaceType.fromString(face_type)
if face_type is None:
if manual_output_debug_fix:
files = pathex.get_image_paths(output_path)
if len(files) != 0:
dflimg = DFLIMG.load(Path(files[0]))
if dflimg is not None and dflimg.has_data():
face_type = FaceType.fromString(dflimg.get_face_type())
input_image_paths = pathex.get_image_unique_filestem_paths(input_path, verbose_print_func=io.log_info)
output_images_paths = pathex.get_image_paths(output_path)
output_debug_path = output_path.parent / (output_path.name + '_debug')
continue_extraction = False
if not manual_output_debug_fix and len(output_images_paths) > 0:
if len(output_images_paths) > 128:
continue_extraction = io.input_bool("Continue extraction?", True, help_message="Extraction can be continued, but you must specify the same options again.")
if len(output_images_paths) > 128 and continue_extraction:
try:
input_image_paths = input_image_paths[[Path(x).stem for x in input_image_paths].index(Path(output_images_paths[-128]).stem.split('_')[0]):]
except:
io.log_err("Error in fetching the last index. Extraction cannot be continued.")
return
elif input_path != output_path:
io.input(f"\n WARNING !!! \n {output_path} contains files! \n They will be deleted. \n Press enter to continue.\n")
for filename in output_images_paths:
Path(filename).unlink()
device_config = nn.DeviceConfig.GPUIndexes(force_gpu_idxs or nn.ask_choose_device_idxs(choose_only_one=detector == 'manual', suggest_all_gpu=True)) \
if not cpu_only else nn.DeviceConfig.CPU()
if face_type is None:
face_type = io.input_str("Face type", 'wf', ['f', 'wf', 'head'],
help_message="Full face / whole face / head. 'Whole face' covers full area of face include forehead. 'head' covers full head, but requires XSeg for src and dst faceset.").lower()
face_type = {'f': FaceType.FULL,
'wf': FaceType.WHOLE_FACE,
'head': FaceType.HEAD}[face_type]
if max_faces_from_image is None:
max_faces_from_image = io.input_int(f"Max number of faces from image", 0,
help_message="If you extract a src faceset that has frames with a large number of faces, it is advisable to set max faces to 3 to speed up extraction. 0 - unlimited")
if image_size is None:
image_size = io.input_int(f"Image size", 512 if face_type < FaceType.HEAD else 768, valid_range=[256, 2048],
help_message="Output image size. The higher image size, the worse face-enhancer works. Use higher than 512 value only if the source image is sharp enough and the face does not need to be enhanced.")
if jpeg_quality is None:
jpeg_quality = io.input_int(f"Jpeg quality", 90, valid_range=[1, 100], help_message="Jpeg quality. The higher jpeg quality the larger the output file size.")
if detector is None:
io.log_info("Choose detector type.")
io.log_info("[0] S3FD")
io.log_info("[1] manual")
detector = {0: 's3fd', 1: 'manual'}[io.input_int("", 0, [0, 1])]
if output_debug is None:
output_debug = io.input_bool(f"Write debug images to {output_debug_path.name}?", False)
if output_debug:
output_debug_path.mkdir(parents=True, exist_ok=True)
if manual_output_debug_fix:
if not output_debug_path.exists():
io.log_err(f'{output_debug_path} not found. Re-extract faces with "Write debug images" option.')
return
else:
detector = 'manual'
io.log_info('Performing re-extract frames which were deleted from _debug directory.')
input_image_paths = DeletedFilesSearcherSubprocessor(input_image_paths, pathex.get_image_paths(output_debug_path)).run()
input_image_paths = sorted(input_image_paths)
io.log_info('Found %d images.' % (len(input_image_paths)))
else:
if not continue_extraction and output_debug_path.exists():
for filename in pathex.get_image_paths(output_debug_path):
Path(filename).unlink()
images_found = len(input_image_paths)
faces_detected = 0
if images_found != 0:
if detector == 'manual':
io.log_info('Performing manual extract...')
data = ExtractSubprocessor([ExtractSubprocessor.Data(Path(filename)) for filename in input_image_paths], 'landmarks-manual', image_size, jpeg_quality,
face_type, output_debug_path if output_debug else None, manual_window_size=manual_window_size, device_config=device_config).run()
io.log_info('Performing 3rd pass...')
data = ExtractSubprocessor(data, 'final', image_size, jpeg_quality, face_type, output_debug_path if output_debug else None, final_output_path=output_path,
device_config=device_config).run()
else:
io.log_info('Extracting faces...')
data = ExtractSubprocessor([ExtractSubprocessor.Data(Path(filename)) for filename in input_image_paths],
'all',
image_size,
jpeg_quality,
face_type,
output_debug_path if output_debug else None,
max_faces_from_image=max_faces_from_image,
final_output_path=output_path,
device_config=device_config).run()
faces_detected += sum([d.faces_detected for d in data])
if manual_fix:
if all(np.array([d.faces_detected > 0 for d in data]) == True):
io.log_info('All faces are detected, manual fix not needed.')
else:
fix_data = [ExtractSubprocessor.Data(d.filepath) for d in data if d.faces_detected == 0]
io.log_info('Performing manual fix for %d images...' % (len(fix_data)))
fix_data = ExtractSubprocessor(fix_data, 'landmarks-manual', image_size, jpeg_quality, face_type, output_debug_path if output_debug else None,
manual_window_size=manual_window_size, device_config=device_config).run()
fix_data = ExtractSubprocessor(fix_data, 'final', image_size, jpeg_quality, face_type, output_debug_path if output_debug else None,
final_output_path=output_path, device_config=device_config).run()
faces_detected += sum([d.faces_detected for d in fix_data])
io.log_info('-------------------------')
io.log_info('Images found: %d' % (images_found))
io.log_info('Faces detected: %d' % (faces_detected))
io.log_info('-------------------------')

复制代码

day270010678 · 发表于 5 天前

优化了gpu并行能力，可以同时处理的图片不是一张，而是批次的多图，还有优化了些计算，等等，忘记了都写了什么了。速度随着显卡的大小和gpu数，呈几何倍数递增。

请大家看清楚，是优化了代码，不是他们那些所谓的靠丢失脸部的提速。

至于具体怎么样，看你们自己的硬件了。

想起来实时监控gpu去动态加载和调整批处理

fghfdg · 发表于 5 天前

day270010678 发表于 2026-1-7 20:11
优化了gpu并行能力，可以同时处理的图片不是一张，而是批次的多图，还有优化了些计算，等等，忘记了都写了 ...

大佬请问提脸主要是这三个py文件吗？如果我想修改提脸的过滤逻辑（我记得原版的过滤好像是过滤10像素一下的人脸），如果我想把提脸的过滤改成用户交互的动态选择。是不是也是修改这三个py文件就行了

day270010678 · 发表于 5 天前

本帖最后由 day270010678 于 2026-1-7 21:00 编辑

fghfdg 发表于 2026-1-7 20:53
大佬请问提脸主要是这三个py文件吗？如果我想修改提脸的过滤逻辑（我记得原版的过滤好像是过滤10像素一下 ...

肯定是这里修改，不是你写交互干嘛？自己设定就行了。你过虑什么？是像他们一样靠丢失提速？在S3FDExtractor.py里面不是有个阀值吗？默认0.5，你随便提高，你切脸速度马上就提升了，只是丢失严重罢了。

DFL小白02 · 发表于 5 天前

你好，我看你早期还发过贴讨论src切脸俯仰角度五官对称问题，现在你有答案了吗？按对称标会出问题吗

day270010678 · 发表于 5 天前

DFL小白02 发表于 2026-1-7 21:05
你好，我看你早期还发过贴讨论src切脸俯仰角度五官对称问题，现在你有答案了吗？按对称标会出问题吗 ...

那个时候刚开始接触，没研究过他代码，现在就这种对称的问题，解决法子多了去了，比方说降低阀值，最佳法子肯定是修改旋转逻辑

day270010678 · 发表于 5 天前

在模型推理内部实现真正的多线程并行（修改模型架构以支持动态批次大小，利用TensorFlow/CUDA的内置并行机制等等），每个子进程可以一次处理多个图像。有人研究过这个方向不？我试验了下TensorFlow的内置并行机制，效果老是不理想，感觉是鸡肋。

长空雷隐 · 发表于 5 天前

楼主能放出文件直接下载吗？小白不会编辑

pasanonic · 发表于 4 天前

A卡有效果吗

		自动登录	找回密码
密码			立即注册（仅限QQ邮箱）

关于提取脸部速度问题的解决

评分

万事如意节日勋章

开心娱乐节日勋章