pensiev 0.25.5__py3-none-any.whl

memos/plugins/ocr/server.py

@@ -0,0 +1,227 @@
+from PIL import Image
+import numpy as np
+import logging
+from fastapi import FastAPI, Body, HTTPException
+from contextlib import asynccontextmanager
+import base64
+import io
+import asyncio
+from pydantic import BaseModel, Field
+from typing import List
+from multiprocessing import Pool
+import threading
+import time
+import uvicorn
+import os
+
+# Configure logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# 创建进程池
+process_pool = None
+
+# 从环境变量中读取参数
+max_workers = int(os.getenv("MAX_WORKERS", 1))
+
+
+def str_to_bool(value):
+    return value.lower() in ("true", "1", "t", "y", "yes")
+
+
+use_gpu = str_to_bool(os.getenv("USE_GPU", "false"))
+
+
+@asynccontextmanager
+async def lifespan(app):
+    global process_pool
+    if process_pool is None:
+        init_process_pool(max_workers=max_workers, use_gpu=use_gpu)
+    yield
+    if process_pool:
+        logger.info("Shutting down process pool...")
+        process_pool.close()
+        process_pool.join()
+        logger.info("Process pool shut down.")
+
+
+app = FastAPI(lifespan=lifespan)
+
+
+def init_worker(use_gpu):
+    global ocr
+    ocr = init_ocr(use_gpu)
+
+
+def init_process_pool(max_workers, use_gpu):
+    global process_pool
+    process_pool = Pool(
+        processes=max_workers, initializer=init_worker, initargs=(use_gpu,)
+    )
+
+
+def init_ocr(use_gpu):
+    if use_gpu:
+        try:
+            from rapidocr_paddle import RapidOCR as RapidOCRPaddle
+
+            ocr = RapidOCRPaddle(
+                det_use_cuda=True, cls_use_cuda=True, rec_use_cuda=True
+            )
+            logger.info("Initialized OCR with RapidOCR Paddle (GPU)")
+        except ImportError:
+            logger.error(
+                "Failed to import rapidocr_paddle. Make sure it's installed for GPU usage."
+            )
+            raise
+    else:
+        try:
+            from rapidocr_onnxruntime import RapidOCR
+
+            ocr = RapidOCR()
+            logger.info("Initialized OCR with RapidOCR ONNX Runtime (CPU)")
+        except ImportError:
+            logger.error(
+                "Failed to import rapidocr_onnxruntime. Make sure it's installed for CPU usage."
+            )
+            raise
+    return ocr
+
+
+def convert_ocr_results(results):
+    if results is None:
+        return []
+
+    converted = []
+    for result in results:
+        item = {"dt_boxes": result[0], "rec_txt": result[1], "score": result[2]}
+        converted.append(item)
+    return converted
+
+
+def predict(image_data):
+    global ocr
+    if ocr is None:
+        raise ValueError("OCR engine not initialized")
+
+    image = Image.open(io.BytesIO(image_data))
+    img_array = np.array(image)
+    results, _ = ocr(img_array)
+    converted_results = convert_ocr_results(results)
+    return converted_results
+
+
+def convert_to_python_type(item):
+    if isinstance(item, np.ndarray):
+        return item.tolist()
+    elif isinstance(item, np.generic):  # This includes numpy scalars like numpy.float32
+        return item.item()
+    elif isinstance(item, list):
+        return [convert_to_python_type(sub_item) for sub_item in item]
+    elif isinstance(item, dict):
+        return {key: convert_to_python_type(value) for key, value in item.items()}
+    else:
+        return item
+
+
+async def async_predict(image_data):
+    loop = asyncio.get_running_loop()
+    results = await loop.run_in_executor(
+        None, process_pool.apply, predict, (image_data,)
+    )
+    return results
+
+
+class OCRResult(BaseModel):
+    dt_boxes: List[List[float]] = Field(..., description="Bounding box coordinates")
+    rec_txt: str = Field(..., description="Recognized text")
+    score: float = Field(..., description="Confidence score")
+
+
+@app.post("/predict", response_model=List[OCRResult])
+async def predict_base64(image_base64: str = Body(..., embed=True)):
+    try:
+        if not image_base64:
+            raise HTTPException(status_code=400, detail="Missing image_base64 field")
+
+        # Remove header part if present
+        if image_base64.startswith("data:image"):
+            image_base64 = image_base64.split(",")[1]
+
+        # Decode the base64 image
+        image_data = base64.b64decode(image_base64)
+
+        # 直接传递图像数据给async_predict
+        ocr_result = await async_predict(image_data)
+
+        return convert_to_python_type(ocr_result)
+
+    except Exception as e:
+        logging.error(f"Error during OCR processing: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+shutdown_event = threading.Event()
+
+
+def signal_handler(signum, frame):
+    logger.info("Received interrupt signal. Initiating shutdown...")
+    shutdown_event.set()
+
+
+def run_server(app, host, port):
+    config = uvicorn.Config(app, host=host, port=port, loop="asyncio")
+    server = uvicorn.Server(config)
+    server.install_signal_handlers = (
+        lambda: None
+    )  # Disable Uvicorn's own signal handlers
+
+    async def serve():
+        await server.serve()
+
+    thread = threading.Thread(target=asyncio.run, args=(serve(),))
+    thread.start()
+
+    try:
+        while not shutdown_event.is_set():
+            time.sleep(1)
+    except KeyboardInterrupt:
+        logger.info("Keyboard interrupt received. Initiating shutdown...")
+    finally:
+        shutdown_event.set()
+        logger.info("Stopping the server...")
+        asyncio.run(server.shutdown())
+        thread.join()
+        logger.info("Server stopped.")
+
+
+if __name__ == "__main__":
+    import uvicorn
+    import argparse
+
+    parser = argparse.ArgumentParser(description="OCR Service")
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=8000,
+        help="Port to run the OCR service on",
+    )
+    parser.add_argument(
+        "--max-workers",
+        type=int,
+        default=1,
+        help="Maximum number of worker threads for OCR processing",
+    )
+    parser.add_argument(
+        "--gpu",
+        action="store_true",
+        help="Use GPU for OCR processing",
+    )
+
+    args = parser.parse_args()
+    port = args.port
+    max_workers = args.max_workers
+    use_gpu = args.gpu
+
+    run_server(app, "0.0.0.0", port)
+

