Update clothGen.py

clothGen.py

@@ -3,6 +3,11 @@ import pandas as pd
 from prompt_gen import prompt_gen
 import requests
 import os
+from openai import OpenAI
+import random
+import shutil
+from pathlib import Path
+from PIL import Image
 
 nv_prompt_file = pd.read_excel('汉服-女词库.xlsx')
 na_prompt_file = pd.read_excel('汉服-男词库.xlsx')
@@ -11,41 +16,51 @@ na_prompt = na_prompt_file.to_string(index=False)
 save_directory = "downloads"
 
 
+def prompt_nan(prompt):
+    client = OpenAI()
+    completion = client.chat.completions.create(
+        model="gpt-4o",
+        messages=[
+            {"role": "system",
+             "content": "You are a helpful assistant.", },
+            {"role": "user",
+             "content": "Please showcase the overall appearance of this Hanfu robe against a contrasting white  "
+                        "background. Highlight its intricate details and unique design elements, including" + prompt,
+             }
+        ]
+
+    )
+    print("change prompt: ")
+    print(completion.choices[0].message.content)
+    return completion.choices[0].message.content
+
+
 def pro_gen(advice, gender, index):
     prompt = prompt_gen(advice, gender)
     start_index = prompt.find("Begin")
     if start_index == -1:
         start_index = prompt.find("begin")
-    intro_index = prompt.find("服饰风格介绍")
-    cloth_intro = ""
     prompt__gen = ""
     if start_index != -1:
         start_index += len("Begin\n")
         end_index = prompt.find("End")
         if end_index != -1:
             prompt__gen = prompt[start_index:end_index]
+            # if gender == "男":
+            #     prompt__gen = prompt_nan(prompt__gen)
             filename = os.path.join(save_directory, f"prompt_{index}.txt")
             with open(filename, "w") as file:
                 file.write(prompt__gen)
-            print(prompt__gen)
+            # print(prompt__gen)
         else:
             print("No 'promptEnd' found after 'prompt'.")
     else:
         print("No 'prompt' found in the text.")
-    if intro_index != -1:
-        intro_index += len("服饰风格介绍\n")
-        cloth_intro = ("汉服，是汉民族的传统服饰。又称衣冠、衣裳、汉装。汉服是中国“衣冠上国”“礼仪之邦”“锦绣中华”的体现，承载了中国的染织绣等杰出"
-                       "工艺和美学，传承了30多项中国非物质文化遗产以及受保护的中国工艺美术。\n") + prompt[intro_index:]
-        filename = os.path.join(save_directory, f"cloth_intro_{index}.txt")
-        with open(filename, "w") as file:
-            file.write(cloth_intro)
-        print(cloth_intro)
-    else:
-        print("No '服饰风格介绍' found.")
     return prompt__gen
 
 
 def generate(lora_path, prompt__gen, index):
