数字无极图像识别实现：基于特征向量与相似度匹配的闭集识别方案

import torch
from torchvision import transforms, models
from PIL import Image

# 定义预处理流水线
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    ),
])

# 加载预训练 ResNet50，去掉最后的全连接层，只取池化后的特征
model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
model.fc = torch.nn.Identity()
model.eval()

def extract_feature(image_path: str) -> torch.Tensor:
    img = Image.open(image_path).convert("RGB")
    tensor = transform(img).unsqueeze(0)
    with torch.no_grad():
        feat = model(tensor)
    return feat.squeeze().numpy()

import faiss
import numpy as np

# features.shape == (N, 2048)，N 是图库图片数量
# faiss 默认用 float32
features = features.astype("float32")

# 构建 HNSW 索引
# M 是每个节点的邻居数，efConstruction 控制构建时的搜索深度
d = features.shape[1]  # 2048
index = faiss.IndexHNSWFlat(d, M=32)
index.hnsw.efConstruction = 200
index.add(features)

# 查询时扩大搜索范围以提升召回
index.hnsw.efSearch = 64

# 搜索 topk
D, I = index.search(query_feat, k=5)

from fastapi import FastAPI, UploadFile, File
from pydantic import BaseModel
from typing import List
import asyncio

app = FastAPI()

class RecognitionResult(BaseModel):
    label: str
    confidence: float
    image_url: str

@app.post("/api/v1/recognize", response_model=List[RecognitionResult])
async def recognize(file: UploadFile = File(...)):
    # 1. 异步读取上传的文件内容
    content = await file.read()

    # 2. 利用线程池做同步的图像解码和特征提取
    feat = await asyncio.to_thread(extract_feature_from_bytes, content)

    # 3. 向量检索（FAISS search 本身很快，跑在线程池里也没问题）
    D, I = await asyncio.to_thread(index.search, feat.reshape(1, -1), 5)

    # 4. 组装返回结果
    results = []
    for score, idx in zip(D[0], I[0]):
        if score < threshold:
            break
        results.append(RecognitionResult(
            label=label_map[idx],
            confidence=float(score),
            image_url=image_urls[idx],
        ))
    return results

from torch.utils.data import DataLoader, Dataset

class ImageFolderDataset(Dataset):
    def __init__(self, paths, transform):
        self.paths = paths
        self.transform = transform

    def __len__(self):
        return len(self.paths)

    def __getitem__(self, idx):
        img = Image.open(self.paths[idx]).convert("RGB")
        return self.transform(img), self.paths[idx]

dataset = ImageFolderDataset(image_paths, transform)
loader = DataLoader(dataset, batch_size=64, num_workers=4, pin_memory=True)

all_features = []
all_paths = []
for batch, paths in loader:
    batch = batch.cuda()
    with torch.no_grad():
        feats = model(batch)
    all_features.append(feats.cpu().numpy())
    all_paths.extend(paths)

features = np.concatenate(all_features, axis=0)

# 用 IDMap 包装索引以支持自定义 ID
id_index = faiss.IndexIDMap(index)
id_index.add_with_ids(features, np.arange(len(features), dtype=np.int64))

wx.chooseMedia({
  count: 1,
  mediaType: ['image'],
  sourceType: ['camera', 'album'],
  success(res) {
    const tempPath = res.tempFiles[0].tempFilePath;
    wx.uploadFile({
      url: 'https://api.example.com/api/v1/recognize',
      filePath: tempPath,
      name: 'file',
      success(uploadRes) {
        const data = JSON.parse(uploadRes.data);
        if (data.length > 0) {
          this.setData({
            resultLabel: data[0].label,
            confidence: (data[0].confidence * 100).toFixed(1) + '%',
          });
        } else {
          this.setData({ resultLabel: '未识别到匹配动物' });
        }
      },
      fail() {
        wx.showToast({ title: '上传失败，请重试', icon: 'none' });
      }
    });
  }
});