3

简要说明

最近开始学习Object Detection，刚开始学习PyTorch，YOLOv5。所以我想为什么不建一个小项目来学习呢？用它来训练检测皮卡丘。

问题

我已经成功地用皮卡丘训练了模型，然后使用训练过的权重和我自己编写的 Python 脚本/代码来使用测试图像检测皮卡丘，现在，问题来了，皮卡丘可以成功检测到，但所有结果都显示为蓝色变色，这应该怎么办变成黄色，全部变成蓝色，蓝色变成黄色。

图 1 结果图像显示为蓝色（少数示例输出）

附加信息

我已经将这个项目推送到了 GitHub，请随意下载或拉取它进行调试。

包含所有文件的 GitHub 存储库

任何解决方案/建议都会有所帮助，谢谢。

编码

"""Object detection using YOLOv5

Pokemon Pikachu detecting

"""

# import os, sys to append YOLOv5 folder path
import os, sys

# import object detection needed modules and libraries
# pillow
from PIL import Image, ImageDraw, ImageFont
import numpy as np
import torch # PyTorch


# YOLOv5 folder path and related folder path settings
cwd = os.getcwd()
root_dir = (cwd + "/yolov5_stable")
sys.path.append(root_dir)


# import methods, functions from YOLOv5
from models.experimental import attempt_load
from utils.datasets import LoadImages
from utils.general import non_max_suppression, scale_coords
from utils.plots import colors



# define a function to show detected pikachu
def show_pikachu(img, det):
    labels = ["pikachu"]
    img = Image.fromarray(img)
    draw = ImageDraw.Draw(img)
    font_size = max(round(max(img.size)/40), 12)
    font = ImageFont.truetype(cwd + "/yolov5_stable/fonts/times.ttf")

    for info in det:
        color = colors(1)
        target, prob = int(info[5].cpu().numpy()), np.round(info[4].cpu().numpy(), 2)
        x_min, y_min, x_max, y_max = info[0], info[1], info[2], info[3]
        draw.rectangle([x_min, y_min, x_max, y_max], width = 3, outline = color)
        draw.text((x_min, y_min), labels[target] + ':' + str(prob), fill = color, font = font)

    # Bug unresolved, pikachu shown in blue discolouration

    return img


if __name__ == "__main__":
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    print("GPU State: ", device)
    
    data_path = (cwd + "/test_data/")
    weight_path = (cwd + "/yolov5_stable/weights/best_v1.pt")
    dataset = LoadImages(data_path)
    model = attempt_load(weight_path, map_location = device)
    model.to(device)
    
    for path, img, im0s, _ in dataset:
        img = torch.from_numpy(img).to(device)
        img = img.float() # uint8 to fp16/32
        img /= 255.0 # 0-255 to 0.0-1.0
        if img.ndimension() == 3:
            img = img.unsqueeze(0)
            
        pred = model(img)[0]
        pred = non_max_suppression(pred, 0.25, 0.45)
        for i, det in enumerate(pred):
            im0 = im0s.copy()
            det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
            result = show_pikachu(im0, det)
            result.show()
4

1 回答 1

3

问题是Image.fromarray需要 RGB 中的图像，而您在 BGR 中提供它们。你只需要改变它。您可以在多个地方执行此操作，例如：

Image.fromarray(img[...,::-1])  # assuming `img` is channel-last

一个证据是鼠标的红色部分（红色为RGB(255, 0, 0)）显示为蓝色（即RGB(0, 0, 255)）。仅供参考，黄色是RGB(255, 255, 0)和青色是RGB(0, 255, 255)，你也可以在你的情况下看到。

于 2021-08-20T17:41:13.650 回答