add batch upload

app/database.py

@@ -129,7 +129,6 @@ class Database:
             )
             
             exists = len(results) > 0
-            logger.info(f"Check file exists - hash: {image_hash}, exists: {exists}, results: {results}")
             return exists
         except Exception as e:
             logger.error(f"Failed to check file exists: {str(e)}")
@@ -149,6 +148,68 @@ class Database:
         except Exception as e:
             logger.error(f"Failed to delete sticker: {str(e)}")
             return f"Failed to delete sticker: {str(e)}"
+
+    def batch_store_stickers(self, stickers: List[Dict[str, Any]], batch_size: int = 100) -> bool:
+        """批量存储贴纸数据到Milvus
+        
+        Args:
+            stickers (List[Dict[str, Any]]): 贴纸数据列表，每个元素包含以下字段：
+                - title: str
+                - description: str
+                - tags: Union[str, List[str]]
+                - file_path: str
+                - image_hash: str (可选)
+            batch_size (int, optional): 每批处理的数量. Defaults to 100.
+            
+        Returns:
+            bool: 是否全部插入成功
+        """
+        try:
+            total_stickers = len(stickers)
+            if total_stickers == 0:
+                logger.warning("No stickers to store")
+                return True
+                
+            logger.info(f"Starting batch store of {total_stickers} stickers")
+            
+            # 分批处理
+            for i in range(0, total_stickers, batch_size):
+                batch = stickers[i:i + batch_size]
+                batch_data = []
+                
+                for sticker in batch:
+                    # 处理标签格式
+                    tags = sticker.get("tags", [])
+                    if isinstance(tags, str):
+                        tags = tags.split(",")
+                    
+                    # 编码描述文本
+                    vector = self.encode_text(sticker.get("description", ""))
+                    
+                    batch_data.append({
+                        "vector": vector,
+                        "title": sticker.get("title", ""),
+                        "description": sticker.get("description", ""),
+                        "tags": tags,
+                        "file_name": sticker.get("file_path", ""),
+                        "image_hash": sticker.get("image_hash")
+                    })
+                
+                # 批量插入
+                if batch_data:
+                    self.client.insert(
+                        collection_name=self.collection_name,
+                        data=batch_data
+                    )
+                    logger.info(f"Batch {i//batch_size + 1} stored successfully - {len(batch_data)} stickers")
+            
+            logger.info("All stickers stored successfully ✅")
+            return True
+            
+        except Exception as e:
+            logger.error(f"Failed to batch store stickers: {str(e)}")
+            return False
+
 # 初始化 Milvus 数据库
 
 # 创建数据库实例

app/image_utils.py

@@ -71,6 +71,19 @@ def upload_to_huggingface(temp_file_path: str) -> tuple:
     logger.info(f"Image uploaded successfully: {file_path}")
     return file_path, image_filename
 
+def upload_folder_to_huggingface(folder_path: str) -> None:
+    """上传目录到 HuggingFace"""
+    logger.info(f"Uploading folder to HuggingFace: {folder_path}")
+    api.upload_folder(
+        folder_path=folder_path,
+        path_in_repo="images/",
+        repo_id=DATASET_ID,
+        token=HUGGING_FACE_TOKEN,
+        repo_type="dataset"
+    )
+    logger.info(f"Image uploaded successfully: {folder_path}")
+    return 
+
 
 def get_image_cdn_url(file_path: str) -> str:
     """获取图片CDN URL"""
@@ -104,4 +117,41 @@ def format_image_url(file_path: str) -> str:
         return os.path.abspath(file_path)
     
     # 如果是HuggingFace路径，返回完整的URL
