app.py

# app.py (Gradio界面)
import gradio as gr
import requests
from PIL import Image
import numpy as np
from sklearn.cluster import KMeans
import time
import random
import os
import torch

# 模型管理器
from models.model_manager import ModelManager

# 初始化模型管理器
model_manager = ModelManager()

def upload_and_analyze(image_path):
    """分析上传的图片"""
    try:
        if image_path is None:
            return {}, {}, []
        
        # 打开图片
        image = Image.open(image_path).convert('RGB')
        
        # 生成图像描述
        caption = model_manager.generate_caption(image)
        
        # 基于图像描述进行智能分析
        analysis_result = analyze_image_content(image, caption)
        
        # 基于分析结果生成个性化建议
        suggestions = generate_personalized_suggestions(analysis_result, caption)
        
        # 创建选择选项
        choices = list(suggestions.keys())
        
        return analysis_result, suggestions, gr.Radio(choices=choices, value=choices[0] if choices else None)
        
    except Exception as e:
        error_result = {"错误": f"分析失败: {str(e)}"}
        return error_result, {}, []

def analyze_image_content(image, caption):
    """基于图像和描述进行深度分析"""
    # 分析图像颜色
    colors = extract_dominant_colors(image)
    
    # 根据描述推断风格类型
    style_type = infer_style_from_caption(caption)
    
    # 根据描述推断服装类别
    clothing_category = infer_clothing_category(caption)
    
    # 根据风格推荐适合场景
    suitable_scenes = get_suitable_scenes(style_type)
    
    return {
        "图像描述": caption,
        "检测到的颜色": colors,
        "风格类型": style_type,
        "服装类别": clothing_category,
        "适合场景": suitable_scenes,
        "图像尺寸": f"{image.width} x {image.height}",
        "分析时间": time.strftime("%Y-%m-%d %H:%M:%S")
    }

def extract_dominant_colors(image):
    """提取图像主要颜色"""
    # 调整图像大小以提高处理速度
    image = image.resize((150, 150))
    
    # 转换为numpy数组
    img_array = np.array(image)
    
    # 重塑为颜色列表
    pixels = img_array.reshape(-1, 3)
    
    # 使用KMeans聚类找到主要颜色
    kmeans = KMeans(n_clusters=3, random_state=42, n_init=10)
    kmeans.fit(pixels)
    
    # 将RGB值转换为颜色名称
    color_names = []
    for color in kmeans.cluster_centers_:
        color_name = rgb_to_color_name(color)
        color_names.append(color_name)
    
    return color_names

def rgb_to_color_name(rgb):
    """将RGB值转换为颜色名称"""
    r, g, b = rgb.astype(int)
    
    # 简单的颜色映射
    if r > 200 and g > 200 and b > 200:
        return "白色"
    elif r < 50 and g < 50 and b < 50:
        return "黑色"
    elif r > g and r > b:
        if r > 150:
            return "红色"
        else:
            return "深红色"
    elif g > r and g > b:
        if g > 150:
            return "绿色"
        else:
            return "深绿色"
    elif b > r and b > g:
        if b > 150:
            return "蓝色"
        else:
            return "深蓝色"
    elif r > 150 and g > 150:
        return "黄色"
    elif r > 100 and b > 100:
        return "紫色"
    elif g > 100 and b > 100:
        return "青色"
    else:
        return "灰色"

def infer_style_from_caption(caption):
    """根据图像描述推断风格类型"""
    caption_lower = caption.lower()
    
    style_keywords = {
        "商务正装": ["suit", "formal", "business", "office", "professional", "tie", "blazer", "西装", "正装", "商务"],
        "休闲风": ["casual", "relaxed", "comfortable", "everyday", "jeans", "t-shirt", "休闲", "日常"],
        "运动风": ["sport", "athletic", "gym", "fitness", "running", "training", "运动", "健身"],
        "时尚潮流": ["fashion", "trendy", "stylish", "modern", "chic", "designer", "时尚", "潮流"],
        "复古风": ["vintage", "retro", "classic", "traditional", "old-fashioned", "复古", "经典"],
        "街头风": ["street", "urban", "hip-hop", "cool", "edgy", "街头", "嘻哈"],
        "优雅风": ["elegant", "sophisticated", "graceful", "refined", "classy", "优雅", "高贵"]
    }
    
    for style, keywords in style_keywords.items():
        if any(keyword in caption_lower for keyword in keywords):
            return style
    
    return "休闲风"  # 默认风格

