YOLOv5检测框crop、MobileNetv3分类网络

在实际深度学习项目中,目标检测算法检测出的目标也会作为分类网络的输入数据,利用目标检测算法的对被检测图像进行抠图,以抠出来的图来扩充分类网络的数据。本文主要讲解yolov5和mobilenetv3结合使用扩展数据样本。

1、yolov5检测框crop

yolov5脚本中自带对检测出来的物体抠图功能,为减少标注、训练等工作量,直接使用官方提供的COCO128.yaml配置文件和yolov5s.pt权重文件,以下是执行的过程。

1)配置文件COCO128.yaml

# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ../datasets/coco128 # dataset root dir
train: images/train2017 # train images (relative to 'path') 128 images
val: images/train2017 # val images (relative to 'path') 128 images
test: # test images (optional)

# Classes
names:
  0: person
  1: bicycle
  2: car
  3: motorcycle
  4: airplane
  5: bus
  6: train
  7: truck
  8: boat
  9: traffic light
  10: fire hydrant
  ...

2)修改detect.py

新增一些交通信号灯的场景图片、视频,将其放在相应相应路径下,更改代码中的以下部分:

@smart_inference_mode()
def run(
    weights=ROOT / "yolov5s.pt",  # model path or triton URL
    source=ROOT / "classify/data/videos",  # file/dir/URL/glob/screen/0(webcam)
    data=ROOT / "data/COCO128.yaml",  # dataset.yaml path
    imgsz=(640, 640),  # inference size (height, width)
    conf_thres=0.25,  # confidence threshold
    iou_thres=0.45,  # NMS IOU threshold
    max_det=1000,  # maximum detections per image
    device="",  # cuda device, i.e. 0 or 0,1,2,3 or cpu
    view_img=False,  # show results
    save_txt=False,  # save results to *.txt
    save_csv=False,  # save results in CSV format
    save_conf=False,  # save confidences in --save-txt labels
    save_crop=False,  # save cropped prediction boxes
    nosave=False,  # do not save images/videos
    classes=None,  # filter by class: --class 0, or --class 0 2 3
    agnostic_nms=False,  # class-agnostic NMS
    augment=False,  # augmented inference
    visualize=False,  # visualize features
    update=False,  # update all models
    project=ROOT / "runs/detect",  # save results to project/name
    name="exp",  # save results to project/name
    exist_ok=False,  # existing project/name ok, do not increment
    line_thickness=3,  # bounding box thickness (pixels)
    hide_labels=False,  # hide labels
    hide_conf=False,  # hide confidences
    half=False,  # use FP16 half-precision inference
    dnn=False,  # use OpenCV DNN for ONNX inference
    vid_stride=1,  # video frame-rate stride
):
...

def parse_opt():
    """Parses command-line arguments for YOLOv5 detection, setting inference options and model configurations."""
    parser = argparse.ArgumentParser()
    parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s.pt", help="model path or triton URL")
    parser.add_argument("--source", type=str, default=ROOT / "/home/mengwen/yolov5/classify/data/videos", help="file/dir/URL/glob/screen/0(webcam)")
    parser.add_argument("--data", type=str, default=ROOT / "data/COCO128.yaml", help="(optional) dataset.yaml path")
    parser.add_argument("--imgsz", "--img", "--img-size", nargs="+", type=int, default=[640], help="inference size h,w")
    parser.add_argument("--conf-thres", type=float, default=0.25, help="confidence threshold")
    parser.add_argument("--iou-thres", type=float, default=0.45, help="NMS IoU threshold")
    parser.add_argument("--max-det", type=int, default=1000, help="maximum detections per image")
    parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu")
    parser.add_argument("--view-img", action="store_true", help="show results")
    parser.add_argument("--save-txt", action="store_true", help="save results to *.txt")
    parser.add_argument("--save-csv", action="store_true", help="save results in CSV format")
    parser.add_argument("--save-conf", action="store_true", help="save confidences in --save-txt labels")
    parser.add_argument("--save-crop", action="store_true", help="save cropped prediction boxes")
    parser.add_argument("--nosave", action="store_true", help="do not save images/videos")
    parser.add_argument("--classes", nargs="+", type=int, help="filter by class: --classes 0, or --classes 0 2 3")
    parser.add_argument("--agnostic-nms", action="store_true", help="class-agnostic NMS")
    parser.add_argument("--augment", action="store_true", help="augmented inference")
    parser.add_argument("--visualize", action="store_true", help="visualize features")
    parser.add_argument("--update", action="store_true", help="update all models")
    parser.add_argument("--project", default=ROOT / "runs/detect", help="save results to project/name")
    parser.add_argument("--name", default="exp", help="save results to project/name")
    parser.add_argument("--exist-ok", action="store_true", help="existing project/name ok, do not increment")
    parser.add_argument("--line-thickness", default=3, type=int, help="bounding box thickness (pixels)")
    parser.add_argument("--hide-labels", default=False, action="store_true", help="hide labels")
    parser.add_argument("--hide-conf", default=False, action="store_true", help="hide confidences")
    parser.add_argument("--half", action="store_true", help="use FP16 half-precision inference")
    parser.add_argument("--dnn", action="store_true", help="use OpenCV DNN for ONNX inference")
    parser.add_argument("--vid-stride", type=int, default=1, help="video frame-rate stride")
    opt = parser.parse_args()
    opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1  # expand
    print_args(vars(opt))
    return opt

