LongArticleTaskServer/app/infra/shared/image.py

import requests
import imagehash
from PIL import Image
from io import BytesIO
from typing import Optional

from app.core.config import GlobalConfigSettings
from app.infra.shared.oss import OssUtils
from app.schemas import ImagePath


class ImageUtils(OssUtils):
    """phash 汉明距离 0~64，越小越相似。低于此阈值视为同一张图。"""

    SAME_IMAGE_PHASH_THRESHOLD = 5

    def __init__(self, config: GlobalConfigSettings):
        super().__init__(config.aliyun_oss)

    def load_image(self, path: str, path_type: Optional[str] = None):
        path_type = path_type or "oss_file"
        match path_type:
            case "filepath":
                img = Image.open(path)

            case "url":
                img = Image.open(BytesIO(requests.get(path, timeout=5).content))

            case "oss_file":
                img = Image.open(BytesIO(self.bucket.get_object(path).read()))

            case _:
                return "file_type error"

        return img.convert("RGB")

    @staticmethod
    def phash_distance(img1: Image.Image, img2: Image.Image):
        h1 = imagehash.phash(img1)
        h2 = imagehash.phash(img2)
        return h1 - h2

    def image_similar(self, file1: ImagePath, file2: ImagePath):
        img1 = self.load_image(
            path=file1.path,
            path_type=file1.path_type,
        )
        img2 = self.load_image(
            path=file2.path,
            path_type=file2.path_type,
        )

        return self.phash_distance(img1, img2)

    def is_same_image(self, file1: ImagePath, file2: ImagePath) -> bool:
        """判断两张图是否视为同一张（phash 距离 <= SAME_IMAGE_PHASH_THRESHOLD）。"""
        return self.image_similar(file1, file2) <= self.SAME_IMAGE_PHASH_THRESHOLD