memos/plugins/ocr/temp_ppocr.yaml

@@ -0,0 +1,42 @@
+Cls:
+  cls_batch_num: 6
+  cls_image_shape:
+  - 3
+  - 48
+  - 192
+  cls_thresh: 0.9
+  label_list:
+  - '0'
+  - '180'
+  model_path: /Users/shanchuanxu/projects/memos/memos/plugins/ocr/models/ch_ppocr_mobile_v2.0_cls_train.onnx
+  use_cuda: false
+Det:
+  box_thresh: 0.3
+  limit_side_len: 1500
+  limit_type: min
+  max_candidates: 1000
+  model_path: /Users/shanchuanxu/projects/memos/memos/plugins/ocr/models/ch_PP-OCRv4_det_infer.onnx
+  score_mode: fast
+  thresh: 0.3
+  unclip_ratio: 1.6
+  use_cuda: false
+  use_dilation: true
+Global:
+  max_side_len: 2000
+  min_height: 30
+  min_side_len: 30
+  print_verbose: false
+  text_score: 0.5
+  use_cls: true
+  use_det: true
+  use_rec: true
+  use_space_char: true
+  width_height_ratio: 40
+Rec:
+  model_path: /Users/shanchuanxu/projects/memos/memos/plugins/ocr/models/ch_PP-OCRv4_rec_infer.onnx
+  rec_batch_num: 6
+  rec_img_shape:
+  - 3
+  - 48
+  - 320
+  use_cuda: false

memos/plugins/vlm/main.py