3) crop

在对检测框进行crop时,可以根据自己的需求选择对应类别,这里以traffic_lights为例。

python detect.py --classes 9 --save-crop

在这里插入图片描述可以看到crop出来的之后交通信号灯种类多样,

2、mobilenetv3分类网络

鉴于数据样本较大,种类较多的情况,现使用原始数据集训练的onnx权重文件对图片进行识别划分。

1)原始数据集类别

将交通信号灯按照颜色、方向等属性分为33类(可根据自己需求调整)

    "0": "composite",
    "1": "green-circle",
    "2": "green-lane-signal",
    "3": "green-left",
    "4": "green-left_turn",
    "5": "green-number",
    "6": "green-others",
    "7": "green-person",
    "8": "green-right",
    "9": "green-straight",
    "10": "green-straight_left",
    "11": "green-turn",
    "12": "others",
    "13": "red-circle",
    "14": "red-left",
    "15": "red-left_turn",
    "16": "red-number",
    "17": "red-others",
    "18": "red-person",
    "19": "red-right",
    "20": "red-straight",
    "21": "red-straight_left",
    "22": "red-turn",  
    "23": "red_signal_lights",
    "24": "yellow-circle",
    "25": "yellow-left",
    "26": "yellow-person",
    "27": "yellow-number",
    "28": "yellow-others",
    "29": "yellow-right",
    "30": "yellow-straight",
    "31": "yellow-straight_left",
    "32": "yellow-turn"

2)mobilenetv3算法

使用轻量级分类算法mobilenetv3对原始数据集(原始数据集可根据类别自行制作)进行训练,根据生成的onnx权重文件进行推理。
(1)onnx权重文件生成见:
mobilenetv3_small实现交通灯识别与推理
(2)分类推理

import os
from PIL import Image
import onnxruntime as ort
import numpy as np

def softmax(x):
    x = x.reshape(-1)
    e_x = np.exp(x - np.max(x))
    return e_x / e_x.sum(axis=0)

def postprocess(result):
    return softmax(np.array(result)).tolist()

class_mapping = {
        "0": "composite",
        "1": "green-circle",
        "2": "green-lane-signal",
        "3": "green-left",
        "4": "green-left_turn",
        "5": "green-number",
        "6": "green-others",
        "7": "green-person",
        "8": "green-right",
        "9": "green-straight",
        "10": "green-straight_left",
        "11": "green-turn",
        "12": "others",
        "13": "red-circle",
        "14": "red-left",
        "15": "red-left_turn",
        "16": "red-number",
        "17": "red-others",
        "18": "red-person",
        "19": "red-right",
        "20": "red-straight",
        "21": "red-straight_left",
        "22": "red-turn",
        "23": "red_signal_lights",
        "24": "yellow-circle",
        "25": "yellow-left",
        "26": "yellow-person",
        "27": "yellow-number",
        "28": "yellow-others",
        "29": "yellow-right",
        "30": "yellow-straight",
        "31": "yellow-straight_left",
        "32": "yellow-turn"
}

if __name__ == "__main__":
    onnx_model_path = "./traffic_light.onnx"
    ort_session = ort.InferenceSession(onnx_model_path)
    onnx_input_name = ort_session.get_inputs()[0].name
    onnx_outputs_names = ort_session.get_outputs()[0].name

    input_folder = 'E:/红绿灯识别/crops/crops-1/traffic light'
    output_folder = 'E:/红绿灯识别/crops/cls-crop-1'

    os.makedirs(output_folder, exist_ok=True)

    for file_name in os.listdir(input_folder):
        if file_name.endswith(('.jpg', '.jpeg', '.png')):
            img_path = os.path.join(input_folder, file_name)
            original_img = Image.open(img_path)
            original_width, original_height = original_img.size

            img = original_img.convert("RGB")
            img = img.resize((224, 224))
            img = np.asarray(img, np.float32) / 255.0
            img = img.transpose(2, 0, 1)
            img = img[np.newaxis, :, :, :]
            input_blob = np.array(img, dtype=np.float32)
            onnx_result = ort_session.run([onnx_outputs_names], input_feed={onnx_input_name: input_blob})
            res = postprocess(onnx_result)
            idx = np.argmax(res)
            predicted_class_name = class_mapping.get(str(idx), "unknown")
            probability = res[idx]

            print("Predicted class:", predicted_class_name)
            print("Probability:", probability)

            # 如果预测概率低于阈值,删除图片
            threshold = 0.5
            if predicted_class_name == "unknown" or probability < threshold:
                os.remove(img_path)
                print("Image removed due to low probability or unknown class.")
            else:
                # 保存结果图片到对应的类别文件夹
                output_class_folder = os.path.join(output_folder, predicted_class_name)
                os.makedirs(output_class_folder, exist_ok=True)
                result_image_path = os.path.join(output_class_folder, f'{file_name}')
                original_img.save(result_image_path)

