Flask加载YOLOv5模型报错：ModuleNotFoundError: No module named 'models'

项目背景 #

我正在开发一个图像目标检测项目，可识别图像中的物体并生成对应语音音频（例如识别到汽车时，音频文件会播放'CAR'）。但运行Flask应用加载YOLOv5的best.pt模型时触发错误。

错误日志 #

PS E:\ObjRec> python app.py
Loading YOLOv5 model...
Traceback (most recent call last):
  File "E:\ObjRec\app.py", line 11, in <module>
    model = torch.load('best.pt')
            ^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Admin\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\LocalCache\local-packages\Python311\site-packages\torch\serialization.py", line 1026, in load
    return _load(opened_zipfile,
           ^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Admin\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\LocalCache\local-packages\Python311\site-packages\torch\serialization.py", line 1438, in _load
    result = unpickler.load()
             ^^^^^^^^^^^^^^^^
  File "C:\Users\Admin\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\LocalCache\local-packages\Python311\site-packages\torch\serialization.py", line 1431, in find_class
    return super().find_class(mod_name, name)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ModuleNotFoundError: No module named 'models'

我的app.py代码 #

from flask import Flask, request, jsonify
import torch
from torchvision import transforms
from PIL import Image
import io

app = Flask(__name__)

print('Loading YOLOv5 model...')
model = torch.load('best.pt')
model.eval()

preprocess = transforms.Compose([
    transforms.Resize((416, 416)),
    transforms.ToTensor(),
])

@app.route('/detect', methods=['POST'])
def detect_objects():
    if 'image' not in request.files:
        return jsonify({'error': 'No image uploaded'}), 400
    image_file = request.files['image']
    image_bytes = image_file.read()
    image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
    image = preprocess(image)
    results = model(image.unsqueeze(0)) 
    labels = results.names
    
    return jsonify({'labels': labels}), 200

if __name__ == '__main__':
    app.run(debug=True)

尝试解决的错误 #

执行pip install models命令时遭遇metadata-generation-failed错误：

PS E:\ObjRec> pip install models
Collecting models
  Using cached models-0.9.3.tar.gz (16 kB)
  Preparing metadata (setup.py) ... error
  error: subprocess-exited-with-error

  × python setup.py egg_info did not run successfully.
  │ exit code: 1
  ╰─> [8 lines of output]
      Traceback (most recent call last):
        File "<string>", line 2, in <module>
        File "<pip-setuptools-caller>", line 34, in <module>
        File "C:\Users\Admin\AppData\Local\Temp\pip-install-8yqo_00s\models_16df2befc82545e7a9071610e1043cde\setup.py", line 25, in <module>
          import models
        File "C:\Users\Admin\AppData\Local\Temp\pip-install-8yqo_00s\models_16df2befc82545e7a9071610e1043cde\models\__init__.py", line 23, in <module>
          from base import *
      ModuleNotFoundError: No module named 'base'
      [end of output]

  note: This error originates from a subprocess, and is likely not a problem with pip.
error: metadata-generation-failed

× Encountered error while generating package metadata.
╰─> See above for output.

note: This is an issue with the package mentioned above, not pip.

解决方案 #

问题根源 #

YOLOv5的best.pt模型依赖YOLOv5项目内置的models模块，直接用torch.load()加载时Python找不到该模块；而你尝试安装的models包和YOLOv5完全无关，因此安装报错。

正确解决步骤 #

1. 安装YOLOv5官方依赖
   安装官方维护的ultralytics包，包含YOLOv5运行所需的所有模块：

   pip install ultralytics
   

2. 修改模型加载代码
   替换原有的torch.load()加载方式，改用YOLOv5的官方API加载模型：

   # 替换原代码中model = torch.load('best.pt')这一行
   from ultralytics import YOLO
   model = YOLO('best.pt')
   model.eval()
   

3. 调整推理逻辑
   YOLOv5的推理返回结果格式与原代码预期不同，修改检测函数：

   @app.route('/detect', methods=['POST'])
   def detect_objects():
       if 'image' not in request.files:
           return jsonify({'error': 'No image uploaded'}), 400
       image_file = request.files['image']
       image_bytes = image_file.read()
       image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
       
       # YOLO模型会自动处理图像预处理，无需手动转换
       results = model(image)
       
       # 提取检测到的物体标签
       detected_labels = []
       for result in results:
           for box in result.boxes:
               label = result.names[int(box.cls)]
               detected_labels.append(label)
       
       return jsonify({'detected_labels': detected_labels}), 200
   

4. 验证运行
   重新执行python app.py，模型即可正常加载并处理检测请求。

内容的提问来源于stack exchange，提问作者Allen07