def infer_clothing_category(caption):
    """根据描述推断服装类别"""
    caption_lower = caption.lower()
    
    categories = {
        "上衣": ["shirt", "blouse", "top", "jacket", "sweater", "hoodie", "blazer", "衬衫", "上衣", "外套"],
        "下装": ["pants", "jeans", "skirt", "shorts", "trousers", "裤子", "短裤", "裙子"],
        "连衣裙": ["dress", "gown", "frock", "连衣裙", "礼服"],
        "外套": ["coat", "jacket", "cardigan", "blazer", "外套", "大衣"],
        "配饰": ["hat", "bag", "shoes", "belt", "jewelry", "帽子", "包", "鞋子", "配饰"],
        "全套搭配": ["outfit", "ensemble", "look", "style", "搭配", "整套"]
    }
    
    for category, keywords in categories.items():
        if any(keyword in caption_lower for keyword in keywords):
            return category
    
    return "服装单品"  # 默认类别

def get_suitable_scenes(style_type):
    """根据风格类型推荐适合场景"""
    scene_mapping = {
        "商务正装": ["办公室", "商务会议", "正式场合", "面试"],
        "休闲风": ["日常出街", "朋友聚会", "购物", "咖啡约会"],
        "运动风": ["健身房", "运动", "户外活动", "晨跑"],
        "时尚潮流": ["时尚派对", "约会", "拍照", "社交活动"],
        "复古风": ["艺术展", "文艺活动", "复古主题活动", "拍摄"],
        "街头风": ["街拍", "音乐节", "朋友聚会", "潮流活动"],
        "优雅风": ["晚宴", "正式聚会", "典礼", "高端场所"]
    }
    
    return scene_mapping.get(style_type, ["日常", "休闲", "约会"])

def generate_personalized_suggestions(analysis_result, caption):
    """基于分析结果生成个性化建议"""
    style_type = analysis_result["风格类型"]
    clothing_category = analysis_result["服装类别"]
    colors = analysis_result["检测到的颜色"]
    
    suggestions = {}
    
    # 根据检测到的风格生成建议
    if style_type == "商务正装":
        suggestions = {
            "经典商务": f"保持{style_type}特色，搭配{colors[0]}系配饰",
            "现代商务": f"在{style_type}基础上加入现代元素",
            "休闲商务": f"将{style_type}与休闲元素结合",
            "时尚商务": f"{style_type}融入时尚潮流元素"
        }
    elif style_type == "休闲风":
        suggestions = {
            "舒适休闲": f"强化{style_type}的舒适感，主色调{colors[0]}",
            "时尚休闲": f"{style_type}加入时尚元素",
            "运动休闲": f"{style_type}融入运动风格",
            "优雅休闲": f"{style_type}提升优雅感"
        }
    elif style_type == "运动风":
        suggestions = {
            "专业运动": f"增强{style_type}的功能性",
            "休闲运动": f"{style_type}与日常穿着结合",
            "时尚运动": f"{style_type}加入潮流设计元素",
            "户外运动": f"强化{style_type}的户外适应性"
        }
    else:
        suggestions = {
            f"经典{style_type}": f"保持原有{style_type}特色",
            f"现代{style_type}": f"{style_type}加入现代元素",
            f"融合风格": f"{style_type}与其他风格混搭",
            f"个性化{style_type}": f"基于{colors[0]}色调的个性化{style_type}"
        }
    
    return suggestions

