读取和检测图像，并保存在文件夹中

Question

首先，我只是简单地读取包含不同图像格式的文件夹的图像。其次，Yolo模型检测该类，绘制矩形并用颜色填充仅检测到的零件，然后将其保存到具有相同名称的另一个文件夹中。 第二种情况，如果该模型未检测到图像中的任何内容，则它将使用相同名称但在其他文件夹中保存相同的图像。

我的代码库被卡在第一个图像上，从不移动到第二张图像。我不知道发生了什么问题。

代码

import torch
import cv2
from matplotlib import pyplot as plt
from utils.plots import Annotator, colors, save_one_box
import os
import glob
# Load Ours Custom Model
model = torch.hub.load('.', 'custom', path='/media/bmvc/CM_1/yolov5/runs/train/exp4/weights/last.pt', source='local')
# Files extension
img_Extension = ['jpg', 'jpeg', 'png']
# Load all testing images
my_path = "/home/bmvc/Documents/hide_info_test_dataset/testing_images/"
# Save images into array
files = []
[files.extend(glob.glob(my_path + '*.' + e)) for e in img_Extension]
# Iteration on all images
images = [cv2.imread(file) for file in files]
total_images = 1
# Taking only image name to save with save name
image_file_name = ''
for file in files:
    for im in images:
        detections = model(im[..., ::-1])
        results = detections.pandas().xyxy[0].to_dict(orient="records")
        if len(results) == 0:
            cv2.imwrite(os.path.join("/home/bmvc/Documents/hide_info_test_dataset/detected/", file), im)
        else:
            for result in results:
                print(result['class'])
                con = result['confidence']
                cs = result['class']
                x1 = int(result['xmin'])
                y1 = int(result['ymin'])
                x2 = int(result['xmax'])
                y2 = int(result['ymax'])
                imagee = cv2.rectangle(im, (x1, y1), (x2, y2), (0, 255, 0), -1)
                cv2.imwrite(os.path.join("/home/bmvc/Documents/hide_info_test_dataset/detected/", file), im)
            total_images += 1

Answer 1

我放置了很多完全没有用的循环，例如读取不同的扩展文件，仅读取图像。我改进了整体实现，仅使用一个循环来解决上述问题。

import torch
import cv2
from PIL import Image

from utils.plots import Annotator, colors, save_one_box
import os
import glob
import numpy as np
# Load Ours Custom Model
model = torch.hub.load('.', 'custom', path='/media/bmvc/CM_1/yolov5/runs/train/exp4/weights/last.pt', source='local')

# Files extension
img_Extension = ['jpg', 'jpeg', 'png']

# Load all testing images
my_path = "/home/bmvc/Documents/hide_info_test_dataset/testing_images/"

# Save images into array
files = []
[files.extend(glob.glob(my_path + '*.' + e)) for e in img_Extension]

# Iteration on all images
images = [cv2.imread(file) for file in files]

total_images = 1

# Taking only image name to save with save name
image_file_name = ''

for img in glob.glob(my_path + '*.*'):
    img_bgr_rgb = cv2.imread(img)
    file_Name = os.path.basename(img)
    detections = model(img_bgr_rgb[:, :, ::-1])
    results = detections.pandas().xyxy[0].to_dict(orient="records")
    if len(results) == 0:
        cv2.imwrite(os.path.join("/home/bmvc/Documents/hide_info_test_dataset/detected/", file_Name), img_bgr_rgb)
    else:
        for result in results:
            print(result['class'])
            con = result['confidence']
            cs = result['class']
            x1 = int(result['xmin'])
            y1 = int(result['ymin'])
            x2 = int(result['xmax'])
            y2 = int(result['ymax'])
            imagee = cv2.rectangle(img_bgr_rgb, (x1, y1), (x2, y2), (255, 87, 51), -1)
            cv2.imwrite(os.path.join("/home/bmvc/Documents/hide_info_test_dataset/detected/", file_Name), img_bgr_rgb)