-    return f"https://huggingface.co/datasets/{DATASET_ID}/resolve/main/{file_path}"
+    return f"https://huggingface.co/datasets/{DATASET_ID}/resolve/main/{file_path}"
+
+def generate_temp_image(temp_dir: str, image: Image.Image, image_filename: str) -> str:
+    """根据图片和文件名生成临时图片文件
+    
+    Args:
+        temp_dir (str): 临时目录路径
+        image (Image.Image): PIL图片对象
+        image_filename (str): 图片文件名，格式为 image_XXXXXX.png
+        
+    Returns:
+        str: 临时图片文件的路径
+        
+    Raises:
+        ValueError: 如果临时目录不存在或无法访问
+        IOError: 如果保存图片失败
+    """
+    try:
+        # 确保临时目录存在
+        if not os.path.exists(temp_dir):
+            os.makedirs(temp_dir)
+            logger.info(f"Created temporary directory: {temp_dir}")
+        
+        # 构建临时文件路径
+        temp_file_path = os.path.join(temp_dir, image_filename)
+        
+        # 保存图片到临时文件
+        image.save(temp_file_path, "PNG")
+        logger.info(f"Generated temporary image: {temp_file_path}")
+        
+        return temp_file_path
+        
+    except OSError as e:
+        logger.error(f"Failed to create temporary directory: {str(e)}")
+        raise ValueError(f"Failed to create temporary directory: {str(e)}")
+    except Exception as e:
+        logger.error(f"Failed to generate temporary image: {str(e)}")
+        raise IOError(f"Failed to generate temporary image: {str(e)}")

app/services.py

@@ -1,5 +1,6 @@
 import os
 import logging
+import random
 import zipfile
 import tempfile
 import shutil