def generate_designs(selected_suggestion, progress=gr.Progress()):
    """根据选择的建议生成设计"""
    try:
        if not selected_suggestion:
            return [], gr.Radio(choices=[])
        
        progress(0.1, desc="准备设计提示...")
        
        # 生成设计图像的提示词
        design_prompts = {
            "经典商务": "professional business suit, modern cut, high-quality fabric, clean lines, neutral colors",
            "现代商务": "contemporary business attire, innovative design, slim fit, premium materials",
            "休闲商务": "business casual outfit, comfortable yet professional, versatile style",
            "时尚商务": "fashion-forward business wear, runway inspired, statement piece",
            "舒适休闲": "casual comfort wear, soft fabrics, relaxed fit, everyday style",
            "时尚休闲": "stylish casual outfit, trendy elements, urban chic",
            "运动休闲": "athleisure wear, sporty elements, comfortable and functional",
            "优雅休闲": "elegant casual attire, sophisticated details, refined look",
            "专业运动": "performance sportswear, technical fabrics, functional design",
            "休闲运动": "casual athletic wear, versatile for sports and daily use",
            "时尚运动": "fashion sportswear, trendy athletic style, streetwear influence",
            "户外运动": "outdoor adventure wear, durable materials, weather-resistant",
            "经典复古风": "vintage retro style, classic silhouette, nostalgic elements",
            "现代复古风": "contemporary take on retro fashion, updated classics",
            "融合风格": "fusion fashion, mixed styles, innovative combination",
            "个性化街头风": "personalized streetwear, unique designs, urban style",
            "经典优雅风": "timeless elegant fashion, sophisticated details, refined look",
            "现代优雅风": "modern elegant attire, contemporary sophistication"
        }
        
        prompt = design_prompts.get(selected_suggestion, "fashion design, stylish clothing")
        
        # 生成设计图像
        design_images = []
        design_choices = []
        
        for i in range(3):  # 生成3个设计
            try:
                progress(0.2 + i*0.25, desc=f"生成设计方案 {i+1}/3...")
                
                # 生成真实的设计图像
                image = model_manager.generate_image(
                    prompt=f"{prompt}, design {i+1}, high detail, fashion illustration",
                    negative_prompt="blurry, low quality, distorted, text, watermark",
                    num_inference_steps=30
                )
                
                if image:
                    design_images.append(image)
                    design_choices.append(f"{selected_suggestion} 设计方案 {i+1}")
            except Exception as e:
                print(f"生成设计 {i+1} 失败: {e}")
                # 创建占位图像
                width, height = 512, 512
                img = Image.new('RGB', (width, height), 
                               color=(random.randint(0, 255), 
                                      random.randint(0, 255), 
                                      random.randint(0, 255)))
                design_images.append(img)
                design_choices.append(f"{selected_suggestion} 设计方案 {i+1}")
        
        progress(0.95, desc="完成设计生成")
        return design_images, gr.Radio(choices=design_choices, value=design_choices[0] if design_choices else None)
        
    except Exception as e:
        print(f"设计生成错误: {e}")
        return [], gr.Radio(choices=[])

def generate_3d_fitting(selected_design, progress=gr.Progress()):
    """生成3D试穿效果"""
    try:
        if not selected_design:
            return None
        
        progress(0.1, desc="准备3D试穿...")
        
        # 生成3D试穿效果的提示词
        fitting_prompt = f"3D fashion fitting, virtual try-on, {selected_design}, realistic human model, full body, studio lighting"
        
        progress(0.3, desc="生成3D模型...")
        
        # 使用模型生成3D试穿图像
        if model_manager.controlnet_pipeline:
            # 如果有ControlNet，使用更高级的生成
            try:
                # 这里简化了ControlNet的使用，实际需要姿势图等输入
                image = model_manager.controlnet_pipeline(
                    prompt=fitting_prompt,
                    negative_prompt="blurry, distorted, low quality, unrealistic, extra limbs",
                    num_inference_steps=35,
                    guidance_scale=8.0
                ).images[0]
                progress(0.9, desc="渲染3D效果")
                return image
            except Exception as e:
                print(f"使用ControlNet生成失败: {e}")
        
        # 回退到普通SD模型
        progress(0.4, desc="使用标准模型生成...")
        image = model_manager.generate_image(
            prompt=fitting_prompt,
            negative_prompt="blurry, distorted, low quality, unrealistic, extra limbs",
            num_inference_steps=35
        )
        
        progress(0.9, desc="完成3D渲染")
        return image
        
    except Exception as e:
        print(f"3D试穿生成错误: {e}")
        return None

