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

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)