@@ -0,0 +1,251 @@
+import base64
+import httpx
+from PIL import Image
+import asyncio
+from typing import Optional
+from fastapi import APIRouter, FastAPI, Request, HTTPException
+from memos.schemas import Entity, MetadataType
+import logging
+import uvicorn
+import os
+import io
+import numpy as np
+
+
+# Configure logger
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+PLUGIN_NAME = "vlm"
+
+router = APIRouter(tags=[PLUGIN_NAME], responses={404: {"description": "Not found"}})
+
+modelname = None
+endpoint = None
+token = None
+concurrency = None
+semaphore = None
+force_jpeg = None
+prompt = None
+
+
+def get_metadata_name() -> str:
+    """Return the metadata field name used by this plugin."""
+    global modelname
+    return f"{modelname.replace('-', '_')}_result"
+
+
+def image2base64(img_path):
+    try:
+        with Image.open(img_path) as img:
+            img.verify()  # Verify the image file
+
+        with Image.open(img_path) as img:
+            # Check image size and skip if it's too small
+            if img.width < 10 or img.height < 10:
+                logger.warning(f"Image is too small: {img.width}x{img.height}. Skipping processing.")
+                return None
+
+            # Convert image to RGB mode (removes alpha channel if present)
+            img = img.convert("RGB")
+            
+            # Convert to numpy array and check shape
+            img_array = np.array(img)
+            logger.info(f"Image shape: {img_array.shape}")
+
+            if img_array.shape[2] != 3:
+                logger.warning(f"Unexpected number of channels: {img_array.shape[2]}. Expected 3. Skipping processing.")
+                return None
+            
+            if force_jpeg:
+                # Save as JPEG in memory
+                buffer = io.BytesIO()
+                img.save(buffer, format="JPEG")
+                buffer.seek(0)
+                encoded_string = base64.b64encode(buffer.getvalue()).decode("utf-8")
+            else:
+                # Use original format, but ensure it's RGB
+                buffer = io.BytesIO()
+                img.save(buffer, format=img.format or "JPEG")
+                buffer.seek(0)
+                encoded_string = base64.b64encode(buffer.getvalue()).decode("utf-8")
+        return encoded_string
+    except Exception as e:
+        logger.error(f"Error processing image {img_path}: {str(e)}")
+        return None
+
+
+async def fetch(endpoint: str, client, request_data, headers: Optional[dict] = None):
+    async with semaphore:
+        try:
+            response = await client.post(
+                f"{endpoint}/v1/chat/completions",
+                json=request_data,
+                timeout=60,
+                headers=headers,
+            )
+            response.raise_for_status()
+            result = response.json()
+            choices = result.get("choices", [])
+            if (
+                choices
+                and "message" in choices[0]
+                and "content" in choices[0]["message"]
+            ):
+                return choices[0]["message"]["content"]
+            return ""
+        except Exception as e:
+            logger.error(f"Exception occurred: {str(e)}")
+            return None
+
+
+async def predict(
+    endpoint: str, modelname: str, img_path: str, token: Optional[str] = None
+) -> Optional[str]:
+    return await predict_remote(endpoint, modelname, img_path, token)
+
+
+async def predict_remote(
+    endpoint: str, modelname: str, img_path: str, token: Optional[str] = None
+) -> Optional[str]:
+    img_base64 = image2base64(img_path)
+    if not img_base64:
+        logger.warning(f"Skipping processing for file: {img_path} due to invalid or small image")
+        return None
+
+    mime_type = (
+        "image/jpeg" if force_jpeg else "image/jpeg"
+    )  # Default to JPEG if force_jpeg is True
+
+    if not force_jpeg:
+        # Only determine MIME type if not forcing JPEG
+        _, file_extension = os.path.splitext(img_path)
+        file_extension = file_extension.lower()[1:]
+        mime_types = {
+            "png": "image/png",
+            "jpg": "image/jpeg",
+            "jpeg": "image/jpeg",
+            "webp": "image/webp",
+        }
+        mime_type = mime_types.get(file_extension, "image/jpeg")
+
+    request_data = {
+        "model": modelname,
+        "messages": [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:{mime_type};base64,{img_base64}"},
+                    },
+                    {"type": "text", "text": prompt},  # Use the global prompt variable here
+                ],
+            }
+        ],
+        "stream": False,
+        "max_tokens": 1024,
+        "temperature": 0.1,
+        "repetition_penalty": 1.1,
+        "top_p": 0.8,
+    }
+    async with httpx.AsyncClient() as client:
+        headers = {}
+        if token:
+            headers["Authorization"] = f"Bearer {token.get_secret_value()}"
+        return await fetch(endpoint, client, request_data, headers=headers)
+
+
+@router.get("/")
+async def read_root():
+    return {"healthy": True}
+
+
+@router.post("", include_in_schema=False)
+@router.post("/")
+async def vlm(entity: Entity, request: Request):
+    global modelname, endpoint, token
+    metadata_field_name = get_metadata_name()
+    if not entity.file_type_group == "image":
+        return {metadata_field_name: ""}
+
+    # Check if the METADATA_FIELD_NAME field is empty or null
+    existing_metadata = entity.get_metadata_by_key(metadata_field_name)
+    if (
+        existing_metadata
+        and existing_metadata.value
+        and existing_metadata.value.strip()
+    ):
+        logger.info(
+            f"Skipping processing for file: {entity.filepath} due to existing metadata"
+        )
+        # If the field is not empty, return without processing
+        return {metadata_field_name: existing_metadata.value}
+
+    # Check if the entity contains the tag "low_info"
+    if any(tag.name == "low_info" for tag in entity.tags):
+        # If the tag is present, return without processing
+        logger.info(
+            f"Skipping processing for file: {entity.filepath} due to 'low_info' tag"
+        )
+        return {metadata_field_name: ""}
+
+    location_url = request.headers.get("Location")
+    if not location_url:
+        raise HTTPException(status_code=400, detail="Location header is missing")
+
+    patch_url = f"{location_url}/metadata"
+
+    vlm_result = await predict(endpoint, modelname, entity.filepath, token=token)
+
+    logger.info(f"VLM result: {vlm_result[:100]}...")
+    if not vlm_result:
+        logger.info(f"No VLM result found for file: {entity.filepath}")
+        return {metadata_field_name: "{}"}
+
+    async with httpx.AsyncClient() as client:
+        response = await client.patch(
+            patch_url,
+            json={
+                "metadata_entries": [
+                    {
+                        "key": metadata_field_name,
+                        "value": vlm_result,
+                        "source": PLUGIN_NAME,
+                        "data_type": MetadataType.TEXT_DATA.value,
+                    }
+                ]
+            },
+            timeout=30,
+        )
+
+    if response.status_code != 200:
+        raise HTTPException(
+            status_code=response.status_code, detail="Failed to patch entity metadata"
+        )
+
+    return {
+        metadata_field_name: vlm_result,
+    }
+
+
+def init_plugin(config):
+    global modelname, endpoint, token, concurrency, semaphore, force_jpeg, prompt
+
+    modelname = config.modelname
+    endpoint = config.endpoint
+    token = config.token
+    concurrency = config.concurrency
+    force_jpeg = config.force_jpeg
+    prompt = config.prompt
+    semaphore = asyncio.Semaphore(concurrency)
+
+    # Print the parameters
+    logger.info("VLM plugin initialized")
+    logger.info(f"Model Name: {modelname}")
+    logger.info(f"Endpoint: {endpoint}")
+    logger.info(f"Token: {token}")
+    logger.info(f"Concurrency: {concurrency}")
+    logger.info(f"Force JPEG: {force_jpeg}")
+    logger.info(f"Prompt: {prompt}")
+

