YOLOV5 模型和代码修改——针对小目标识别

mulpicplus = "3"     #1 for normal,2 for 4pic plus,3 for 9pic plus and so on

        assert(int(mulpicplus)>=1)

        if mulpicplus == "1":

            pred = model(img,

                          augment=augment,

                          visualize=increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False)[0]

        else:

            xsz = img.shape[2]

            ysz = img.shape[3]

            mulpicplus = int(mulpicplus)

            x_smalloccur = int(xsz / mulpicplus * 1.2)

            y_smalloccur = int(ysz / mulpicplus * 1.2)

            for i in range(mulpicplus):

                x_startpoint = int(i * (xsz / mulpicplus))

                for j in range(mulpicplus):

                    y_startpoint = int(j * (ysz / mulpicplus))

                    x_real = min(x_startpoint + x_smalloccur, xsz)

                    y_real = min(y_startpoint + y_smalloccur, ysz)

                    if (x_real - x_startpoint) % 64 != 0:

                        x_real =外汇跟单gendan5.com x_real - (x_real-x_startpoint) % 64

                    if (y_real - y_startpoint) % 64 != 0:

                        y_real = y_real - (y_real - y_startpoint) % 64

                    dicsrc = img[:, :, x_startpoint:x_real,

                                                    y_startpoint:y_real]

                    pred_temp = model(dicsrc,

                                    augment=augment,

                                    visualize=increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False)[0]

                    pred_temp[..., 0] = pred_temp[..., 0] + y_startpoint

                    pred_temp[..., 1] = pred_temp[..., 1] + x_startpoint

                    if i==0 and j == 0:

                        pred = pred_temp

                    else:

                        pred = torch.cat([pred, pred_temp], dim=1)

        # Apply NMS

        pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)