hangzhou_mat_manage/imgsearchimg/func/image_search.py

import numpy as np
from PIL import Image
from typing import  Optional
import torch
import torchvision.transforms as transforms
import torchvision.models as models
from torchvision.models import ResNet50_Weights

import torch.nn.functional as F
import time

class ImageSearchEngine:

    def __init__(self):
        # 检查GPU是否可用（仅用于PyTorch模型）
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"使用设备: {self.device}")
        
        # 定义基础预处理转换
        self.base_transform = transforms.Compose([
            transforms.Grayscale(num_output_channels=3),  # 转换为灰度图但保持3通道
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])
        
    
        # 加载预训练的ResNet模型
        self.model = models.resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)
        # 移除最后的全连接层
        self.model = torch.nn.Sequential(*list(self.model.children())[:-1])
        self.model = self.model.to(self.device)
        self.model.eval()
        
        # 初始化FAISS索引（2048是ResNet50的特征维度）
        self.dimension = 2048
        # self.index = faiss.IndexFlatL2(self.dimension)

        # 改为支持删除的索引
        # base_index = faiss.IndexFlatL2(self.dimension)
        # self.index = faiss.IndexIDMap(base_index) 
        

    def _process_image(self, image_path: str) -> Optional[torch.Tensor]:
        """处理单张图片并提取特征。
        
        Args:
            image_path: 图片路径
            
        Returns:
            处理后的特征向量，如果处理失败返回None
        """
        try:
            # 读取图片
            image = Image.open(image_path)
            
            # 确保图片是RGB模式
            if image.mode != 'RGB':
                image = image.convert('RGB')
            

            start_ms_time = time.time()
            # 提取多尺度特征
            multi_scale_features = self._extract_multi_scale_features(image)
            end_ms_time = time.time()
            print(f"提取多尺度特征耗时: { end_ms_time - start_ms_time } s",)
            if multi_scale_features is None:
                return None
            
            start_sw_time = time.time()
            # 提取滑动窗口特征
            sliding_window_features = self._extract_sliding_window_features(image)
            end_sw_time = time.time()
            print(f"提取滑动窗口耗时: { end_sw_time - start_sw_time } s",)
            if sliding_window_features is None:
                return None
            
            # 组合特征（加权平均）
            combined_feature = multi_scale_features * 0.6 + sliding_window_features * 0.4
            
            # 标准化特征
            combined_feature = F.normalize(combined_feature, p=2, dim=0)
            
            return combined_feature
            
        except Exception as e:
            print(f"处理图片时出错: {e}")
            return None

    def _extract_multi_scale_features(self, image: Image.Image) -> Optional[torch.Tensor]:
        """基于原图分辨率的多尺度特征提取（智能动态调整版）
        
        Args:
            image: PIL图片对象
            
        Returns:
            多尺度特征向量，处理失败返回None
        """
        try:
            # 获取原图信息
            orig_w, orig_h = image.size
            max_edge = max(orig_w, orig_h)
            aspect_ratio = orig_w / orig_h

            # 动态调整策略 -------------------------------------------
            # 策略1：根据最大边长确定基准尺寸
            base_size = min(max_edge, 3000)  # 不超过模型支持的最大尺寸
            
            # 策略2：自动生成窗口尺寸（等比数列）
            min_size = 224  # 最小特征尺寸
            num_scales = 4  # 固定采样点数
            scale_factors = np.logspace(0, 1, num_scales, base=2)
            window_sizes = [int(base_size * f) for f in scale_factors]
            window_sizes = sorted({min(max(s, min_size), 3000) for s in window_sizes})
            
            # 策略3：根据长宽比调整尺寸组合
            if aspect_ratio > 1.5:  # 宽幅图像
                window_sizes = [int(s*aspect_ratio) for s in window_sizes]
            elif aspect_ratio < 0.67:  # 竖幅图像
                window_sizes = [int(s/aspect_ratio) for s in window_sizes]

            # 预处理优化 --------------------------------------------
            # 选择最优基准尺寸（最接近原图尺寸的2的幂次）
            base_size = 2 ** int(np.log2(base_size))
            base_transform = transforms.Compose([
                transforms.Resize((base_size, base_size),
                                interpolation=transforms.InterpolationMode.LANCZOS),
                self.base_transform
            ])
            
            # 半精度加速
            self.model.half()
            img_base = base_transform(image).unsqueeze(0).to(self.device).half()

            # 动态特征提取 ------------------------------------------
            features = []
            for size in window_sizes:
                # 保持长宽比的重采样
                target_size = (int(size*aspect_ratio), size) if aspect_ratio > 1 else (size, int(size/aspect_ratio))
                
                # GPU加速的智能插值
                img_tensor = torch.nn.functional.interpolate(
                    img_base, 
                    size=target_size,
                    mode= 'area' if size < base_size else 'bicubic',  # 下采样用area，上采样用bicubic
                    align_corners=False
                )

                # 自适应归一化（保持原图统计特性）
                if hasattr(self, 'adaptive_normalize'):
                    img_tensor = self.adaptive_normalize(img_tensor)

                # 混合精度推理
                with torch.no_grad(), torch.cuda.amp.autocast():
                    feature = self.model(img_tensor)
                
                features.append(feature.squeeze().float())

            # 动态权重分配 ------------------------------------------
            # 基于尺寸差异的权重（尺寸越接近原图权重越高）
            size_diffs = [abs(size - base_size) for size in window_sizes]
            weights = 1 / (torch.tensor(size_diffs, device=self.device) + 1e-6)
            weights = weights / weights.sum()

            # 加权融合
            final_feature = torch.stack([f * w for f, w in zip(features, weights)]).sum(dim=0)
            
            return final_feature

        except Exception as e:
            print(f"智能特征提取失败: {e}")
            return None

    def _extract_sliding_window_features(self, image: Image.Image) -> Optional[torch.Tensor]:
        """优化版滑动窗口特征提取（动态调整+批量处理）
        
        Args:
            image: PIL图片对象
            
        Returns:
            滑动窗口特征向量，处理失败返回None
        """
        try:
            # 获取原图信息
            orig_w, orig_h = image.size
            aspect_ratio = orig_w / orig_h
            
            # 动态窗口配置 -------------------------------------------
            # 根据原图尺寸自动选择关键窗口尺寸（示例逻辑，需根据实际调整）
            max_dim = max(orig_w, orig_h)
            window_sizes = sorted({
                int(2 ** np.round(np.log2(max_dim * 0.1))),  # 约10%尺寸
                int(2 ** np.floor(np.log2(max_dim * 0.5))),  # 约50%尺寸
                int(2 ** np.ceil(np.log2(max_dim)))          # 接近原图尺寸
            } & {256, 512, 1024, 2048, 3000})  # 与预设尺寸取交集
            
            # 智能步长调整（窗口尺寸越大步长越大）
            stride_ratios = {256:0.5, 512:0.4, 1024:0.3, 2048:0.2, 3000:0.15}
            
            # 预处理优化 --------------------------------------------
            # 生成基准图像（最大窗口尺寸）
            max_win_size = max(window_sizes)
            base_size = (int(max_win_size * aspect_ratio), max_win_size) if aspect_ratio > 1 else \
                        (max_win_size, int(max_win_size / aspect_ratio))
            
            transform = transforms.Compose([
                transforms.Resize(base_size[::-1], interpolation=transforms.InterpolationMode.LANCZOS),
                self.base_transform
            ])
            base_img = transform(image).to(self.device)
            
            # 半精度加速
            self.model.half()
            base_img = base_img.half()

            # 批量特征提取 ------------------------------------------
            all_features = []
            for win_size in window_sizes:
                # 动态步长选择
                stride = int(win_size * stride_ratios.get(win_size, 0.3))
                
                # 生成窗口坐标（考虑边缘填充）
                h, w = base_img.shape[1:]
                num_h = (h - win_size) // stride + 1
                num_w = (w - win_size) // stride + 1
                
                # 调整窗口数量上限（防止显存溢出）
                MAX_WINDOWS = 32  # 根据显存调整
                if num_h * num_w > MAX_WINDOWS:
                    stride = int(np.sqrt(h * w * win_size**2 / MAX_WINDOWS))
                    num_h = (h - win_size) // stride + 1
                    num_w = (w - win_size) // stride + 1

                # 批量裁剪窗口
                windows = []
                for i in range(num_h):
                    for j in range(num_w):
                        top = i * stride
                        left = j * stride
                        window = base_img[:, top:top+win_size, left:left+win_size]
                        windows.append(window)
                
                if not windows:
                    continue

                # 批量处理（自动分块防止OOM）
                BATCH_SIZE = 8  # 根据显存调整
                with torch.no_grad(), torch.cuda.amp.autocast():
                    for i in range(0, len(windows), BATCH_SIZE):
                        batch = torch.stack(windows[i:i+BATCH_SIZE])
                        features = self.model(batch)
                        all_features.append(features.cpu().float())  # 转移至CPU释放显存

            # 特征融合 ---------------------------------------------
            if not all_features:
                return None
            
            final_feature = torch.cat([f.view(-1, f.shape[-1]) for f in all_features], dim=0)
            final_feature = final_feature.mean(dim=0).to(self.device)

            return final_feature

        except Exception as e:
            print(f"滑动窗口特征提取失败: {e}")
            return None