343 lines
13 KiB
Python
343 lines
13 KiB
Python
import os
|
||
import re
|
||
import shutil
|
||
from datetime import datetime
|
||
from PIL import ImageDraw, ImageFont
|
||
from fastapi import UploadFile
|
||
import cv2
|
||
from PIL import Image
|
||
import numpy as np
|
||
|
||
# 上传根目录
|
||
UPLOAD_ROOT = "upload"
|
||
PRE = "/api/file/download/"
|
||
|
||
# 确保上传根目录存在
|
||
os.makedirs(UPLOAD_ROOT, exist_ok=True)
|
||
|
||
|
||
|
||
def save_detect_file(client_ip: str, image_np: np.ndarray, file_type: str) -> str:
|
||
"""保存numpy数组格式的PNG图片到detect目录，返回下载路径"""
|
||
today = datetime.now()
|
||
year = today.strftime("%Y")
|
||
month = today.strftime("%m")
|
||
day = today.strftime("%d")
|
||
|
||
# 构建目录路径: upload/detect/客户端IP/type/年/月/日（包含UPLOAD_ROOT）
|
||
file_dir = os.path.join(
|
||
UPLOAD_ROOT,
|
||
"detect",
|
||
client_ip,
|
||
file_type,
|
||
year,
|
||
month,
|
||
day
|
||
)
|
||
|
||
# 创建目录（确保目录存在）
|
||
os.makedirs(file_dir, exist_ok=True)
|
||
|
||
# 生成当前时间戳作为文件名，确保唯一性
|
||
timestamp = datetime.now().strftime("%Y%m%d%H%M%S%f")
|
||
filename = f"{timestamp}.png"
|
||
|
||
# 1. 完整路径：用于实际保存文件（包含UPLOAD_ROOT）
|
||
full_path = os.path.join(file_dir, filename)
|
||
# 2. 相对路径：用于返回给前端（移除UPLOAD_ROOT前缀）
|
||
relative_path = full_path.replace(UPLOAD_ROOT, "", 1).lstrip(os.sep)
|
||
|
||
# 保存numpy数组为PNG图片
|
||
try:
|
||
# -------- 新增/修改：处理颜色通道和数据类型 --------
|
||
# 1. 数据类型转换：确保是uint8（若为float32且范围0-1，需转成0-255的uint8）
|
||
if image_np.dtype != np.uint8:
|
||
image_np = (image_np * 255).astype(np.uint8)
|
||
|
||
# 2. 通道顺序转换：若为OpenCV的BGR格式，转成PIL需要的RGB格式
|
||
image_rgb = cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB)
|
||
|
||
# 3. 转换为PIL Image并保存
|
||
img = Image.fromarray(image_rgb)
|
||
img.save(full_path, format='PNG')
|
||
except Exception as e:
|
||
# 处理可能的异常（如数组格式不正确）
|
||
raise Exception(f"保存图片失败: {str(e)}")
|
||
|
||
# 统一路径分隔符为/，拼接前缀返回
|
||
return PRE + relative_path.replace(os.sep, "/")
|
||
|
||
|
||
def save_detect_yolo_file(
|
||
client_ip: str,
|
||
image_np: np.ndarray,
|
||
detection_results: list,
|
||
file_type: str = "yolo"
|
||
) -> str:
|
||
|
||
|
||
print("......................")
|
||
"""
|
||
保存YOLO检测结果图片（在原图上绘制边界框+标签），返回前端可访问的下载路径
|
||
"""
|
||
# 输入参数验证
|
||
if not isinstance(image_np, np.ndarray):
|
||
raise ValueError(f"输入image_np必须是numpy数组，当前类型：{type(image_np)}")
|
||
if image_np.ndim != 3 or image_np.shape[-1] != 3:
|
||
raise ValueError(f"输入图像必须是 (h, w, 3) 的BGR数组，当前shape：{image_np.shape}")
|
||
|
||
if not isinstance(detection_results, list):
|
||
raise ValueError(f"detection_results必须是列表，当前类型：{type(detection_results)}")
|
||
for idx, result in enumerate(detection_results):
|
||
required_keys = {"class", "confidence", "bbox"}
|
||
if not isinstance(result, dict) or not required_keys.issubset(result.keys()):
|
||
raise ValueError(
|
||
f"detection_results第{idx}个元素格式错误，需包含键：{required_keys}，"
|
||
f"当前元素：{result}"
|
||
)
|
||
bbox = result["bbox"]
|
||
if not (isinstance(bbox, (tuple, list)) and len(bbox) == 4 and all(isinstance(x, int) for x in bbox)):
|
||
raise ValueError(
|
||
f"detection_results第{idx}个元素的bbox格式错误，需为(x1,y1,x2,y2)整数元组，"
|
||
f"当前bbox：{bbox}"
|
||
)
|
||
|
||
#图像预处理（数据类型+通道）
|
||
draw_image = image_np.copy()
|
||
if draw_image.dtype != np.uint8:
|
||
draw_image = np.clip(draw_image * 255, 0, 255).astype(np.uint8)
|
||
|
||
#绘制边界框+标签
|
||
# 遍历所有检测结果，逐个绘制
|
||
for result in detection_results:
|
||
class_name = result["class"]
|
||
confidence = result["confidence"]
|
||
x1, y1, x2, y2 = result["bbox"]
|
||
cv2.rectangle(draw_image, (x1, y1), (x2, y2), color=(0, 255, 0), thickness=2)
|
||
label = f"{class_name}: {confidence:.2f}"
|
||
font = cv2.FONT_HERSHEY_SIMPLEX
|
||
font_scale = 0.5
|
||
font_thickness = 2
|
||
(label_width, label_height), baseline = cv2.getTextSize(
|
||
label, font, font_scale, font_thickness
|
||
)
|
||
|
||
bg_top_left = (x1, y1 - label_height - 10)
|
||
bg_bottom_right = (x1 + label_width, y1)
|
||
if bg_top_left[1] < 0:
|
||
bg_top_left = (x1, 0)
|
||
bg_bottom_right = (x1 + label_width, label_height + 10)
|
||
cv2.rectangle(draw_image, bg_top_left, bg_bottom_right, color=(0, 0, 0), thickness=-1)
|
||
|
||
text_origin = (x1, y1 - 5)
|
||
if bg_top_left[1] == 0:
|
||
text_origin = (x1, label_height + 5)
|
||
cv2.putText(
|
||
draw_image, label, text_origin,
|
||
font, font_scale, color=(255, 255, 255), thickness=font_thickness
|
||
)
|
||
|
||
#保存图片
|
||
try:
|
||
today = datetime.now()
|
||
year = today.strftime("%Y")
|
||
month = today.strftime("%m")
|
||
day = today.strftime("%d")
|
||
file_dir = os.path.join(
|
||
UPLOAD_ROOT, "detect", client_ip, file_type, year, month, day
|
||
)
|
||
|
||
#创建目录（若不存在则创建，支持多级目录）
|
||
os.makedirs(file_dir, exist_ok=True)
|
||
|
||
#生成唯一文件名
|
||
timestamp = today.strftime("%Y%m%d%H%M%S%f")
|
||
filename = f"{timestamp}.png"
|
||
|
||
# 4.4 构建完整保存路径和前端访问路径
|
||
full_path = os.path.join(file_dir, filename) # 本地完整路径
|
||
# 相对路径：移除UPLOAD_ROOT前缀，统一用"/"作为分隔符（兼容Windows/Linux）
|
||
relative_path = full_path.replace(UPLOAD_ROOT, "", 1).lstrip(os.sep)
|
||
download_path = PRE + relative_path.replace(os.sep, "/")
|
||
|
||
# 4.5 保存图片（CV2绘制的是BGR，需转RGB后用PIL保存，与原逻辑一致）
|
||
image_rgb = cv2.cvtColor(draw_image, cv2.COLOR_BGR2RGB)
|
||
img_pil = Image.fromarray(image_rgb)
|
||
img_pil.save(full_path, format="PNG", quality=95) # PNG格式无压缩，quality可忽略
|
||
|
||
print(f"YOLO检测图片保存成功 | 本地路径：{full_path} | 下载路径：{download_path}")
|
||
return download_path
|
||
|
||
except Exception as e:
|
||
raise Exception(f"YOLO检测图片保存失败：{str(e)}") from e
|
||
|
||
|
||
def save_detect_face_file(
|
||
client_ip: str,
|
||
image_np: np.ndarray,
|
||
face_result: str,
|
||
file_type: str = "face",
|
||
matched_color: tuple = (0, 255, 0)
|
||
) -> str:
|
||
"""
|
||
保存人脸识别结果图片（仅为「匹配成功」的人脸画框，标签不包含“匹配”二字）
|
||
"""
|
||
#输入参数验证
|
||
if not isinstance(image_np, np.ndarray) or image_np.ndim != 3 or image_np.shape[-1] != 3:
|
||
raise ValueError(f"输入图像需为 (h, w, 3) 的BGR数组，当前shape：{image_np.shape}")
|
||
if not isinstance(face_result, str) or face_result.strip() == "":
|
||
raise ValueError("face_result必须是非空字符串")
|
||
|
||
# 解析face_result提取人脸信息
|
||
face_info_list = []
|
||
if face_result.strip() != "未检测到人脸":
|
||
face_pattern = re.compile(
|
||
r"(匹配|未匹配):\s*([^\s(]+)\s*\(相似度:\s*(\d+\.\d+),\s*边界框:\s*\[(\d+,\s*\d+,\s*\d+,\s*\d+)\]\)"
|
||
)
|
||
for part in [p.strip() for p in face_result.split(";") if p.strip()]:
|
||
match = face_pattern.match(part)
|
||
if match:
|
||
status, name, similarity, bbox_str = match.groups()
|
||
bbox = list(map(int, bbox_str.replace(" ", "").split(",")))
|
||
if len(bbox) == 4:
|
||
face_info_list.append({
|
||
"status": status,
|
||
"name": name,
|
||
"similarity": float(similarity),
|
||
"bbox": bbox
|
||
})
|
||
|
||
# 图像格式转换（OpenCV→PIL）
|
||
image_rgb = cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB)
|
||
pil_img = Image.fromarray(image_rgb)
|
||
draw = ImageDraw.Draw(pil_img)
|
||
|
||
# 绘制边界框和标签
|
||
font_size = 12
|
||
try:
|
||
font = ImageFont.truetype("simhei", font_size)
|
||
except:
|
||
try:
|
||
font = ImageFont.truetype("simsun", font_size)
|
||
except:
|
||
font = ImageFont.load_default()
|
||
print("警告：未找到指定中文字体，使用PIL默认字体（可能影响中文显示）")
|
||
|
||
for face_info in face_info_list:
|
||
status = face_info["status"]
|
||
if status != "匹配":
|
||
print(f"跳过未匹配人脸：{face_info['name']}（相似度：{face_info['similarity']:.2f}）")
|
||
continue
|
||
|
||
name = face_info["name"]
|
||
similarity = face_info["similarity"]
|
||
x1, y1, x2, y2 = face_info["bbox"]
|
||
|
||
# 4.1 绘制边界框（绿色）
|
||
img_cv = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
|
||
cv2.rectangle(img_cv, (x1, y1), (x2, y2), color=matched_color, thickness=2)
|
||
pil_img = Image.fromarray(cv2.cvtColor(img_cv, cv2.COLOR_BGR2RGB))
|
||
draw = ImageDraw.Draw(pil_img)
|
||
|
||
label = f"{name} (相似度: {similarity:.2f})"
|
||
|
||
# 4.3 计算标签尺寸（文本变短后会自动适配，无需额外调整）
|
||
label_bbox = draw.textbbox((0, 0), label, font=font)
|
||
label_width = label_bbox[2] - label_bbox[0]
|
||
label_height = label_bbox[3] - label_bbox[1]
|
||
|
||
# 4.4 计算标签背景位置（避免超出图像）
|
||
bg_x1, bg_y1 = x1, y1 - label_height - 10
|
||
bg_x2, bg_y2 = x1 + label_width, y1
|
||
if bg_y1 < 0:
|
||
bg_y1, bg_y2 = y2 + 5, y2 + label_height + 15
|
||
|
||
# 4.5 绘制标签背景（黑色）和文本（白色）
|
||
draw.rectangle([(bg_x1, bg_y1), (bg_x2, bg_y2)], fill=(0, 0, 0))
|
||
text_x = bg_x1
|
||
text_y = bg_y1 if bg_y1 >= 0 else bg_y1 + label_height
|
||
draw.text((text_x, text_y), label, font=font, fill=(255, 255, 255))
|
||
|
||
#保存图片
|
||
try:
|
||
today = datetime.now()
|
||
file_dir = os.path.join(
|
||
UPLOAD_ROOT, "detect", client_ip, file_type,
|
||
today.strftime("%Y"), today.strftime("%m"), today.strftime("%d")
|
||
)
|
||
os.makedirs(file_dir, exist_ok=True)
|
||
|
||
timestamp = today.strftime("%Y%m%d%H%M%S%f")
|
||
filename = f"{timestamp}.png"
|
||
full_path = os.path.join(file_dir, filename)
|
||
|
||
pil_img.save(full_path, format="PNG", quality=100)
|
||
|
||
relative_path = full_path.replace(UPLOAD_ROOT, "", 1).lstrip(os.sep)
|
||
download_path = PRE + relative_path.replace(os.sep, "/")
|
||
|
||
matched_count = sum(1 for info in face_info_list if info["status"] == "匹配")
|
||
print(f"人脸检测图片保存成功 | 客户端IP：{client_ip} | 匹配人脸数：{matched_count} | 保存路径：{download_path}")
|
||
return download_path
|
||
|
||
except Exception as e:
|
||
raise Exception(f"人脸检测图片保存失败（客户端IP：{client_ip}）：{str(e)}") from e
|
||
|
||
def save_source_file(upload_file: UploadFile, file_type: str) -> str:
|
||
"""保存上传的文件到source目录，返回下载路径"""
|
||
today = datetime.now()
|
||
year = today.strftime("%Y")
|
||
month = today.strftime("%m")
|
||
day = today.strftime("%d")
|
||
|
||
# 生成精确到微秒的时间戳，确保文件名唯一
|
||
timestamp = today.strftime("%Y%m%d%H%M%S%f")
|
||
# 构建新文件名：时间戳_原文件名
|
||
unique_filename = f"{timestamp}_{upload_file.filename}"
|
||
|
||
# 构建目录路径: upload/source/type/年/月/日（包含UPLOAD_ROOT）
|
||
file_dir = os.path.join(
|
||
UPLOAD_ROOT,
|
||
"source",
|
||
file_type,
|
||
year,
|
||
month,
|
||
day
|
||
)
|
||
|
||
# 创建目录（确保目录存在）
|
||
os.makedirs(file_dir, exist_ok=True)
|
||
|
||
# 1. 完整路径：用于实际保存文件（使用带时间戳的唯一文件名）
|
||
full_path = os.path.join(file_dir, unique_filename)
|
||
# 2. 相对路径：用于返回给前端
|
||
relative_path = full_path.replace(UPLOAD_ROOT, "", 1).lstrip(os.sep)
|
||
|
||
# 保存文件（使用完整路径）
|
||
try:
|
||
with open(full_path, "wb") as buffer:
|
||
shutil.copyfileobj(upload_file.file, buffer)
|
||
finally:
|
||
upload_file.file.close()
|
||
|
||
# 统一路径分隔符为/
|
||
return PRE + relative_path.replace(os.sep, "/")
|
||
|
||
|
||
def get_absolute_path(relative_path: str) -> str:
|
||
"""
|
||
根据相对路径获取服务器上的绝对路径
|
||
"""
|
||
path_without_pre = relative_path.replace(PRE, "", 1)
|
||
|
||
# 将相对路径转换为系统兼容的格式
|
||
normalized_path = os.path.normpath(path_without_pre)
|
||
|
||
# 拼接基础路径和相对路径，得到绝对路径
|
||
absolute_path = os.path.abspath(os.path.join(UPLOAD_ROOT, normalized_path))
|
||
|
||
# 安全检查：确保生成的路径在UPLOAD_ROOT目录下，防止路径遍历
|
||
if not absolute_path.startswith(os.path.abspath(UPLOAD_ROOT)):
|
||
raise ValueError("无效的相对路径，可能试图访问上传目录之外的内容")
|
||
|
||
return absolute_path |