yoloV5核心代码

以下是YOLOv5的核心代码，包括模型定义和推断：

模型定义：

class YOLOv5(nn.Module):
    def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None):  # model, input channels, number of classes
        super(YOLOv5, self).__init__()

        # Model
        self.model, self.save = parse_model(deepcopy(model), ch=[ch], nc=nc)  # model, savelist = model(items)

        # Save and reduce
        self.nc, self.nde = nc, next(i for i, x in enumerate(
            reversed([getattr(self.model, 'stride', [None])[-1] if hasattr(self.model, 'stride') else None for _ in
                      range(2)])))  # number of classes, depth
        self.stride = int(2 ** self.nde)
        self.names = ['P{}'.format(i) for i in range(self.nde + 1)]
        self.names.append('anchor')
        self.names.append('stride')
        self.names.append('indices')
        self.names.append('nl')
        self.names.append('nc')
        self.names.append('version')

        # Anchors
        self.nl = len(self.model)
        self.na = self.model[self.nl - 1].na
        self.no = self.na * (self.nc + 5)

        # Detect
        self.detect = Detect(self.nc, anchors=self.model[self.nl - 1].anchor)

    def forward(self, x):
        # x = self.forward_features(x)
        # x = self.forward_anchors(x)
        # x = self.forward_detect(x)
        return self.detect(self.model(x))

推断：

def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False):
    """Performs Non-Maximum Suppression (NMS) on inference results

    Returns:
         detections with shape: nx6 (x1, y1, x2, y2, conf, cls)
    """

    # Settings
    xc = prediction[..., 4] > conf_thres  # candidates
    min_wh, max_wh = 2, 4096  # (pixels) minimum and maximum box width and height
    max_det = 300  # maximum number of detections per image
    time_limit = 10.0  # seconds to quit after
    redundant = True  # require redundant detections
    multi_label = classes is not None  # multiple labels per box (adds 0.5ms/img)

    t = time.time()
    output = [torch.zeros((0, 6), device=prediction.device)] * prediction.shape[0]
    for xi, x in enumerate(prediction):  # image index, image inference
        # Apply constraints
        x = x[xc[xi]]  # confidence
        x[:, :4] = clip_coords(x[:, :4], (height, width))
        if multi_label:
            i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T
            x = torch.cat((x[i], x[j]), 0) if i.numel() and j.numel() else x[i + j]
        else:  # best class only
            conf, j = x[:, 5:].max(1, keepdim=True)
            x = torch.cat((x[:, :5], conf, j.float()), 1)[conf.view(-1) > conf_thres]

        # If none remain process next image
        n = x.shape[0]  # number of boxes
        if not n:
            continue

        # Sort by confidence
        x = x[x[:, 4].argsort(descending=True)]

        # Batched NMS
        c = x[:, 5:6] * (0 if agnostic else max_wh)  # classes
        boxes, scores = x[:, :4] + c, x[:, 4]  # boxes (offset by class), scores
        i = torchvision.ops.boxes.nms(boxes, scores, iou_thres)
        if i.shape[0] > max_det:  # limit detections
            i = i[:max_det]
        if redundant:  # redundant detections
            j = torchvision.ops.boxes.box_iou(boxes[i], boxes).view(-1, n) > iou_thres
            i = i[j.sum(1) == 1]

        # Append detections
        output[xi] = torch.cat((x[i], boxes[i]), 1)

        # Break if time limit exceeded
        if (time.time() - t) > time_limit:
            break  # time limit exceeded

    return output if len(output) > 1 else output[0]