+    # print(prompt__gen)
     handler = fal_client.submit(
         "fal-ai/fast-sdxl",
         arguments={
@@ -56,7 +71,7 @@ def generate(lora_path, prompt__gen, index):
             "image_size": "portrait_4_3",
             "num_inference_steps": 28,
             "guidance_scale": 7.5,
-            "num_images": 2,
+            "num_images": 1,
             "loras": [{"path": lora_path, "scale": 0.7}],
             "embeddings": [],
             "safety_checker_version": "v1",
@@ -70,28 +85,120 @@ def generate(lora_path, prompt__gen, index):
     for image in result['images']:
         response = requests.get(image['url'])
         if response.status_code == 200:
-            filename = os.path.join(save_directory, f"gen_cloth_{image_index}.jpeg")
+            filename = os.path.join(save_directory, f"cloth_{image_index}.jpeg")
             with open(filename, 'wb') as f:
                 f.write(response.content)
             image_index += 1
         else:
             print(f"Failed to download image from {image['url']}")
 
+    client = OpenAI()
+    completion = client.chat.completions.create(
+        model="gpt-4o",
+        messages=[
+            {"role": "system",
+             "content": "You are a helpful assistant.", },
+            {"role": "user",
+             "content": prompt__gen + "以上是一段对于一套汉服的描述，请根据描述内容对该套汉服进行介绍。要求以介绍的口吻输出内容",
+             }
+        ]
+
+    )
+    cloth_intro = ("汉服，是汉民族的传统服饰。又称衣冠、衣裳、汉装。汉服是中国“衣冠上国”“礼仪之邦”“锦绣中华”的体现，承载了中国的染织绣等杰出"
+                   "工艺和美学，传承了30多项中国非物质文化遗产以及受保护的中国工艺美术。\n") + completion.choices[
+                      0].message.content
+    filename = os.path.join(save_directory, f"cloth_intro_{index * 2 - 1}.txt")
+    with open(filename, "w") as file:
+        file.write(cloth_intro)
+
+
+def convert_image_to_jpeg(input_path, output_path):
+    try:
+        image = Image.open(input_path)
+        if image.mode in ('RGBA', 'LA'):
+            image = image.convert('RGB')
+        image.save(output_path, 'JPEG')
+    except Exception as e:
+        print(f"转换图像时出错: {e}")
+
+
+def pic_match(prompt__gen, cates, folder_path, intro_path, index):
+    client = OpenAI()
+    completion = client.chat.completions.create(
+        model="gpt-4o",
+        messages=[
+            {"role": "system",
+             "content": "You are a helpful assistant.", },
+            {"role": "user",
+             "content": prompt__gen + "以上是关于一套汉族服饰的描述，请根据描述内容从以下几种颜色中选择最符合描述的一种，可选颜色包括：" + cates
+                        + ". 仅需输出一种颜色名称，不要带任何符号",
+             }
+        ]
+
+    )
+    print(f"Selected color: {completion.choices[0].message.content}")
+
+    folder_path = os.path.join(folder_path, completion.choices[0].message.content)
+    files = os.listdir(folder_path)
+    random_file = random.choice(files)
+    source_file_path = os.path.join(folder_path, random_file)
+    file_prefix, file_ext = os.path.splitext(random_file)
+    target_file_path = os.path.join(save_directory, f"cloth_{index * 2}.jpeg")
+    convert_image_to_jpeg(source_file_path, target_file_path)
+
+    file_extension = ".txt"
+    search_path = Path(intro_path)
+    for file in search_path.glob(f"{file_prefix}*{file_extension}"):
+        if file.is_file():
+            with open(file, "r") as f:
+                content = f.read()
+                client = OpenAI()
+                completion = client.chat.completions.create(
+                    model="gpt-4o",
+                    messages=[
+                        {"role": "system",
+                         "content": "You are a helpful assistant.", },
+                        {"role": "user",
+                         "content": content + "以上是一段对于一套汉服的描述，请根据描述内容对该套汉服进行介绍。要求以介绍的口吻输出内容",
+                         }
+                    ]
+
+                )
+                cloth_intro = ("汉服，是汉民族的传统服饰。又称衣冠、衣裳、汉装。汉服是中国“衣冠上国”“礼仪之邦”“锦绣中华”的体现，承载了中国的染织绣等杰出"
+                               "工艺和美学，传承了30多项中国非物质文化遗产以及受保护的中国工艺美术。\n") + \
+                              completion.choices[0].message.content
+                filename = os.path.join(save_directory, f"cloth_intro_{index * 2}.txt")
+                with open(filename, "w") as file:
+                    file.write(cloth_intro)
+
+    return target_file_path
+
 
 def cloth_gen(gender):
+    cates = "Black, Blue, Green, Orange, Pink, Red, Violet, White, Yellow"
     lora_path = "https://huggingface.co/PPSharks/PPSharksModels/resolve/main/NV.safetensors"
+    folder_path = "database/female"
+    intro_path = "database/female_intro"
     if gender == "男":
         lora_path = "https://huggingface.co/PPSharks/PPSharksModels/resolve/main/NA.safetensors"
-    else:
+        cates = "Black, Blue, Green, Brown, Red, Violet"
+        folder_path = "database/male"
+        intro_path = "database/male_intro"
+    elif gender == "女":
         lora_path = "https://huggingface.co/PPSharks/PPSharksModels/resolve/main/NV.safetensors"
+        cates = "Black, Blue, Green, Orange, Pink, Red, Violet, White, Yellow"
+        folder_path = "database/female"
+        intro_path = "database/female_intro"
 
     cloth_image = []
     for i in range(1, 4):
         with open(os.path.join(save_directory, f"prompt_{i}.txt"), "r") as file:
             prompt__gen = file.read()
         generate(lora_path, prompt__gen, i)
-        cloth_image.append(os.path.join(save_directory, f"gen_cloth_{i*2-1}.jpeg"))
-        cloth_image.append(os.path.join(save_directory, f"gen_cloth_{i*2}.jpeg"))
+        cloth_image.append(os.path.join(save_directory, f"cloth_{i*2-1}.jpeg"))
+        pic_path = pic_match(prompt__gen, cates, folder_path, intro_path, i)
+        cloth_image.append(pic_path)
+
     with open(os.path.join(save_directory, f"cloth_intro_1.txt"), "r") as file:
         cloth_intro = file.read()
     return cloth_image, cloth_image[0], cloth_intro