pytorch加载10个yolov5模型并推理1万图片

慕容洛凝

于 2023-03-25 12:33:57 发布

阅读量698

点赞数

文章标签： pytorch 深度学习人工智能 python 目标检测

本文链接：https://blog.youkuaiyun.com/weixin_38132142/article/details/129765512

版权

该代码段导入了torch,torchvision和PIL库，加载了10个YOLOv5模型，并对指定输入文件夹中的10000张图像进行推理。每个图像都通过所有模型进行处理，然后将结果保存到输出文件夹中。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

import torch
import torchvision
from PIL import Image
import os

# Load the models
models = []
for i in range(10):
    model_path = f"yolov5_{i}.pt"
    model = torch.hub.load('ultralytics/yolov5', 'custom', path=model_path)
    models.append(model)

# Perform inference on the images
input_folder = "path/to/input/folder"
output_folder = "path/to/output/folder"
for i in range(10000):
    image_path = os.path.join(input_folder, f"{i}.jpg")
    image = Image.open(image_path)
    for model in models:
        results = model(image)
        # Process the results here
    image.save(os.path.join(output_folder, f"{i}.jpg"))