3)效果

总体分类效果效果还可以,存在一些伪样本可以人为剔除,大大减少工作量,效果如下:

(1)green_circle
在这里插入图片描述(3)red_left
在这里插入图片描述(3)yellow_right
在这里插入图片描述

  • 10
    点赞
  • 15
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
### 回答1: `save-crop`是YOLOv5中的一个选项,用于将检测到的物体从原始图像中裁剪出来并保存为单独的图像文件。这个选项可以用于训练数据集的准备,也可以用于实际应用中。 使用`save-crop`选项的方法如下: 1. 在运行YOLOv5检测脚本时,加上`--save-crop`选项,例如: ``` python detect.py --weights yolov5s.pt --img 640 --conf 0.25 --source data/images/ --save-crop ``` 2. 运行脚本后,YOLOv5检测输入源中的所有图像,并将检测到的物体裁剪出来,保存为单独的图像文件。这些图像文件将保存在`runs/detect/exp/crops`目录下。 注意,`save-crop`选项需要在YOLOv5的源代码中进行配置才能正常使用。如果你想自定义裁剪后的图像文件名、保存路径等参数,可以在`detect.py`中进行修改。 ### 回答2: YOLOv5的save-crop功能是指在目标检测过程中,将检测到的目标所在的区域裁剪出来并保存成独立的图像文件。这个功能可以帮助我们更方便地提取和保存目标区域,以便后续分析或使用。 在YOLOv5中,使用save-crop功能可以通过在命令行中添加--save-crop参数来启用。启用该功能后,在检测到目标后,会将目标的位置信息和对应的图像通过指定的输出文件夹保存成裁剪后的图片文件。我们也可以通过修改源代码自定义输出文件夹的路径。 使用save-crop功能可以帮助我们进行一系列的后处理操作。例如,当需要进一步对检测到的目标进行分类、分析或其他处理时,我们可以通过save-crop先提取出目标区域,然后将这些裁剪后的图像输入到其他模型或算法中进行处理,从而提高后续任务的准确性和效率。 此外,save-crop功能还可以用于数据集的制作。在进行目标检测数据集标注时,我们可以使用这个功能将目标区域剪切下来并保存成单独的图像文件,作为标注的一部分。这样,我们可以更方便地获取高质量的目标区域图像,并用于后续的数据集训练和测试中。 总之,YOLOv5中的save-crop功能提供了一种便捷的方式来保存裁剪后的目标区域图像。它具有许多实际应用的潜力,可以帮助我们进行数据处理、特征提取、后处理等一系列任务,并提高目标检测任务的准确性和效率。 ### 回答3: YOLOv5是一种流行的目标检测算法,而save-crop是其其中一个重要的功能。save-crop是指在YOLOv5模型训练或推理过程中,将检测到的目标从原始图像中剪切出来,并保存为独立的图像文件。 save-crop的使用具有以下几个步骤: 1. 配置模型参数:在使用YOLOv5进行训练或推理之前,需要设置save-crop选项为true,以指示算法需要保存剪切后的目标图像。 2. 目标检测:模型会对输入图像进行目标检测,找出其中的目标物体。 3. 目标裁剪:当检测到目标后,save-crop会自动将这些目标从原始图像中剪切出来,并保存为独立的图像文件。每个目标都会保存为一个独立的图像文件。 4. 图像保存:被裁剪出的目标图像会被保存在指定的目录中,通常是模型所在的文件夹中。 通过使用save-crop功能,可以方便地将检测到的目标物体保存下来,这对于一些特定的应用场景非常有用。例如,对于目标跟踪或者目标分割等任务,可以使用保存的目标图像进行进一步的分析和处理。 需要注意的是,为了成功使用save-crop功能,需要确保模型已经经过了训练,以便正确地检测和裁剪目标。此外,应该指定一个合适的保存目录,以便保存剪切后的目标图像。 总而言之,YOLOv5中的save-crop功能可以将检测到的目标物体从原始图像中剪切出来,并保存为独立的图像文件,提供了更多方便快捷的目标处理选项。

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值