开发逆向

使用YOLOWorld与相似度比较(概率)识别腾讯混元点选

本文将会涉及到的内容包括: YOLOv8-World的使用 YOLOv8导出为ONNX并正常使用 根据YOLOv8识别后得到的Classes进行选择 基于OpenCV的(直方图)色系识别

前言

本文仅作学习和交流使用, 如有侵权, 请联系删除.

本文将会涉及到的内容包括:

  1. YOLOv8-World的使用

  2. YOLOv8导出为ONNX并正常使用

  3. 根据YOLOv8识别后得到的Classes进行选择

  4. 基于OpenCV的(直方图)色系识别

继文字点选, 图标点选之后, 腾讯又有了基于AIGC图像的点选 (混元AI)

本文将从识别角度进行探讨, 以YOLOv8-World为主进行识别, 对于没有的描述则使用相似度识别然后硬蒙(猜)

也就是P(Correct) = 1/3

YOLOv8

识别效果

下载模型

Python环境, 不再赘述

地址: github.com/ultralytics/assets/releases 找到适合你的权重文件(.pt)下载即可 (注: 模型名称中需要有-world )

本文使用权重文件yolov8m-world.pt (版本8.3.0)

安装ultralytics (pip install ultralytics )

选择适合自己的模型(s/m/l/x)

不同的模型有着不同的大小, 推理速度也有很大差距, 如果需要选一个适合自己电脑的模型, 请参照下表

(注: YOLO-World没有n, 故省略)

Model

size
(pixels)

mAPval
50-95

Speed
CPU ONNX
(ms)

Speed
A100 TensorRT
(ms)

params
(M)

FLOPs
(B)

YOLOv8s

640

44.9

128.4

1.20

11.2

28.6

YOLOv8m

640

50.2

234.7

1.83

25.9

78.9

YOLOv8l

640

52.9

375.2

2.39

43.7

165.2

YOLOv8x

640

53.9

479.1

3.53

68.2

257.8

笔者建议使用m模型, 速度适中且准确率高

验证

命令行执行: yolo detect predict model=yolov8m-world.pt source=YourImageHere.png

如果一切正常, 则在同级目录下会生成目录runs/detect/predict 此时应能做到开头所给效果

导出为ONNX

先安装好onnxruntime, onnxslim

pip install onnxruntime onnxslim

随后执行 (注: imgsz为要识别图像的大小, 格式为 宽,长 不要搞反了!!! )

yolo export model=yolov8m-world.pt format=onnx imgsz=480,672

稍等片刻就会生成yolov8m-world.onnx

封装

from PIL import Image
import onnxruntime
import numpy as np
import cv2

# YOLOv8
class YOLOv8:
    def __init__(self, OnnxPath: str, CFThresh: float = 0.5, IOUThresh: float = 0.3) -> None:
        self.OSession = onnxruntime.InferenceSession(OnnxPath)
        self.CFThresh = CFThresh
        self.IOUThresh = IOUThresh

    # 预处理
    def Pre(self, Img: Image.Image) -> np.ndarray:
        return np.expand_dims(np.transpose(np.array(Img.convert('RGB')) / 255.0, (2, 0, 1)), axis=0).astype(np.float32)

    # 后处理, 返回 (Box, ClassID, Score)
    def PostP(self, Output: list[np.ndarray]) -> list[tuple[list[int], int, float]]:
        Outputs = np.transpose(np.squeeze(Output[0]))
        Boxes = []
        Scores = []
        ClassIDs = []
        Gain = 1.0

        for i in range(Outputs.shape[0]):
            CS = Outputs[i][4:]
            MS = np.amax(CS)
            if MS >= self.CFThresh:
                CID = np.argmax(CS)

                X, Y, W, H = Outputs[i][0], Outputs[i][1], Outputs[i][2], Outputs[i][3]

                ClassIDs.append(CID)
                Scores.append(MS)
                # Left, Top, Width, Height
                Boxes.append([int((X - W / 2) / Gain), int((Y - H / 2) / Gain), int(W / Gain), int(H / Gain)])

        Res = [(Boxes[i], ClassIDs[i], Scores[i]) for i in cv2.dnn.NMSBoxes(Boxes, Scores, self.CFThresh, self.IOUThresh)]
        return Res

    # 检测, 返回结果为[(Box(X, Y, W, H), ClassID, 置信度), ...]
    def Inference(self, Img: Image.Image) -> list[tuple[list[int], int, float]]:
        return self.PostP(self.OSession.run(None, {self.OSession.get_inputs()[0].name: self.Pre(Img)}))