memos/prepare_dataset.py

@@ -0,0 +1,107 @@
+import json
+import os
+import shutil
+from tqdm import tqdm
+from PIL import Image
+import io
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+
+def compress_and_convert_image(image_path, output_path, max_size_mb=2):
+    with Image.open(image_path) as img:
+        if (
+            image_path.lower().endswith(".png")
+            and os.path.getsize(image_path) > max_size_mb * 1024 * 1024
+        ):
+            # 将 PNG 转换为 JPG 并压缩
+            img = img.convert("RGB")
+            output_path = output_path.rsplit(".", 1)[0] + ".jpg"
+            img.save(output_path, "JPEG", quality=85)
+        else:
+            # 直接复制原图
+            img.save(output_path)
+    return output_path
+
+
+def process_image(data, shots_dir):
+    image_path = data.get("image")
+    if image_path and os.path.exists(image_path):
+        image_name = os.path.basename(image_path)
+        new_image_path = os.path.join(shots_dir, image_name)
+        if new_image_path == image_path:
+            return None
+        try:
+            new_image_path = compress_and_convert_image(image_path, new_image_path)
+            data["image"] = os.path.relpath(new_image_path, shots_dir)
+            return data, os.path.basename(new_image_path)
+        except Exception as e:
+            print(f"Error processing image {image_path}: {e}")
+    elif not image_path:
+        return data, None
+    return None
+
+
+def update_image_paths(input_file, output_file, shots_dir, max_workers=8):
+    if not os.path.exists(shots_dir):
+        os.makedirs(shots_dir)
+
+    updated_data = []
+    seen_images = set()
+
+    with open(input_file, "r") as infile:
+        data_list = [json.loads(line) for line in infile]
+
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        futures = [
+            executor.submit(process_image, data, shots_dir) for data in data_list
+        ]
+
+        for future in tqdm(
+            as_completed(futures), total=len(futures), desc="Processing images"
+        ):
+            result = future.result()
+            if result:
+                data, image_name = result
+                if image_name and image_name not in seen_images:
+                    seen_images.add(image_name)
+                    updated_data.append(data)
+                elif not image_name:
+                    updated_data.append(data)
+
+    with open(output_file, "w") as outfile:
+        for item in tqdm(updated_data, desc="Writing updated data", unit="item"):
+            json.dump(item, outfile, ensure_ascii=False)
+            outfile.write("\n")
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Update image paths in the dataset.")
+    parser.add_argument(
+        "--input_file",
+        type=str,
+        default="dataset.updated.jsonl",
+        help="Path to the input JSONL file.",
+    )
+    parser.add_argument(
+        "--output_file",
+        type=str,
+        default="dataset.updated.formatted.jsonl",
+        help="Path to the output JSONL file.",
+    )
+    parser.add_argument(
+        "--shots_dir",
+        type=str,
+        default="shots",
+        help="Directory to save processed images.",
+    )
+    parser.add_argument(
+        "--max_workers", type=int, default=8, help="Maximum number of worker threads."
+    )
+
+    args = parser.parse_args()
+
+    update_image_paths(
+        args.input_file, args.output_file, args.shots_dir, args.max_workers
+    )

