81 lines
2.5 KiB
Python
81 lines
2.5 KiB
Python
"""
|
|
图像处理工具模块
|
|
提供图像tensor转换、哈希计算等通用图像处理方法
|
|
"""
|
|
|
|
import base64
|
|
import hashlib
|
|
from io import BytesIO
|
|
from typing import Optional
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
|
|
|
|
def tensor_to_base64(image_tensor: torch.Tensor, quality: int = 95) -> str:
|
|
"""
|
|
将图像tensor转换为base64编码的data URL
|
|
|
|
参数:
|
|
image_tensor: 图像tensor,形状为 [H, W, C] 或 [B, H, W, C]
|
|
quality: JPEG压缩质量 (1-100),默认95
|
|
|
|
返回:
|
|
base64编码的data URL,格式: "data:image/jpeg;base64,..."
|
|
"""
|
|
# 如果是4D tensor,取第一张图片
|
|
if len(image_tensor.shape) == 4:
|
|
image_tensor = image_tensor[0]
|
|
|
|
# 转换为numpy数组并缩放到0-255范围
|
|
image_np = (image_tensor.cpu().numpy() * 255).astype(np.uint8)
|
|
|
|
# 创建PIL图像
|
|
image = Image.fromarray(image_np)
|
|
|
|
# 转换为JPEG格式的字节流
|
|
buffer = BytesIO()
|
|
image.save(buffer, format="JPEG", quality=quality)
|
|
|
|
# 转换为base64编码
|
|
encoded_image = base64.b64encode(buffer.getvalue()).decode('utf-8')
|
|
|
|
# 返回带有MIME类型的data URL
|
|
return f"data:image/jpeg;base64,{encoded_image}"
|
|
|
|
|
|
def compute_image_hash(image_tensor: Optional[torch.Tensor]) -> str:
|
|
"""
|
|
计算图像tensor的哈希值,用于 IS_CHANGED 检测
|
|
|
|
算法: 只使用第一帧的中心区域作为哈希计算依据,避免计算量过大
|
|
|
|
参数:
|
|
image_tensor: 图像tensor,形状为 [B, H, W, C] 或 None
|
|
|
|
返回:
|
|
MD5哈希值的十六进制字符串,如果输入为None或计算失败则返回 "0"
|
|
"""
|
|
if image_tensor is None:
|
|
return "0"
|
|
|
|
try:
|
|
if len(image_tensor.shape) == 4:
|
|
# 取第一帧的中心区域作为哈希计算依据
|
|
h, w = image_tensor.shape[1:3]
|
|
center_h, center_w = h // 2, w // 2
|
|
size = min(100, h // 4, w // 4) # 限制计算区域大小
|
|
|
|
img_data = image_tensor[0,
|
|
max(0, center_h - size):min(h, center_h + size),
|
|
max(0, center_w - size):min(w, center_w + size),
|
|
0].cpu().numpy().tobytes()
|
|
return hashlib.md5(img_data).hexdigest()
|
|
else:
|
|
# 如果不是4D tensor,使用整个tensor的哈希
|
|
img_data = image_tensor.cpu().numpy().tobytes()
|
|
return hashlib.md5(img_data).hexdigest()
|
|
except Exception:
|
|
return "0"
|