检测指定

现在我们已经可以识别图片中的物体了, 接下来我们需要做的就是找到指定描述的物体及其所在位置

例如, 该验证的指令(返回参数instruction )为: "红色玻璃花瓶"

且YOLOv8可以识别花瓶, 但官方Classes的名称为英文, 此时我们需要将官方给出的Classes转为中文, 并且对应同一个物体要有不同的名称(不同的表述方式表示同一个物体, 且考虑近似物体识别时的误差)

以下为本人修改后用于识别的Classes:

# Classes, 将推理后得到的ClassesID直接转为下标从此中取即可
Classes = ['人', '自行车', '车', '摩托', '飞机', '公交', '火车 地铁', '卡车', '船', '交通', '消防栓', '停车', '停车收费', '凳', ' 鸟', '猫', '狗', '马', '羊', '牛', '象', '熊', '斑马', '长颈鹿', '包', '雨伞', '包', '领带', '箱', '飞盘', '滑雪', '滑雪', '球', '风筝', '棒球', '棒球', '滑板', '冲浪板', '网球', '瓶', '杯', '杯', '叉', '刀', '勺', '碗', '香蕉', '苹果', '三明治 面包片 蛋糕', '橘子 黄色水果', '西兰花', '萝卜', '热狗', '披萨', '甜甜圈', '蛋糕', '椅子 沙发', '椅子 沙发', '盆栽 植物', '床', '桌', '厕所 马桶', '电视', '笔记本', '鼠', '遥控', '键盘', '手机', '微波炉', '烤', '面包机', '洗', '冰箱', '书 本', '钟 表', '瓶', '剪刀', '熊', '吹风机', '牙刷']

相似度识别

对于模型无法识别的情况, 我们只能通过猜测来尽可能增大识别成功的概率

但是吧, 这猜也不是瞎猜捏

原理

对比图像可知, 每一组验证小图片中都有一些特征与关联: 共有6张小图片, 其中图片每两张可视为一组, 每一组中的图片都相似.

也就是说, 6张图片我们可以分为三组, 则有P(Correct) = 1/3

图片处理

我们可以先将此大图片裁剪为6张小图

from PIL import Image
Img: Image.Image
Imgs = [Img.crop((0, 34, 219, 253)), Img.crop((226, 34, 445, 253)), Img.crop((452, 34, 671, 253)), Img.crop((0, 260, 219, 479)), Img.crop((226, 260, 445, 479)), Img.crop((452, 260, 671, 479))]

并转换为OpenCV图像 (当然, 一开始就是OpenCV的话也可以)

from PIL import Image
import cv2
import numpy as np

# 转cv2
def ToCV2(Img: Image.Image) -> cv2.Mat:
    Img = Img.convert('RGB')
    cv_img = np.array(Img)
    # 处理可能的通道问题（PIL 采用 RGB, OpenCV 采用 BGR）
    if Img.mode == "RGB":
        cv_img = cv2.cvtColor(cv_img, cv2.COLOR_RGB2BGR)
    elif Img.mode == "RGBA":
        cv_img = cv2.cvtColor(cv_img, cv2.COLOR_RGBA2BGRA)
    return cv_img

色系匹配(三通道直方图)

import cv2

# 色系匹配
def ClassifyHistCompare(Img1: cv2.Mat, Img2: cv2.Mat) -> float:
    Similar = 0
    Degree = 0
    for i, a in zip(Img1, Img2):
        Hist1 = cv2.calcHist([i], [0], None, [256], [0.0, 255.0])
        Hist2 = cv2.calcHist([a], [0], None, [256], [0.0, 255.0])
        for i in range(len(Hist1)):
            if Hist1[i] != Hist2[i]:
                Degree = Degree + (1 - abs(Hist1[i] - Hist2[i]) / max(Hist1[i], Hist2[i]))
            else:
                Degree = Degree + 1
        Degree = Degree / len(Hist1)
        Similar += Degree
    Similar = Similar / 3
    return Similar

最后返回值范围: 0-1

此时对于第三张图片有如下匹配结果

第n个图片

1

2

3

4

5

6

1

0.6900668

0.78154117

0.581816

0.60272217

0.52412134

2

0.6379582

0.65864617

0.56729025

0.54956526

3

0.5587124

0.58857036

0.51522857

4

0.6517611

0.5770158

5

0.5889302

可以将其明显分组 (虽然分的并不准)

后记

TDC补环境补麻了

我一时间竟不知道如何评价YOLO的中文文档......

Comment