Merge pull request #22 from DanielZhangyc/dev

Feat: Upload Images

.gitignore

@@ -5,4 +5,5 @@
 .env
 .env.*
 !.env.example
-config/config.yaml
+config/config.yaml
+cache/

README.md

@@ -38,10 +38,26 @@ pic_example:
   type: "名称"
 ```
 
+## 上传图片
+
+选择 `upload images` 页面，点击`添加表情包`按钮，上传图片即可。
+
+可以选择多张表情包。
+
+可以启用`使用VLM自动生成文件名`，程序会使用VLM为每张图片生成文件名。
+
+每次上传完成后，需要叉掉所有已上传的文件或者刷新页面，否则会重复保存。
+
+上传完成后，需要重新生成缓存。
+
 ## 图片打标 (WIP)
 
 选择 `label images` 页面，选择图片文件夹，点击 `使用VLM生成描述` ，选择需要的描述后重命名文件即可。
 
+程序使用文件名生成embedding用于检索。如果你希望给一个图片多个embedding，你可以在文件名中使用`-`分隔。
+
+
+
 ## Demo
 
 <img width="256" alt="3bfb772e239f3437a13d46252aab1e1d" src="https://github.com/user-attachments/assets/d7e02f8f-205d-42ef-9c80-49f98aff64a6" />

app.py

@@ -1,8 +1,28 @@
 import streamlit as st
 import random
+import os
 import yaml
 from services.image_search import ImageSearch
 from config.settings import config, reload_config
+from services.utils import verify_folder
 
-pg = st.navigation([st.Page("pages/VVQuest.py"), st.Page("pages/label_images.py")])
+
+def delete_all_files_in_folder(folder_path):
+    try:
+        # 遍历文件夹中的所有文件和子文件夹
+        for root, dirs, files in os.walk(folder_path, topdown=False):
+            # 删除所有文件
+            for file in files:
+                file_path = os.path.join(root, file)
+                os.remove(file_path)
+            # 删除所有空文件夹
+            for dir in dirs:
+                dir_path = os.path.join(root, dir)
+                os.rmdir(dir_path)
+    except Exception as e:
+        print(f"删除过程中出现错误: {e}")
+
+verify_folder(os.path.join(config.base_dir, "cache"))
+delete_all_files_in_folder(os.path.join(config.base_dir, "cache"))
+pg = st.navigation([st.Page("pages/VVQuest.py"), st.Page("pages/label_images.py"), st.Page("pages/upload_images.py")])
 pg.run()

pages/label_images.py

@@ -14,19 +14,7 @@ from config.settings import config, reload_config
 from pages.utils import *
 from services.label_memes import LabelMemes
 
-IMAGE_DIRS_PATH = 'data/image_dirs'
-if not os.path.exists(IMAGE_DIRS_PATH):
-    os.makedirs(IMAGE_DIRS_PATH)
 
-# 获取image_dirs下的所有文件夹
-def get_image_dirs():
-    dirs = ['data/images']  # 默认包含原始图片目录
-    if os.path.exists(IMAGE_DIRS_PATH):
-        for item in os.listdir(IMAGE_DIRS_PATH):
-            item_path = os.path.join(IMAGE_DIRS_PATH, item)
-            if os.path.isdir(item_path):
-                dirs.append(item_path)
-    return dirs
 
 st.set_page_config(
     page_title="LabelImages",
@@ -88,10 +76,10 @@ def pregenerate_label(img_path, label_obj:LabelMemes, result_dict):
             return True
         except Exception as e:
             print(f'pregenerate_label failed: {str(e)}')
+            time.sleep(1)
 
 
 with st.sidebar:
-    # 使用selectbox替代text_input
     st.selectbox(
         '选择图片文件夹',
         options=get_image_dirs(),
@@ -118,12 +106,12 @@ if os.path.exists(st.session_state.image_folder_name):
 
     st.write(st.session_state.image_folder_name)
     img_path = st.session_state.all_images_path[st.session_state.image_index]
-    st.write(img_path)
-    with Image.open(img_path) as img:
-        img_obj = img.copy()
-    img_obj = np.array(img_obj)
-    img_obj = resize_image(img_obj, 256)
-    st.image(img_obj)
+    # st.write(img_path)
+    # with Image.open(img_path) as img:
+    #     img_obj = img.copy()
+    # img_obj = np.array(img_obj)
+    # img_obj = resize_image(img_obj, 256)
+    st.image(img_path, width=256)
 
     col3, col4, col5 = st.columns([1, 1, 1])
 
@@ -148,15 +136,15 @@ if os.path.exists(st.session_state.image_folder_name):
     def use_vlm_result_to_generate_buttons():
         try:
             name_list = st.session_state.img_analyse_result
-            colB1, colB2, colB3 = st.columns([1, 1, 1])
-            for index, i in enumerate([colB1, colB2, colB3]):
+            colB1, colB2, colB3, colB4 = st.columns([1, 1, 1, 1])
+            for index, i in enumerate([colB1, colB2, colB3, colB4]):
                 with i:
                     def create_onc(inner_index):
                         def onc():
                             st.session_state.new_file_name += name_list[inner_index]
                         return onc
-
-                    st.button(f"添加 \"{name_list[index]}\" 到文件名", on_click=create_onc(index),key=f'generate_clicked_{index}')
+                    if not name_list[index] == '':
+                        st.button(f"添加 \"{name_list[index]}\" 到文件名", on_click=create_onc(index),key=f'generate_clicked_{index}')
 
                     # auto mode
                     if st.session_state.auto_generate_labels:

pages/upload_images.py

@@ -0,0 +1,84 @@
+import datetime
+import os
+import shutil
+import time
+
+import numpy as np
+import streamlit as st
+import random
+import yaml
+from PIL import Image
+import threading
+
+from services.image_search import ImageSearch
+from config.settings import config, reload_config
+from pages.utils import *
+from services.label_memes import LabelMemes
+from services.utils import verify_folder
+
+if 'label_meme_obj' not in st.session_state:
+    st.session_state.label_meme_obj = LabelMemes()
+if 'image_folder_name' not in st.session_state:
+    st.session_state.image_folder_name = 'data/images'  # 默认使用原始图片目录
+
+with st.sidebar:
+    st.selectbox(
+        '选择图片文件夹',
+        options=get_image_dirs(),
+        key='image_folder_name',
+        help='选择导入后的图片文件夹'
+    )
+
+    st.checkbox(
+        '使用VLM自动生成文件名',
+        key='auto_generate_labels',
+        value=True
+    )
+
+    uploaded_images = st.file_uploader(label='添加表情包', accept_multiple_files=True, type=['png', 'jpg', 'jpeg', 'gif'])
+
+CACHE_PATH = os.path.join(config.base_dir, 'cache')
+verify_folder(CACHE_PATH)
+
+def label_image(image_path, show_result_area):
+    for i in range(15):
+        try:
+            l = st.session_state.label_meme_obj.label_image(image_path)
+            return f'{l[0]}-{l[1]}-{l[2]}-{l[3]}'
+        except Exception as e:
+            show_result_area.error(f'第{i + 1}次尝试生成描述失败, 错误信息: {e}')
+            time.sleep(1)
+    show_result_area.error('生成描述失败')
+    return False
+
+if uploaded_images:
+    show_image_area = st.empty()
+    show_result_area= st.empty()
+    for uploaded_image in uploaded_images:
+        try:
+            img_name = uploaded_image.name
+        except Exception as e:
+            img_name = f'{datetime.datetime.now()}.gif'
+        # 使用 PIL 打开图片并显示
+        img = Image.open(uploaded_image)
+        show_image_area.image(img, caption=uploaded_image.name, width=200)
+        cache_path = os.path.join(CACHE_PATH, img_name)
+        if img.mode == 'P': #TODO: 检查这行代码是否不需要
+            img = img.convert('RGB')
+        img.save(os.path.join(cache_path))
+        # 生成图片描述
+        if st.session_state.auto_generate_labels:
+            res = label_image(cache_path, show_result_area)
+            if res:
+                img_name = res + os.path.splitext(img_name)[1]
+
+        # 保存上传的图片到指定文件夹
+        save_image_path = os.path.join(st.session_state.image_folder_name,img_name )
+        if os.path.exists(save_image_path):
+            show_result_area.error(f'图片 {save_image_path} 已存在')
+        else:
+            if img.mode == 'P':
+                img = img.convert('RGB')
+            img.save(save_image_path)
+            show_result_area.success(f'图片 {save_image_path} 已保存')
+    uploaded_images = [] # 这行代码没用，那个组件仍然会保存已上传的文件

pages/utils.py

@@ -1,7 +1,7 @@
 import os
 
 import cv2
-
+from config.settings import config
 
 def get_all_file_paths(folder_path):
     # 用于存储所有文件的绝对路径
@@ -24,4 +24,20 @@ def resize_image(img, max_size=1024):
         new_w = int(w * scale)
         new_h = int(h * scale)
         return cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
-    return img
+    return img
+
+IMAGE_DIRS_PATH = 'data/image_dirs'
+if not os.path.exists(IMAGE_DIRS_PATH):
+    os.makedirs(IMAGE_DIRS_PATH)
+
+# 获取image_dirs下的所有文件夹
+def get_image_dirs():
+    dirs = ['data/images']  # 默认包含原始图片目录
+    if os.path.exists(IMAGE_DIRS_PATH):
+        for item in os.listdir(IMAGE_DIRS_PATH):
+            item_path = os.path.join(IMAGE_DIRS_PATH, item)
+            if os.path.isdir(item_path):
+                dirs.append(item_path)
+    for i in config.paths.image_dirs.keys():
+        dirs.append(i)
+    return dirs

services/image_search.py

@@ -152,23 +152,30 @@ class ImageSearch:
                             break
 
                     if full_filename:
-                        if filepath in generated_files:
-                            # 使用已经存在的embedding
-                            embedding = self.image_data[generated_files.index(filepath)]['embedding']
-                            embedding_name = self.image_data[generated_files.index(filepath)]['embedding_name']
-                        else:
-                            embedding_name = filename
-                            if replace_patterns_regex is not None:
-                                for pattern, replacement in replace_patterns_regex.items():
-                                    embedding_name = re.sub(pattern, replacement, embedding_name)
+                        # if filepath in generated_files:
+                        #     # 使用已经存在的embedding
+                        #     embedding = self.image_data[generated_files.index(filepath)]['embedding']
+                        #     embedding_name = self.image_data[generated_files.index(filepath)]['embedding_name']
+                        # else:
+
+                        # 在service那边已经有缓存了，这边直接开干，同时也是为了适配一个图片多个embedding。
+                        raw_embedding_name = filename
+                        if replace_patterns_regex is not None:
+                            for pattern, replacement in replace_patterns_regex.items():
+                                raw_embedding_name = re.sub(pattern, replacement, raw_embedding_name)
+                        embedding_names = raw_embedding_name.split('-')
+                        for embedding_name in embedding_names:
+                            if embedding_name == '':
+                                print()
+                                continue
                             embedding = self.embedding_service.get_embedding(embedding_name)
-                        embeddings.append({
-                            "filename": full_filename,
-                            "filepath": filepath,
-                            "embedding": embedding,
-                            "embedding_name": embedding_name,
-                            "type": image_type if image_type is not None else 'Normal'
-                        })
+                            embeddings.append({
+                                "filename": full_filename,
+                                "filepath": filepath,
+                                "embedding": embedding,
+                                "embedding_name": embedding_name,
+                                "type": image_type if image_type is not None else 'Normal'
+                            })
 
                     progress_bar.progress((index + 1) / length, text=f"处理图片 {index + 1}/{length}")
 
@@ -209,6 +216,7 @@ class ImageSearch:
             return []
         
         similarities = []
+        exists_imgs_path = []
         for img in self.image_data:
             if 'filepath' not in img and config.misc.adapt_for_old_version:
                 img['filepath'] = os.path.join(config.get_absolute_image_dirs()[0], img["filename"])
@@ -219,5 +227,14 @@ class ImageSearch:
             return []
             
         # 按相似度降序排序并返回前top_k个结果
-        sorted_items = sorted(similarities, key=lambda x: x[1], reverse=True)[:top_k]
-        return [item[0] for item in sorted_items] 
+        sorted_items = sorted(similarities, key=lambda x: x[1], reverse=True)
+        return_list = []
+        count = 0
+        for i in sorted_items:
+            if count >= top_k:
+                break
+            if i[0] not in exists_imgs_path:
+                return_list.append(i[0])
+                exists_imgs_path.append(i[0])
+                count += 1
+        return return_list

services/label_memes.py

@@ -9,9 +9,7 @@ import io
 PROMOTE = """你是一位表情包分类专家。请分析这个表情包，要求：
 
 1. 简要分析表情包的内容，含义，主体和可能的使用场景；如果表情包有文字，你应该仔细分析考虑，因为表情包的文字通常有幽默感而难以理解；
-2. 按格式要求输出表情包的文本描述，格式如下： **表情包含义**:(几个关键词概括表情包；如果表情包有文字，你要写出); **表情包主体**:(提取表情包的主角); **表情包使用场景**:(几个关键词描述表情包可能的使用场景)
-
-
+2. 按格式要求输出表情包的文本描述，格式如下： **表情包含义**:(几个关键词概括表情包；); **表情包主体**:(几个关键词描述表情包的主角); **表情包使用场景**:(几个关键词描述表情包可能的使用场景);**表情包文字**:(提取表情包中的所有文字；如果没有文字，输出"无文字")
 """
 
 class LabelMemes():
@@ -43,22 +41,32 @@ class LabelMemes():
 
     def _analyze_result_text(self, text:str):
         """分析并格式化模型返回的文本"""
-        if not '**表情包含义**' in text or not '**表情包主体**' in text or not '**表情包使用场景**' in text:
-            raise Exception(f'analyze result text error: {text}; 模型太蠢,换个模型或者重试')
+        if not '**表情包含义**' in text or not '**表情包主体**' in text or not '**表情包使用场景**' in text or not '**表情包文字**' in text:
+            raise Exception(f'analyze result text error: {text}: \n 模型太蠢（输出不符合要求）,换个模型或者重试')
         desc = text.split('**表情包含义**')[-1]
         character = desc.split('**表情包主体**')[-1]
         usage = character.split('**表情包使用场景**')[-1]
-        def clean_some_characters(x, l):
+        texts = character.split('**表情包文字**')[-1]
+
+        def clean_some_characters(x, l, r=''):
             for i in l:
-                x = x.replace(i, '')
+                x = x.replace(i, r)
             return x
         desc = desc.replace(character, '')
         character = character.replace(usage, '')
-        laji = ['表情包主体', '表情包使用场景', ':', '**(', ')；**', ');**', '**', ');', ')', '；', '(', ')', '\n', '：', '（', '）']
-        desc = clean_some_characters(desc, laji).replace('/', ' ').replace('\\', ' ')
-        character = clean_some_characters(character, laji).replace('/', ' ').replace('\\', ' ')
-        usage = clean_some_characters(usage, laji).replace('/', ' ').replace('\\', ' ')
-        return desc, character, usage
+        usage = usage.replace(texts, '')
+        laji = ['表情包主体', '表情包使用场景', '表情包文字', ':', '**(', ')；**', ');**', '**', ');', ')', '；', '(', ')', '\n', '：', '（', '）']
+        seperator = ['/', '\\', ',', '，', '、']
+        desc = clean_some_characters(clean_some_characters(desc, laji), seperator, ' ')
+        character = clean_some_characters(clean_some_characters(character, laji), seperator, ' ')
+        usage = clean_some_characters(clean_some_characters(usage, laji), seperator, ' ')
+        texts = clean_some_characters(clean_some_characters(texts, laji), seperator, ' ')
+        if '无文字' in texts:
+            texts = ''
+        for i in [desc, character, usage]:
+            if len(i) > 20:
+                raise Exception(f'analyze result text error: {text}: \n 模型太蠢（字数过多）， 换个模型或者重试')
+        return desc, character, usage, texts
 
     def _resize_image(self, img):
         """尺寸调整"""
@@ -147,7 +155,14 @@ class LabelMemes():
                         }
                     ]
                 }
-            ]
+            ],
+            "stream": False,
+            "max_tokens": 1024,
+            "stop": ["null"],
+            "temperature": 0.5,
+            "top_p": 0.5,
+            "top_k": 50,
+            "frequency_penalty": 0.0,
         }
         headers = {
             "Authorization": f"Bearer {config.api.silicon_api_key}",