@@ -16,6 +17,8 @@ from app.config import (
 from app.database import db
 from app.image_utils import (
     save_image_temp, 
+    generate_temp_image,
+    upload_folder_to_huggingface,
     upload_to_huggingface, 
     get_image_cdn_url, 
     get_image_description,
@@ -78,35 +81,31 @@ class StickerService:
     def import_stickers(
         sticker_dataset: str,
         upload: bool = False,
-        gen_description: bool = False,
         progress_callback: callable = None,
-        total_files: int = 0
     ) -> List[str]:
         """导入表情包数据集
         
         Args:
             sticker_dataset (str): 表情包数据集路径
             upload (bool, optional): 是否上传到HuggingFace. Defaults to False.
-            gen_description (bool, optional): 是否生成AI描述. Defaults to False.
             progress_callback (callable, optional): 进度回调函数. Defaults to None.
-            total_files (int, optional): 总文件数. Defaults to 0.
         """
         results = []
         descriptions = {}
         
         try:
             # 创建临时目录
-            temp_dir = TEMP_DIR
-            logger.info(f"Created temporary directory: {temp_dir}")
-            
+            cache_folder = os.path.join(TEMP_DIR, 'cache/')
+            img_folder = os.path.join(TEMP_DIR, 'images/')
+            data_json_path = os.path.join(cache_folder, 'data.json')
+
+            logger.info(f"start import dataset")
             # 解压数据集
             with zipfile.ZipFile(sticker_dataset, 'r') as zip_ref:
-                zip_ref.extractall(temp_dir)
-            
-            logger.info(f"Extracted dataset to: {temp_dir}")
+                zip_ref.extractall(cache_folder)
+            logger.info(f"Extracted dataset to: {cache_folder}")
             
             # 尝试读取data.json文件
-            data_json_path = os.path.join(temp_dir, 'data.json')
             if os.path.exists(data_json_path):
                 with open(data_json_path, 'r', encoding='utf-8') as f:
                     data = json.load(f)
@@ -114,7 +113,7 @@ class StickerService:
                 logger.info(f"Loaded descriptions from data.json")
             
             # 遍历解压后的目录
-            for root, dirs, files in os.walk(temp_dir):
+            for root, dirs, files in os.walk(cache_folder):
                 for file in files:
                     if file.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.webp')):
                         image_path = os.path.join(root, file)
@@ -122,8 +121,6 @@ class StickerService:
                         try:
                             # 打开图片
                             image = Image.open(image_path)
-                            
-                            # 检查文件名是否已存在
                             image_hash = calculate_image_hash(image)
                             if db.check_image_exists(image_hash):
                                 results.append(f"跳过已存在的图片: {file}")
@@ -131,40 +128,23 @@ class StickerService:
                                     progress_callback(file, "Skipped (exists)")
                                 continue
                             
-                            # 保存到临时文件
-                            temp_file_path = save_image_temp(image)
-                            
-                            # 上传到 HuggingFace
-                            if upload:
-                                file_path, image_filename = upload_to_huggingface(temp_file_path)
-                            else:
-                                file_path = temp_file_path
-                                image_url = file_path
-                            
                             # 获取图片描述
                             description = None
                             if file in descriptions:
                                 description = descriptions[file]
-                            elif gen_description:
-                                image_cdn_url = ''
-                                if (PUBLIC_URL):
-                                    image_cdn_url = f'{PUBLIC_URL}/gradio_api/file={temp_file_path}'
-                                else:
-                                    image_cdn_url = get_image_cdn_url(file_path)
-                                description = get_image_description(image_cdn_url)
-                            
+
                             if not description:
-                                logger.warning(f"No description available for {file}")
+                                results.append(f"跳过无描述的图片: {file}")
                                 if progress_callback:
                                     progress_callback(file, "Skipped (no description)")
                                 continue
                             
-                            # 清理临时文件
-                            if upload:
-                                os.unlink(temp_file_path)
-                            
+                            image_filename = f"image_{random.randint(100000, 999999)}.png"
+                            file_path = f"images/{image_filename}"
+                            generate_temp_image(img_folder, image, image_filename)
+
                             db.store_sticker("", description, "", file_path, image_hash)
-                            results.append(f"成功导入: {file}")
+                            results.append(f"成功导入: {image_filename}")
                             
                             if progress_callback:
                                 progress_callback(file, "Imported")
@@ -174,7 +154,11 @@ class StickerService:
                             results.append(f"处理失败 {file}: {str(e)}")
                             if progress_callback:
                                 progress_callback(file, f"Failed: {str(e)}")
-            
+
+            # 上传到 HuggingFace
+            if upload:
+                upload_folder_to_huggingface(img_folder)
+
             return results
             
         except Exception as e:
@@ -184,9 +168,13 @@ class StickerService:
             
         finally:
             # 清理临时目录
-            if temp_dir and os.path.exists(temp_dir):
-                shutil.rmtree(temp_dir)
-                logger.info(f"Cleaned up temporary directory: {temp_dir}")
+            if cache_folder and os.path.exists(cache_folder):
+                shutil.rmtree(cache_folder)
+                logger.info(f"Cleaned up temporary directory: {cache_folder}")
+
+            if img_folder and os.path.exists(img_folder):
+                shutil.rmtree(img_folder)
+                logger.info(f"Cleaned up temporary directory: {img_folder}")
 
     @staticmethod
     def search_stickers(description: str, limit: int = 2, reranking : bool = False) -> List[Dict[str, Any]]:

app/ui.py

@@ -144,7 +144,6 @@ class StickerUI:
             
             with gr.Row():
                 self.upload_checkbox = gr.Checkbox(label="Upload to HuggingFace", value=False)
-                self.gen_desc_checkbox = gr.Checkbox(label="Generate AI Descriptions", value=False)
             
             with gr.Row():
                 self.import_button.render()
@@ -155,13 +154,12 @@ class StickerUI:
                 fn=self._import_stickers_with_progress,
                 inputs=[
                     self.dataset_input,
-                    self.upload_checkbox,
-                    self.gen_desc_checkbox
+                    self.upload_checkbox
                 ],
                 outputs=self.import_output
             )
     
-    def _import_stickers_with_progress(self, dataset_path, upload, gen_description, progress=gr.Progress()):
+    def _import_stickers_with_progress(self, dataset_path, upload, progress=gr.Progress()):
         """Import stickers with progress tracking."""
         try:
             # Count total files first
@@ -189,7 +187,6 @@ class StickerUI:
             results = sticker_service.import_stickers(
                 dataset_path,
                 upload=upload,
-                gen_description=gen_description,
                 progress_callback=update_progress,
                 total_files=total_files
             )

main.py

@@ -72,7 +72,7 @@ async def api_delete_stickers(request: dict):
 
 
 @app.post("/api/import_dataset")
-async def api_import_dataset(file: UploadFile = File(...), upload: bool = False, gen_description: bool = False):
+async def api_import_dataset(file: UploadFile = File(...), upload: bool = False):
     """Import sticker dataset from ZIP file"""
     try:
         # Create a temporary file to store the uploaded ZIP
@@ -82,7 +82,7 @@ async def api_import_dataset(file: UploadFile = File(...), upload: bool = False,
             temp_file_path = temp_file.name
         
         # Import the dataset
-        results = sticker_service.import_stickers(temp_file_path, upload, gen_description)
+        results = sticker_service.import_stickers(temp_file_path, upload)
         
         # Clean up the temporary file
         os.unlink(temp_file_path)