data_juicer.ops.mapper.image_detection_yolo_mapper module

class data_juicer.ops.mapper.image_detection_yolo_mapper.ImageDetectionYoloMapper(imgsz=640, conf=0.05, iou=0.5, model_path='yolo11n.pt', *args, **kwargs)[源代码]

基类:Mapper

Perform object detection using YOLO on images and return bounding boxes and class labels.

This operator uses a YOLO model to detect objects in images. It processes each image in the sample, returning the bounding boxes and class labels for detected objects. The operator sets the bbox_tag and class_label_tag fields in the sample's metadata. If no image is present or no objects are detected, it sets bbox_tag to an empty array and class_label_tag to -1. The operator uses a confidence score threshold and IoU (Intersection over Union) score threshold to filter detections.

__init__(imgsz=640, conf=0.05, iou=0.5, model_path='yolo11n.pt', *args, **kwargs)[源代码]

Initialization method.

参数:
  • imgsz -- resolution for image resizing

  • conf -- confidence score threshold

  • iou -- IoU (Intersection over Union) score threshold

  • model_path -- the path to the YOLO model.

process_single(sample, rank=None, context=False)[源代码]

For sample level, sample --> sample

参数:

sample -- sample to process

返回:

processed sample