524 lines
25 KiB
Python
524 lines
25 KiB
Python
# pattern_to_tshirt.py
|
||
import sys
|
||
import os
|
||
import json
|
||
import io
|
||
import cv2
|
||
import numpy as np
|
||
from PIL import Image, ImageFilter, ImageDraw, ImageChops
|
||
import uuid
|
||
import urllib.request
|
||
import urllib3
|
||
import requests
|
||
from pydantic import BaseModel
|
||
from typing import Optional
|
||
import base64
|
||
import asyncio
|
||
import warnings
|
||
import tempfile
|
||
from urllib.parse import urlparse
|
||
import torch
|
||
import time
|
||
import gc
|
||
|
||
# 关闭不必要的警告
|
||
warnings.filterwarnings("ignore", category=UserWarning)
|
||
warnings.filterwarnings("ignore", category=FutureWarning)
|
||
|
||
class PtnToTshirtService:
|
||
# 默认配置
|
||
DEFAULT_CONFIG = {
|
||
'background_removed_marker': "_rmbg",
|
||
'ptt_exclude_markers': ["_upscaled", "_vector", "_processed", "_tshirt", "_tryon"],
|
||
'tshirt_marker': "_tshirt",
|
||
'processed_pattern_marker': "_processed",
|
||
'tshirt_image_path': 'home/tshirt',
|
||
'tshirt_urls': [], # 新增: T恤图片URL列表
|
||
'alpha': 1, # 透明度，0表示全透明，1表示全不透明
|
||
'ptt_design_size_ratio': 0.4, # 设计图像占T恤图像的比例
|
||
'ptt_design_offset': [0.5, 0.45], # 设计图像在T恤图像中的相对位置 [x, y]
|
||
'ptt_design_rotation': 0, # 设计图案旋转角度
|
||
|
||
'enable_gradient_effect': True, # 是否启用渐变效果
|
||
'gradient_width': 512, # 渐变宽度
|
||
'gradient_direction': 'outward', # 渐变方向: 'outward', 'inward'
|
||
'gradient_type': 'linear', # 渐变类型: 'linear', 'radial'
|
||
'gradient_max_alpha': 150, # 渐变的最大透明度值，0-255
|
||
'gradient_start_alpha': 0, # 渐变起始处的透明度，0-255
|
||
'gradient_color': [255, 255, 255, 255], # 渐变颜色
|
||
'gradient_blur_intensity': 10, # 渐变模糊强度
|
||
'gradient_center': [0.5, 0.5], # 渐变中心位置，相对于设计图案的 [x, y]
|
||
'gradient_repeat_count': 1, # 渐变重复次数
|
||
|
||
'ptt_enable_texture_effect': False, # 是否启用纹理效果
|
||
'ptt_texture_type': 'lines', # 纹理类型: 'noise', 'lines'
|
||
'ptt_texture_blend_mode': 'multiply', # 纹理混合模式
|
||
|
||
'enable_save_processed_design': True, # 是否单独保存处理后的设计图案
|
||
'ptt_design_output_format': 'png', # 设计图案保存格式: 'png' 或 'tiff'
|
||
|
||
'ptt_enable_color_matching': True, # 是否启用颜色匹配
|
||
'ptt_enable_lighting_effect': False, # 是否启用光效
|
||
'ptt_enable_monochrome': False,
|
||
'ptt_light_intensity': 0.5, # 光照强度
|
||
'ptt_light_position': [0.5, 0.3], # 光源位置 [相对位置 x, y]
|
||
'ptt_light_radius_ratio': [0.4, 0.25], # 光源半径相对比例 [宽, 高]
|
||
'ptt_light_angle': 45, # 光源角度（度）
|
||
'ptt_light_blur': 91, # 光源模糊程度
|
||
'ptt_light_shape': 'ellipse', # 光源形状: 'ellipse', 'circle', 'rect'
|
||
}
|
||
|
||
def __init__(self):
|
||
"""初始化图案到T恤服务"""
|
||
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
print(f"使用设备: {self.device}")
|
||
|
||
if torch.cuda.is_available():
|
||
print(f"GPU: {torch.cuda.get_device_name(0)}")
|
||
print(f"显存总量: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f}GB")
|
||
|
||
def overlay_image_alpha(self, img, img_overlay, pos, alpha_mask):
|
||
"""在图像上叠加另一个具有透明度的图像"""
|
||
x, y = pos
|
||
y1, y2 = max(0, y), min(img.shape[0], y + img_overlay.shape[0])
|
||
x1, x2 = max(0, x), min(img.shape[1], x + img_overlay.shape[1])
|
||
y1o, y2o = max(0, -y), min(img_overlay.shape[0], img.shape[0] - y)
|
||
x1o, x2o = max(0, -x), min(img_overlay.shape[1], img.shape[1] - x)
|
||
|
||
if y1 >= y2 or x1 >= x2 or y1o >= y2o or x1o >= x2o:
|
||
return
|
||
|
||
img_crop = img[y1:y2, x1:x2]
|
||
img_overlay_crop = img_overlay[y1o:y2o, x1o:x2o]
|
||
alpha = alpha_mask[y1o:y2o, x1o:x2o, np.newaxis]
|
||
img_crop[:] = alpha * img_overlay_crop + (1 - alpha) * img_crop
|
||
|
||
def color_transfer(self, source, target):
|
||
"""颜色匹配：将源图像的颜色转换为目标图像的颜色风格"""
|
||
source = cv2.cvtColor(source, cv2.COLOR_BGR2LAB)
|
||
target = cv2.cvtColor(target, cv2.COLOR_BGR2LAB)
|
||
|
||
src_mean, src_std = cv2.meanStdDev(source)
|
||
tgt_mean, tgt_std = cv2.meanStdDev(target)
|
||
|
||
src_mean = src_mean.reshape(1, 1, 3)
|
||
src_std = src_std.reshape(1, 1, 3)
|
||
tgt_mean = tgt_mean.reshape(1, 1, 3)
|
||
tgt_std = tgt_std.reshape(1, 1, 3)
|
||
|
||
result = (source - src_mean) * (tgt_std / src_std) + tgt_mean
|
||
result = np.clip(result, 0, 255)
|
||
result = result.astype(np.uint8)
|
||
|
||
return cv2.cvtColor(result, cv2.COLOR_LAB2BGR)
|
||
|
||
def apply_lighting_effect(self, image, light_intensity=0.5, light_position=[0.5, 0.3],
|
||
light_radius_ratio=[0.4, 0.25], light_angle=45, light_blur=91, light_shape='ellipse'):
|
||
"""应用光照效果到图像"""
|
||
height, width = image.shape[:2]
|
||
light_position = (int(light_position[0] * width), int(light_position[1] * height))
|
||
light_radius = (int(light_radius_ratio[0] * width), int(light_radius_ratio[1] * height))
|
||
mask = np.zeros((height, width), dtype=np.uint8)
|
||
|
||
if light_shape == 'ellipse':
|
||
cv2.ellipse(mask, light_position, light_radius, light_angle, 0, 360, 255, -1)
|
||
elif light_shape == 'circle':
|
||
cv2.circle(mask, light_position, min(light_radius), 255, -1)
|
||
elif light_shape == 'rect':
|
||
rect_top_left = (light_position[0] - light_radius[0] // 2, light_position[1] - light_radius[1] // 2)
|
||
rect_bottom_right = (light_position[0] + light_radius[0] // 2, light_position[1] + light_radius[1] // 2)
|
||
cv2.rectangle(mask, rect_top_left, rect_bottom_right, 255, -1)
|
||
|
||
mask = cv2.GaussianBlur(mask, (light_blur, light_blur), 0)
|
||
mask = mask.astype(np.float32) / 255
|
||
result = image.astype(np.float32)
|
||
for i in range(3):
|
||
result[:, :, i] = result[:, :, i] * (1 - light_intensity + mask * light_intensity)
|
||
return result.astype(np.uint8)
|
||
|
||
def apply_monochrome(self, image):
|
||
"""将图像转换为单色"""
|
||
gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
||
monochrome_image = cv2.cvtColor(gray_image, cv2.COLOR_GRAY2BGR)
|
||
return monochrome_image
|
||
|
||
def enhance_design(self, design_image, tshirt_image, config):
|
||
"""增强设计图像，应用各种效果"""
|
||
if config.get('ptt_enable_color_matching', self.DEFAULT_CONFIG['ptt_enable_color_matching']):
|
||
design_image = self.color_transfer(design_image, tshirt_image)
|
||
|
||
if config.get('ptt_enable_lighting_effect', self.DEFAULT_CONFIG['ptt_enable_lighting_effect']):
|
||
design_image = self.apply_lighting_effect(
|
||
design_image,
|
||
config.get('ptt_light_intensity', self.DEFAULT_CONFIG['ptt_light_intensity']),
|
||
config.get('ptt_light_position', self.DEFAULT_CONFIG['ptt_light_position']),
|
||
config.get('ptt_light_radius_ratio', self.DEFAULT_CONFIG['ptt_light_radius_ratio']),
|
||
config.get('ptt_light_angle', self.DEFAULT_CONFIG['ptt_light_angle']),
|
||
config.get('ptt_light_blur', self.DEFAULT_CONFIG['ptt_light_blur']),
|
||
config.get('ptt_light_shape', self.DEFAULT_CONFIG['ptt_light_shape'])
|
||
)
|
||
|
||
if config.get('ptt_enable_monochrome', self.DEFAULT_CONFIG['ptt_enable_monochrome']):
|
||
design_image = self.apply_monochrome(design_image)
|
||
|
||
return design_image
|
||
|
||
def add_edge_gradient(self, image, gradient_width, gradient_direction, gradient_type,
|
||
gradient_max_alpha, gradient_start_alpha, gradient_color,
|
||
gradient_blur_intensity, gradient_center):
|
||
"""添加边缘渐变效果"""
|
||
alpha = image.getchannel('A')
|
||
width, height = alpha.size
|
||
mask = Image.new('L', (width, height), 0)
|
||
draw = ImageDraw.Draw(mask)
|
||
|
||
# 确保渐变宽度不超过图像尺寸的一半
|
||
gradient_width = min(gradient_width, width // 2, height // 2)
|
||
|
||
if gradient_type == 'linear':
|
||
if gradient_direction == 'outward':
|
||
for i in range(gradient_width):
|
||
if i >= width // 2 or i >= height // 2:
|
||
break # 避免无效的矩形坐标
|
||
fill_value = int(gradient_start_alpha + (gradient_max_alpha - gradient_start_alpha) * (i / gradient_width))
|
||
draw.rectangle([i, i, width - i - 1, height - i - 1], fill=fill_value)
|
||
elif gradient_direction == 'inward':
|
||
for i in range(gradient_width):
|
||
if i >= width // 2 or i >= height // 2:
|
||
break # 避免无效的矩形坐标
|
||
fill_value = int(gradient_start_alpha + (gradient_max_alpha - gradient_start_alpha) * ((gradient_width - i) / gradient_width))
|
||
draw.rectangle([i, i, width - i - 1, height - i - 1], fill=fill_value)
|
||
elif gradient_type == 'radial':
|
||
center_x = int(width * gradient_center[0])
|
||
center_y = int(height * gradient_center[1])
|
||
max_radius = min(center_x, center_y, width - center_x, height - center_y)
|
||
for i in range(gradient_width):
|
||
radius = max_radius * (i / gradient_width)
|
||
if radius <= 0:
|
||
continue
|
||
fill_value = int(gradient_start_alpha + (gradient_max_alpha - gradient_start_alpha) * (i / gradient_width))
|
||
draw.ellipse([center_x - radius, center_y - radius, center_x + radius, center_y + radius], fill=fill_value)
|
||
|
||
if gradient_blur_intensity < 1:
|
||
gradient_blur_intensity = 1
|
||
mask = mask.filter(ImageFilter.GaussianBlur(gradient_blur_intensity))
|
||
alpha = ImageChops.multiply(alpha, mask)
|
||
image.putalpha(alpha)
|
||
|
||
if gradient_color != [255, 255, 255, 255]:
|
||
colored_mask = Image.new('RGBA', image.size, tuple(gradient_color))
|
||
colored_mask.putalpha(mask)
|
||
image = Image.alpha_composite(image, colored_mask)
|
||
|
||
return image
|
||
|
||
def add_gradient_repeat(self, image, gradient_repeat_count, *args, **kwargs):
|
||
"""重复应用渐变效果"""
|
||
for _ in range(max(gradient_repeat_count, 1)): # 确保至少执行一次
|
||
image = self.add_edge_gradient(image, *args, **kwargs)
|
||
return image
|
||
|
||
def generate_noise_texture(self, size, intensity=64):
|
||
"""生成噪点纹理"""
|
||
noise = np.random.randint(0, intensity, (size, size, 4), dtype=np.uint8)
|
||
noise[..., 3] = 255 # 设置 alpha 通道为不透明
|
||
return Image.fromarray(noise)
|
||
|
||
def generate_line_texture(self, size, line_width=4, spacing=20, color=(0, 0, 0, 255)):
|
||
"""生成线条纹理"""
|
||
texture = Image.new('RGBA', (size, size), (255, 255, 255, 0))
|
||
draw = ImageDraw.Draw(texture)
|
||
for y in range(0, size, spacing):
|
||
draw.line([(0, y), (size, y)], fill=color, width=line_width)
|
||
for x in range(0, size, spacing):
|
||
draw.line([(x, 0), (x, size)], fill=color, width=line_width)
|
||
return texture
|
||
|
||
def add_texture(self, image, texture_type, texture_blend_mode):
|
||
"""添加纹理效果到图像"""
|
||
if texture_type == 'noise':
|
||
texture = self.generate_noise_texture(image.size[0])
|
||
elif texture_type == 'lines':
|
||
texture = self.generate_line_texture(image.size[0])
|
||
else:
|
||
return image
|
||
|
||
if texture_blend_mode == 'multiply':
|
||
return ImageChops.multiply(image, texture)
|
||
elif texture_blend_mode == 'overlay':
|
||
return ImageChops.overlay(image, texture)
|
||
else:
|
||
return image
|
||
|
||
def rotate_image_with_transparency(self, image, angle):
|
||
"""旋转带有透明度的图像"""
|
||
rotated_image = image.rotate(angle, expand=True)
|
||
return rotated_image
|
||
|
||
def save_processed_design_image(self, design_image, output_format='png'):
|
||
"""保存处理后的设计图像"""
|
||
try:
|
||
img_bytes = io.BytesIO()
|
||
# 确保使用包含透明背景的BGRA格式
|
||
design_image_pil = Image.fromarray(cv2.cvtColor(design_image, cv2.COLOR_BGRA2RGBA)).convert('RGBA')
|
||
|
||
if output_format == 'tiff':
|
||
design_image_pil.save(img_bytes, format='TIFF', save_all=True, compression='tiff_deflate')
|
||
else:
|
||
design_image_pil.save(img_bytes, format='PNG')
|
||
|
||
img_bytes.seek(0)
|
||
return img_bytes
|
||
except Exception as e:
|
||
print(f"保存处理后的设计图像时发生错误: {e}")
|
||
return None
|
||
|
||
def generate_tshirt_image(self, design_image, tshirt_image, config):
|
||
"""将花型图案合成到T恤图像上"""
|
||
# 合并默认配置和用户配置
|
||
config = {**self.DEFAULT_CONFIG, **config}
|
||
|
||
# 将设计图像从RGBA转换为BGRA（如果需要）
|
||
if isinstance(design_image, np.ndarray) and design_image.shape[2] == 4:
|
||
if design_image.dtype != np.uint8:
|
||
design_image = design_image.astype(np.uint8)
|
||
else:
|
||
# 如果输入是PIL图像，转换为OpenCV格式
|
||
if isinstance(design_image, Image.Image):
|
||
design_image = cv2.cvtColor(np.array(design_image), cv2.COLOR_RGBA2BGRA)
|
||
else:
|
||
raise ValueError("设计图像必须是PIL Image或带Alpha通道的NumPy数组")
|
||
|
||
# 对设计图像应用渐变效果
|
||
if config.get('enable_gradient_effect', self.DEFAULT_CONFIG['enable_gradient_effect']):
|
||
design_image_pil = Image.fromarray(cv2.cvtColor(design_image, cv2.COLOR_BGRA2RGBA)).convert("RGBA")
|
||
design_image_pil = self.add_gradient_repeat(
|
||
design_image_pil,
|
||
config.get('gradient_repeat_count', self.DEFAULT_CONFIG['gradient_repeat_count']),
|
||
config.get('gradient_width', self.DEFAULT_CONFIG['gradient_width']),
|
||
config.get('gradient_direction', self.DEFAULT_CONFIG['gradient_direction']),
|
||
config.get('gradient_type', self.DEFAULT_CONFIG['gradient_type']),
|
||
config.get('gradient_max_alpha', self.DEFAULT_CONFIG['gradient_max_alpha']),
|
||
config.get('gradient_start_alpha', self.DEFAULT_CONFIG['gradient_start_alpha']),
|
||
config.get('gradient_color', self.DEFAULT_CONFIG['gradient_color']),
|
||
config.get('gradient_blur_intensity', self.DEFAULT_CONFIG['gradient_blur_intensity']),
|
||
config.get('gradient_center', self.DEFAULT_CONFIG['gradient_center'])
|
||
)
|
||
design_image = cv2.cvtColor(np.array(design_image_pil), cv2.COLOR_RGBA2BGRA)
|
||
|
||
# 应用纹理效果到设计图案
|
||
if config.get('ptt_enable_texture_effect', self.DEFAULT_CONFIG['ptt_enable_texture_effect']):
|
||
design_image_pil = Image.fromarray(cv2.cvtColor(design_image, cv2.COLOR_BGRA2RGBA)).convert("RGBA")
|
||
design_image_pil = self.add_texture(
|
||
design_image_pil,
|
||
config.get('ptt_texture_type', self.DEFAULT_CONFIG['ptt_texture_type']),
|
||
config.get('ptt_texture_blend_mode', self.DEFAULT_CONFIG['ptt_texture_blend_mode'])
|
||
)
|
||
design_image = cv2.cvtColor(np.array(design_image_pil), cv2.COLOR_RGBA2BGRA)
|
||
|
||
# 进行设计图像增强处理
|
||
design_image_enhanced = self.enhance_design(design_image[:, :, :3], tshirt_image, config)
|
||
|
||
# 应用旋转效果到设计图案
|
||
ptt_design_rotation = config.get('ptt_design_rotation', self.DEFAULT_CONFIG['ptt_design_rotation'])
|
||
if ptt_design_rotation != 0:
|
||
design_image_pil = Image.fromarray(cv2.cvtColor(design_image, cv2.COLOR_BGRA2RGBA)).convert("RGBA")
|
||
design_image_pil = self.rotate_image_with_transparency(design_image_pil, ptt_design_rotation)
|
||
processed_design_image_with_alpha = cv2.cvtColor(np.array(design_image_pil), cv2.COLOR_RGBA2BGRA)
|
||
else:
|
||
processed_design_image_with_alpha = cv2.merge((design_image_enhanced, design_image[:, :, 3]))
|
||
|
||
# 保存处理后的设计图像
|
||
processed_design_io = None
|
||
if config.get('enable_save_processed_design', self.DEFAULT_CONFIG['enable_save_processed_design']):
|
||
processed_design_io = self.save_processed_design_image(
|
||
processed_design_image_with_alpha,
|
||
config.get('ptt_design_output_format', self.DEFAULT_CONFIG['ptt_design_output_format'])
|
||
)
|
||
|
||
# 调整设计图像大小
|
||
tshirt_height, tshirt_width = tshirt_image.shape[:2]
|
||
design_width = int(tshirt_width * config.get('ptt_design_size_ratio', self.DEFAULT_CONFIG['ptt_design_size_ratio']))
|
||
aspect_ratio = processed_design_image_with_alpha.shape[0] / processed_design_image_with_alpha.shape[1]
|
||
design_height = int(design_width * aspect_ratio)
|
||
design_image_resized = cv2.resize(processed_design_image_with_alpha, (design_width, design_height))
|
||
|
||
# 提取Alpha通道
|
||
alpha_channel = design_image_resized[:, :, 3] / 255.0
|
||
|
||
# 计算设计图像在T恤上的位置
|
||
ptt_design_offset = config.get('ptt_design_offset', self.DEFAULT_CONFIG['ptt_design_offset'])
|
||
design_position = (
|
||
int((tshirt_width - design_width) * ptt_design_offset[0]),
|
||
int((tshirt_height - design_height) * ptt_design_offset[1])
|
||
)
|
||
|
||
# 将设计图像叠加到T恤图像上
|
||
result_image = tshirt_image.copy()
|
||
self.overlay_image_alpha(result_image, design_image_resized[:, :, :3], design_position, alpha_channel)
|
||
|
||
# 返回结果图像和处理后的设计图像
|
||
return result_image, processed_design_io
|
||
|
||
def image_to_base64(self, image, format='png'):
|
||
"""将图像转换为base64字符串"""
|
||
try:
|
||
if isinstance(image, np.ndarray):
|
||
# 如果是OpenCV图像（NumPy数组），转换为PIL图像
|
||
if image.shape[2] == 3:
|
||
# BGR转RGB
|
||
image_pil = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
|
||
else:
|
||
# BGRA转RGBA
|
||
image_pil = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGRA2RGBA))
|
||
else:
|
||
# 已经是PIL图像
|
||
image_pil = image
|
||
|
||
# 保存为BytesIO对象
|
||
buffered = io.BytesIO()
|
||
image_pil.save(buffered, format=format.upper())
|
||
img_str = base64.b64encode(buffered.getvalue()).decode()
|
||
return img_str
|
||
except Exception as e:
|
||
print(f"将图像转换为base64时发生错误: {e}")
|
||
return None
|
||
|
||
def download_tshirt_images(self, config):
|
||
"""下载T恤图像列表"""
|
||
try:
|
||
# 首先检查是否提供了T恤图片URL列表
|
||
tshirt_urls = config.get('tshirt_urls', self.DEFAULT_CONFIG['tshirt_urls'])
|
||
if tshirt_urls and isinstance(tshirt_urls, list) and len(tshirt_urls) > 0:
|
||
tshirt_images = []
|
||
for url in tshirt_urls:
|
||
if self.is_valid_url(url):
|
||
tshirt_io = self.download_image(url)
|
||
if tshirt_io:
|
||
tshirt_image = cv2.imdecode(np.frombuffer(tshirt_io.getvalue(), np.uint8), cv2.IMREAD_COLOR)
|
||
if tshirt_image is not None:
|
||
tshirt_images.append(tshirt_image)
|
||
|
||
if tshirt_images:
|
||
return tshirt_images
|
||
|
||
# 如果没有提供URL或URL下载失败，则尝试使用本地模板
|
||
sample_tshirt_path = os.path.join(config.get('tshirt_image_path', self.DEFAULT_CONFIG['tshirt_image_path']), 'sample_tshirt.jpg')
|
||
if os.path.exists(sample_tshirt_path):
|
||
tshirt_image = cv2.imread(sample_tshirt_path)
|
||
return [tshirt_image]
|
||
else:
|
||
# 创建一个纯白色的示例T恤图像作为最后的备选
|
||
tshirt_image = np.ones((800, 600, 3), dtype=np.uint8) * 255
|
||
return [tshirt_image]
|
||
except Exception as e:
|
||
print(f"下载T恤图像时发生错误: {e}")
|
||
# 创建一个纯白色的示例T恤图像作为最后的备选
|
||
tshirt_image = np.ones((800, 600, 3), dtype=np.uint8) * 255
|
||
return [tshirt_image]
|
||
|
||
def is_valid_url(self, url):
|
||
"""检查URL是否有效"""
|
||
try:
|
||
result = urlparse(url)
|
||
return all([result.scheme, result.netloc])
|
||
except:
|
||
return False
|
||
|
||
def download_image(self, url):
|
||
"""下载图像"""
|
||
try:
|
||
if self.is_valid_url(url):
|
||
response = requests.get(url, verify=False, timeout=10)
|
||
if response.status_code == 200:
|
||
return io.BytesIO(response.content)
|
||
return None
|
||
except Exception as e:
|
||
print(f"下载图像失败: {str(e)}")
|
||
return None
|
||
|
||
async def pattern_to_tshirt(self, image_url, config=None):
|
||
"""将花型图案添加到T恤上（URL输入）"""
|
||
if not config:
|
||
config = {}
|
||
|
||
try:
|
||
# 下载花型图案
|
||
design_io = self.download_image(image_url)
|
||
if not design_io:
|
||
return {"status": "error", "message": "无法下载图像"}
|
||
|
||
return await self.pattern_to_tshirt_from_file(design_io.getvalue(), config)
|
||
except Exception as e:
|
||
import traceback
|
||
error_trace = traceback.format_exc()
|
||
print(f"处理图像时发生错误: {str(e)}\n{error_trace}")
|
||
return {"status": "error", "message": f"处理图像失败: {str(e)}"}
|
||
|
||
async def pattern_to_tshirt_from_file(self, file_content, config=None):
|
||
"""将花型图案添加到T恤上（文件输入）"""
|
||
if not config:
|
||
config = {}
|
||
|
||
try:
|
||
# 加载花型图案
|
||
design_io = io.BytesIO(file_content)
|
||
design_image = Image.open(design_io).convert("RGBA")
|
||
|
||
# 获取T恤图像列表
|
||
tshirt_images = self.download_tshirt_images(config)
|
||
if not tshirt_images:
|
||
return {"status": "error", "message": "无法获取T恤图像模板"}
|
||
|
||
results = []
|
||
processed_design_base64 = None
|
||
|
||
# 处理每个T恤图像
|
||
for tshirt_image in tshirt_images:
|
||
try:
|
||
# 生成合成图像
|
||
result_image, processed_design_io = self.generate_tshirt_image(design_image, tshirt_image, config)
|
||
|
||
# 转换为base64
|
||
result_base64 = self.image_to_base64(result_image)
|
||
|
||
# 如果有处理后的设计图像，也转换为base64
|
||
if processed_design_io and processed_design_base64 is None:
|
||
processed_design_image = Image.open(processed_design_io)
|
||
processed_design_base64 = self.image_to_base64(processed_design_image)
|
||
|
||
results.append({
|
||
"tshirt_image": result_base64
|
||
})
|
||
except Exception as e:
|
||
import traceback
|
||
error_trace = traceback.format_exc()
|
||
print(f"处理单个T恤图像时发生错误: {str(e)}\n{error_trace}")
|
||
# 继续处理下一个T恤图像
|
||
|
||
if not results:
|
||
return {"status": "error", "message": "所有T恤图像处理均失败"}
|
||
|
||
response = {
|
||
"status": "success",
|
||
"results": results
|
||
}
|
||
|
||
# 如果有处理后的设计图像，添加到响应中
|
||
if processed_design_base64:
|
||
response["processed_design"] = processed_design_base64
|
||
|
||
return response
|
||
|
||
except Exception as e:
|
||
import traceback
|
||
error_trace = traceback.format_exc()
|
||
print(f"处理图像时发生错误: {str(e)}\n{error_trace}")
|
||
return {"status": "error", "message": f"处理图像失败: {str(e)}"}
|
||
|
||
def cleanup(self):
|
||
"""清理资源"""
|
||
if torch.cuda.is_available():
|
||
torch.cuda.empty_cache()
|
||
gc.collect()
|