memos/process_webp.py

@@ -0,0 +1,55 @@
+from pathlib import Path
+from PIL import Image
+import piexif
+import json
+from memos.utils import write_image_metadata, get_image_metadata
+from tqdm import tqdm
+
+def convert_webp_metadata(directory):
+    webp_files = list(Path(directory).glob('**/*.webp'))
+
+    for webp_file in tqdm(webp_files, desc="Converting WebP metadata", unit="file"):
+        try:
+            # Try to get metadata using the new method
+            new_metadata = get_image_metadata(webp_file)
+
+            if new_metadata:
+                tqdm.write(f"Skipping {webp_file}: Already in new format")
+                continue
+
+            # If new method fails, try to get old metadata
+            img = Image.open(webp_file)
+            old_metadata = img.info.get("exif", None)
+
+            if old_metadata is None:
+                tqdm.write(f"Skipping {webp_file}: No metadata found")
+                continue
+
+            if isinstance(old_metadata, bytes):
+                try:
+                    old_metadata = old_metadata.decode('utf-8')
+                except UnicodeDecodeError:
+                    tqdm.write(f"Skipping {webp_file}: Unable to decode metadata")
+                    continue
+
+            try:
+                metadata = json.loads(old_metadata)
+            except json.JSONDecodeError:
+                tqdm.write(f"Skipping {webp_file}: Invalid metadata format")
+                continue
+
+            # Convert to new format
+            write_image_metadata(webp_file, metadata)
+            tqdm.write(f"Converted metadata for {webp_file}")
+
+        except Exception as e:
+            tqdm.write(f"Error processing {webp_file}: {str(e)}")
+
+if __name__ == "__main__":
+    import sys
+    if len(sys.argv) != 2:
+        print("Usage: python convert_webp_metadata.py <directory>")
+        sys.exit(1)
+
+    directory = sys.argv[1]
+    convert_webp_metadata(directory)