diff --git a/.gitignore b/.gitignore index 6f2fce4..93d0353 100644 --- a/.gitignore +++ b/.gitignore @@ -5,4 +5,5 @@ .env .env.* !.env.example -config/config.yaml \ No newline at end of file +config/config.yaml +cache/ \ No newline at end of file diff --git a/README.md b/README.md index 6eae97e..7a66390 100644 --- a/README.md +++ b/README.md @@ -38,10 +38,26 @@ pic_example: type: "名称" ``` +## 上传图片 + +选择 `upload images` 页面，点击`添加表情包`按钮，上传图片即可。 + +可以选择多张表情包。 + +可以启用`使用VLM自动生成文件名`，程序会使用VLM为每张图片生成文件名。 + +每次上传完成后，需要叉掉所有已上传的文件或者刷新页面，否则会重复保存。 + +上传完成后，需要重新生成缓存。 + ## 图片打标 (WIP) 选择 `label images` 页面，选择图片文件夹，点击 `使用VLM生成描述` ，选择需要的描述后重命名文件即可。 +程序使用文件名生成embedding用于检索。如果你希望给一个图片多个embedding，你可以在文件名中使用`-`分隔。 + + + ## Demo 3bfb772e239f3437a13d46252aab1e1d diff --git a/app.py b/app.py index f053d0e..124c9c7 100644 --- a/app.py +++ b/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() \ No newline at end of file diff --git a/pages/label_images.py b/pages/label_images.py index 5d0a2e7..afc08d3 100644 --- a/pages/label_images.py +++ b/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: diff --git a/pages/upload_images.py b/pages/upload_images.py new file mode 100644 index 0000000..ca2bc00 --- /dev/null +++ b/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 = [] # 这行代码没用，那个组件仍然会保存已上传的文件 \ No newline at end of file diff --git a/pages/utils.py b/pages/utils.py index e2a4365..8b35d75 100644 --- a/pages/utils.py +++ b/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 \ No newline at end of file + 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 \ No newline at end of file diff --git a/services/image_search.py b/services/image_search.py index 6877e05..f413df7 100644 --- a/services/image_search.py +++ b/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] \ No newline at end of file + 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 \ No newline at end of file diff --git a/services/label_memes.py b/services/label_memes.py index 4d24744..0a41cd9 100644 --- a/services/label_memes.py +++ b/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}",