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]:
        """提取多尺度特征。"""
        try:
            features_list = []
            width, height = image.size
            min_dim = min(width, height)
            max_dim = max(width, height)

            # 动态生成候选尺寸，基于原图尺寸
            scales = [0.25, 0.5, 0.75, 1.0, 1.5, 2.0]
            fixed_sizes = [256, 512, 1024, 2048]
            candidate_sizes = [int(min_dim * s) for s in scales] + fixed_sizes
            max_allowed = int(max_dim * 1.5)
            window_sizes = [size for size in candidate_sizes if 64 <= size <= max_allowed]
            window_sizes = sorted(list(set(window_sizes)))

            if not window_sizes:
                return None

            for size in window_sizes:
                # 保持宽高比调整较小边，并中心裁剪
                transform = transforms.Compose([
                    transforms.Resize(size, interpolation=transforms.InterpolationMode.LANCZOS),
                    transforms.CenterCrop(size),
                    self.base_transform
                ])
                img_tensor = transform(image).unsqueeze(0).to(self.device)
                
                with torch.no_grad():
                    feature = self.model(img_tensor)
                features_list.append(feature.squeeze())

            # 加权平均（较大尺度权重更高）
            weights = torch.linspace(1, 2, len(features_list), device=self.device)
            weights /= weights.sum()
            final_feature = torch.stack(features_list) * weights[:, None]
            return final_feature.sum(dim=0)
            
        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:
            # 基础图片检查
            if image is None or image.size[0] < 64 or image.size[1] < 64:
                print("图片无效或尺寸过小")
                return None
                
            # 获取原图信息
            orig_w, orig_h = image.size
            aspect_ratio = orig_w / orig_h
            max_dim = max(orig_w, orig_h)
            
            # 动态窗口配置 -------------------------------------------
            # 使用对数尺度生成窗口尺寸，确保合理的尺寸分布
            base_sizes = {256, 512, 1024, 2048}
            log_size = np.log2(max_dim)
            dynamic_sizes = {
                int(2 ** size) for size in [
                    np.floor(log_size - 1),  # 约50%原尺寸
                    np.ceil(log_size),       # 接近原尺寸
                ]
            }
            window_sizes = sorted(base_sizes & dynamic_sizes)
            
            if not window_sizes:
                # 如果没有合适的预设尺寸，选择最接近的基础尺寸
                closest_size = min(base_sizes, key=lambda x: abs(np.log2(x) - log_size))
                window_sizes = [closest_size]
                
            # 智能步长配置（窗口越大，步长比例越大）
            def get_stride_ratio(size):
                # 使用线性插值计算步长比例
                size_ratio = np.clip(size / 2048, 0.2, 0.8)
                return 0.2 + size_ratio * 0.3  # 步长比例范围：0.2-0.5
                
            # 预处理优化 --------------------------------------------
            # 生成基准图像（使用最大窗口尺寸）
            max_win_size = max(window_sizes)
            if aspect_ratio > 1:
                base_size = (int(max_win_size * aspect_ratio), max_win_size)
            else:
                base_size = (max_win_size, int(max_win_size / aspect_ratio))
                
            # 图像转换和加载
            transform = transforms.Compose([
                transforms.Resize(base_size[::-1], interpolation=transforms.InterpolationMode.LANCZOS),
                self.base_transform
            ])
            
            try:
                base_img = transform(image).to(self.device)
            except Exception as e:
                print(f"图像转换失败: {e}")
                return None
                
            # 特征提取 ---------------------------------------------
            all_features = []
            total_windows = 0
            
            for win_size in window_sizes:
                # 计算动态步长
                stride_ratio = get_stride_ratio(win_size)
                stride = max(int(win_size * stride_ratio), 16)  # 确保最小步长
                
                # 计算窗口数量
                h, w = base_img.shape[1:]
                num_h = (h - win_size) // stride + 1
                num_w = (w - win_size) // stride + 1
                
                # 内存优化：控制单个尺寸下的最大窗口数
                MAX_WINDOWS_PER_SIZE = 64
                if num_h * num_w > MAX_WINDOWS_PER_SIZE:
                    adjusted_stride = int(np.sqrt((h * w) / MAX_WINDOWS_PER_SIZE))
                    stride = max(stride, adjusted_stride)
                    num_h = (h - win_size) // stride + 1
                    num_w = (w - win_size) // stride + 1
                
                print(f"处理窗口 {win_size}x{win_size}, 步长 {stride}, 窗口数 {num_h * num_w}")
                
                # 批量处理窗口
                batch = []
                batch_size = min(16, num_h * num_w)  # 动态批次大小
                
                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]
                        
                        if torch.isnan(window).any() or torch.isinf(window).any():
                            continue
                            
                        batch.append(window)
                        total_windows += 1
                        
                        if len(batch) >= batch_size:
                            with torch.no_grad():
                                try:
                                    batch_tensor = torch.stack(batch)
                                    features = self.model(batch_tensor)
                                    all_features.append(features.cpu())  # 转移到CPU释放显存
                                except RuntimeError as e:
                                    print(f"批处理失败，尝试减小批次大小: {e}")
                                    if batch_size > 4:
                                        batch_size //= 2
                                        continue
                                batch = []
                
                # 处理剩余的窗口
                if batch:
                    with torch.no_grad():
                        try:
                            batch_tensor = torch.stack(batch)
                            features = self.model(batch_tensor)
                            all_features.append(features.cpu())
                        except RuntimeError as e:
                            print(f"处理剩余窗口失败: {e}")
            
            # 特征融合 ---------------------------------------------
            if not all_features:
                print("未能提取到有效特征")
                return None
                
            print(f"总处理窗口数: {total_windows}")
            
            # 合并所有特征
            try:
                final_features = torch.cat([f.view(-1, f.shape[-1]) for f in all_features], dim=0)
                
                # 如果特征数量过多，进行随机采样
                if final_features.size(0) > 1000:
                    indices = torch.randperm(final_features.size(0))[:1000]
                    final_features = final_features[indices]
                    
                return final_features.mean(dim=0).to(self.device)
                
            except Exception as e:
                print(f"特征融合失败: {e}")
                return None

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