搭建python的Flask作为http API接口(结合基于ResNet的图像识别)

项目目录接口：

安装flask:

pip install Flask

image_matching.py:

import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
import os
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
from flask import Flask, request, jsonify

app = Flask(__name__)

# 1. 加载预训练的 ResNet 模型，并去掉分类层
def get_resnet_model():
    model = models.resnet50(pretrained=True)  # 使用 ResNet-50 作为基础模型
    # 去掉最后的全连接层
    model = nn.Sequential(*list(model.children())[:-1])
    model.eval()  # 切换到评估模式
    return model

# 2. 图像预处理
def preprocess_image(image_path):
    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]
        )
    ])
    image = Image.open(image_path).convert("RGB")
    image = transform(image).unsqueeze(0)  # 增加批次维度
    return image

# 3. 提取图像特征
def extract_features(model, image):
    with torch.no_grad():
        features = model(image).squeeze()  # 提取特征
        features = features.flatten()  # 扁平化
    return features.numpy()

# 4. 计算相似度（使用余弦相似度）
def calculate_similarity(feature1, feature2):
    return cosine_similarity([feature1], [feature2])[0][0]

# 5. 加载数据库中的图像并提取特征
def build_image_database(image_dir, model):
    features_database = {}
    for image_name in os.listdir(image_dir):
        image_path = os.path.join(image_dir, image_name)
        if os.path.isfile(image_path):
            image = preprocess_image(image_path)
            feature = extract_features(model, image)
            features_database[image_name] = feature
    return features_database

# 6. 查询数据库，找出最匹配的图像
def find_most_similar_image(query_image_path, image_dir, model):
    query_image = preprocess_image(query_image_path)
    query_feature = extract_features(model, query_image)

    features_database = build_image_database(image_dir, model)

    best_match = None
    best_similarity = -1  # 初始时设定最差的相似度（-1）

    for image_name, feature in features_database.items():
        similarity = calculate_similarity(query_feature, feature)
        if similarity > best_similarity:
            best_similarity = similarity
            best_match = image_name

    return best_match, best_similarity

# 7. 定义 API 路由
@app.route('/find_match', methods=['POST'])
def find_match():
    # 从请求中获取查询图像的文件名
    data = request.get_json()
    query_image_name = data.get('query_image_name', '')

    if not query_image_name:
        return jsonify({'error': 'No query image name provided'}), 400

    # 拼接查询图像的路径和图像库路径
    query_image_path = os.path.join('images', query_image_name)
    image_dir = 'images_database'

    if not os.path.exists(query_image_path):
        return jsonify({'error': 'Query image not found'}), 404

    # 加载模型
    model = get_resnet_model()

    # 查找最匹配的图像
    best_match, similarity = find_most_similar_image(query_image_path, image_dir, model)

    if best_match:
        # 转换 similarity 为 Python 原生类型 float
        return jsonify({
            'best_match': best_match,
            'similarity': float(similarity)  # 转换为 float
        })
    else:
        return jsonify({'error': 'No match found'}), 404



if __name__ == '__main__':
    # 启动 Flask API
    app.run(host='0.0.0.0', port=3000)

执行文件：

注意这里的3000端口是docker 容器的3000端口

容器外的对应关系是：8001-->3000

所以真正访问的地址： 使用postman：