Learning Object Detection 检测结果显示变色

Question

简要说明

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

问题

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

Fig-1 Result images showed in blue discolouration(few example outputs)

补充信息

我已经将这个项目推送到GitHub，欢迎下载或拉取调试。

GitHub repository where it contains all the files

任何 solution/suggestion 都会有帮助，谢谢。

代码

"""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()

Answer 1

问题是 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)，您也可以在您的案例中看到。

Learning Object Detection 检测结果显示变色

Learning Object Detection Detected result showed in discolouration

python

colors

object-detection

pytorch

yolov5