def create_gradio_interface():
    """创建Gradio用户界面"""
    
    with gr.Blocks(title="AI时尚设计师", theme="soft") as demo:
        gr.Markdown("# 🎨 AI时尚设计师")
        gr.Markdown("上传图片，获得专业的服装设计建议和3D试穿效果")
        
        with gr.Row():
            with gr.Column(scale=1):
                image_input = gr.Image(type="filepath", label="上传参考图片", height=400)
                analyze_btn = gr.Button("分析风格", variant="primary")
            
            with gr.Column(scale=2):
                analysis_output = gr.JSON(label="风格分析结果")
        
        with gr.Tab("设计建议"):
            suggestions_output = gr.JSON(label="个性化设计建议")
            suggestion_choice = gr.Radio(label="选择设计方向", interactive=True)
            generate_designs_btn = gr.Button("生成设计方案", variant="primary")
        
        with gr.Tab("样衣设计"):
            designs_gallery = gr.Gallery(label="样衣设计图", columns=3, height=400)
            design_choice = gr.Radio(label="选择设计方案", interactive=True)
            generate_3d_btn = gr.Button("生成3D试穿效果", variant="primary")
        
        with gr.Tab("3D试穿效果"):
            fitting_result = gr.Image(label="3D试穿效果", height=500)
        
        # 添加示例图片
        examples_dir = "examples"
        if os.path.exists(examples_dir):
            example_files = [os.path.join(examples_dir, f) for f in os.listdir(examples_dir) 
                            if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
            
            gr.Examples(
                examples=example_files[:4],  # 最多显示4个示例
                inputs=image_input,
                outputs=[analysis_output, suggestions_output, suggestion_choice],
                fn=upload_and_analyze,
                cache_examples=True,
                label="示例图片"
            )
        
        # 事件绑定
        analyze_btn.click(
            fn=upload_and_analyze,
            inputs=[image_input],
            outputs=[analysis_output, suggestions_output, suggestion_choice]
        )
        
        generate_designs_btn.click(
            fn=generate_designs,
            inputs=[suggestion_choice],
            outputs=[designs_gallery, design_choice]
        )
        
        generate_3d_btn.click(
            fn=generate_3d_fitting,
            inputs=[design_choice],
            outputs=[fitting_result]
        )
        
        # 添加清理按钮
        clear_btn = gr.Button("清理内存", variant="secondary")
        clear_btn.click(
            fn=model_manager.cleanup,
            inputs=[],
            outputs=[]
        )
        gr.Markdown("> **提示**: 生成图像后，点击'清理内存'按钮可以释放GPU资源")
    
    return demo

if __name__ == "__main__":
    # 检查并创建示例目录
    examples_dir = "examples"
    if not os.path.exists(examples_dir):
        os.makedirs(examples_dir)
        print(f"创建了示例目录: {examples_dir}")
        print("请在此目录中添加示例图片以便在界面中使用")
    
    demo = create_gradio_interface()
    demo.queue(concurrency_count=1)  # 限制并发以避免内存问题
    demo.launch(
        server_name="0.0.0.0", 
        server_port=7860, 
        share=True,
        favicon_path="favicon.ico" if os.path.exists("favicon.ico") else None
    )