Update app.py

app.py

@@ -1,11 +1,15 @@
 # app.py (Gradio界面)
 import gradio as gr
-from main import app
 import requests
 from PIL import Image
-import json
+import numpy as np
+from sklearn.cluster import KMeans
+import time
+import random
+import os
+import torch
 
-# 导入模型管理器
+# 模型管理器
 from models.model_manager import ModelManager
 
 # 初始化模型管理器
@@ -18,27 +22,16 @@ def upload_and_analyze(image_path):
             return {}, {}, []
         
         # 打开图片
-        image = Image.open(image_path)
+        image = Image.open(image_path).convert('RGB')
         
         # 生成图像描述
         caption = model_manager.generate_caption(image)
         
-        # 模拟风格分析结果
-        analysis_result = {
-            "图像描述": caption,
-            "检测到的颜色": ["蓝色", "白色", "黑色"],
-            "风格类型": "休闲风",
-            "服装类别": "上衣",
-            "适合场景": ["日常", "休闲", "约会"]
-        }
+        # 基于图像描述进行智能分析
+        analysis_result = analyze_image_content(image, caption)
         
-        # 生成设计建议
-        suggestions = {
-            "建议1": "现代简约风格搭配",
-            "建议2": "复古经典款式",
-            "建议3": "运动休闲风格",
-            "建议4": "商务正装风格"
-        }
+        # 基于分析结果生成个性化建议
+        suggestions = generate_personalized_suggestions(analysis_result, caption)
         
         # 创建选择选项
         choices = list(suggestions.keys())
@@ -49,18 +42,207 @@ def upload_and_analyze(image_path):
         error_result = {"错误": f"分析失败: {str(e)}"}
         return error_result, {}, []
 
-def generate_designs(selected_suggestion):
+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 = {
-            "建议1": "modern minimalist clothing design, clean lines, neutral colors",
-            "建议2": "vintage classic fashion design, retro style, elegant",
-            "建议3": "sporty casual wear design, comfortable, athletic",
-            "建议4": "business formal attire, professional, sophisticated"
+            "经典商务": "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")
@@ -71,44 +253,75 @@ def generate_designs(selected_suggestion):
         
         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}",
-                    negative_prompt="blurry, low quality, distorted",
-                    num_inference_steps=20
+                    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"设计方案 {i+1}")
+                    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}")
         
-        # 如果没有成功生成图像，返回空结果
-        if not design_images:
-            return [], gr.Radio(choices=[])
-        
+        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):
+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"
+        fitting_prompt = f"3D fashion fitting, virtual try-on, {selected_design}, realistic human model, full body, studio lighting"
+        
+        progress(0.3, desc="生成3D模型...")
         
         # 使用模型生成3D试穿图像
-        fitting_image = model_manager.generate_image(
+        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",
-            num_inference_steps=25
+            negative_prompt="blurry, distorted, low quality, unrealistic, extra limbs",
+            num_inference_steps=35
         )
         
-        return fitting_image
+        progress(0.9, desc="完成3D渲染")
+        return image
         
     except Exception as e:
         print(f"3D试穿生成错误: {e}")
@@ -117,27 +330,45 @@ def generate_3d_fitting(selected_design):
 def create_gradio_interface():
     """创建Gradio用户界面"""
     
-    with gr.Blocks(title="AI时尚设计师") as demo:
+    with gr.Blocks(title="AI时尚设计师", theme="soft") as demo:
         gr.Markdown("# 🎨 AI时尚设计师")
         gr.Markdown("上传图片，获得专业的服装设计建议和3D试穿效果")
         
-        with gr.Tab("1. 图片上传与分析"):
-            image_input = gr.Image(type="filepath", label="上传参考图片")
-            analyze_btn = gr.Button("分析风格")
-            analysis_output = gr.JSON(label="风格分析结果")
+        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("2. 设计建议"):
-            suggestions_output = gr.JSON(label="设计建议")
-            suggestion_choice = gr.Radio(label="选择设计建议")
-            generate_designs_btn = gr.Button("生成样衣设计")
+        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("3. 样衣设计"):
-            designs_gallery = gr.Gallery(label="样衣设计图")
-            design_choice = gr.Radio(label="选择设计")
-            generate_3d_btn = gr.Button("生成3D试穿")
+        with gr.Tab("3D试穿效果"):
+            fitting_result = gr.Image(label="3D试穿效果", height=500)
         
-        with gr.Tab("4. 3D试穿"):
-            fitting_result = gr.Image(label="3D试穿效果")
+        # 添加示例图片
+        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(
@@ -157,9 +388,31 @@ def create_gradio_interface():
             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.launch(server_name="0.0.0.0", server_port=7860, share=True)
+    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
+    )