python环境

版本要求

目的

torch

>= 1.8.0

搭建、训练网络并进行预测

torchvision

>= 0.9.0

用来生成图片数据集和加载一些预训练模型

torchsummary

>= 1.5.1

输出、可视化网络结果

TensorboardX

>= 2.5.1

网络结构、训练过程参数可视化

PyQt

>= 5.15.7

构建软件界面

Matplotlib

>=3.2.2

数据可视化，二维、三维绘图

numpy

>=1.18.5

提供数据结构以及相应高效的处理函数

Opencv-python

>=4.1.2

基于开源发行的跨平台计算机视觉库

pillow

==8.4.0

用于图像的基本处理

pyyaml

>=5.3.1

Python的YAML解析器和生成器。

requests

>=2.23.0

用作发送网络请求

scipy

>=1.4.1

高级的科学计算库

tqdm

>=4.41.0

方便且易于扩展的Python进度条

Tensorboard

>=2.11.1

用于数据可视化的工具

pandas

>=1.1.4

用于管理数据集，数据分析和探索工具

seaborn

>=0.11.0

用于生成统计图形

pycocotools-windows

>=2.0.0

大型的图像数据集

层数

form

moudule

arguments

input

output

0

-1

Conv

[3, 32, 6, 2, 2]

[3, 640, 640]

[32, 320, 320]

1

-1

Conv

[32, 64, 3, 2]

[32, 320, 320]

[64, 160, 160]

2

-1

C3

[64, 64, 1]

[64, 160, 160]

[64, 160, 160]

3

-1

Conv

[64, 128, 3, 2]

[64, 160, 160]

[128, 80, 80]

4

-1

C3

[128, 128, 2]

[128, 80, 80]

[128, 80, 80]

5

-1

Conv

[128, 256, 3, 2]

[128, 80, 80]

[256, 40, 40]

6

-1

C3

[256, 256, 3]

[256, 40, 40]

[256, 40, 40]

7

-1

Conv

[256, 512, 3, 2]

[256, 40, 40]

[512, 20, 20]

8

-1

C3

[512, 512, 1]

[512, 20, 20]

[512, 20, 20]

9

-1

SPPF

[512, 512, 5]

[512, 20, 20]

[512, 20, 20]

10

-1

Conv

[512, 256, 1, 1]

[512, 20, 20]

[256, 20, 20]

11

-1

Upsample

[None, 2, 'nearest']

[256, 20, 20]

[256, 40, 40]

12

[-1, 6]

Concat

[1]

[1, 256, 40, 40],[1, 256, 40, 40]

[512, 40, 40]

13

-1

C3

[512, 256, 1, False]

[512, 40, 40]

[256, 40, 40]

14

-1

Conv

[256, 128, 1, 1]

[256, 40, 40]

[128, 40, 40]

15

-1

Upsample

[None, 2, 'nearest']

[128, 40, 40]

[128, 80, 80]

16

[-1, 4]

Concat

[1]

[1, 128, 80, 80],[1, 128, 80, 80]

[256, 80, 80]

17

-1

C3

[256, 128, 1, False]

[256, 80, 80]

[128, 80, 80]

18

-1

Conv

[128, 128, 3, 2]

[128, 80, 80]

[128, 40, 40]

19

[-1, 14]

Concat

[1]

[1, 128, 40, 40],[1, 128, 40, 40]

[256, 40, 40]

20

-1

C3

[256, 256, 1, False]

[256, 40, 40]

[256, 40, 40]

21

-1

Conv

[256, 256, 3, 2]

[256, 40, 40]

[256, 20, 20]

22

[-1, 10]

Concat

[1]

[1, 256, 20, 20],[1, 256, 20, 20]

[512, 20, 20]

23

-1

C3

[512, 512, 1, False]

[512, 20, 20]

[512, 20, 20]

24

[17, 20, 23]

Detect

[80, [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]], [128, 256, 512]]

[1, 128, 80, 80],[1, 256, 40, 40],[1, 512, 20, 20]

[1, 3, 80, 80, 85],[1, 3, 40, 40, 85],[1, 3, 20, 20, 85]

class Model(nn.Module):

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

        super().__init__()

        if isinstance(cfg, dict):

            self.yaml = cfg  # model dict

        else:  # is *.yaml

            import yaml  # for torch hub

            self.yaml_file = Path(cfg).name

            with open(cfg, encoding='ascii', errors='ignore') as f:

                self.yaml = yaml.safe_load(f)  # model dict

        # Define model

        ch = self.yaml['ch'] = self.yaml.get('ch', ch)  # input channels

        if nc and nc != self.yaml['nc']:

            LOGGER.info(f"Overriding model.yaml nc={self.yaml['nc']} with nc={nc}")

            self.yaml['nc'] = nc  # override yaml value

        if anchors:

            LOGGER.info(f'Overriding model.yaml anchors with anchors={anchors}')

            self.yaml['anchors'] = round(anchors)  # override yaml value

        self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch])  # model, savelist

        self.names = [str(i) for i in range(self.yaml['nc'])]  # default names

        self.inplace = self.yaml.get('inplace', True)

        # Build strides, anchors

        m = self.model[-1]  # Detect()

        if isinstance(m, Detect):

            s = 256  # 2x min stride

            m.inplace = self.inplace

            m.stride = torch.tensor([s / x.shape[-2] for x in self.forward(torch.zeros(1, ch, s, s))])  # forward

            m.anchors /= m.stride.view(-1, 1, 1)

            check_anchor_order(m)

            self.stride = m.stride

            self._initialize_biases()  # only run once

        # Init weights, biases

        initialize_weights(self)

        self.info()

        LOGGER.info('')

    def forward(self, x, augment=False, profile=False, visualize=False):

        if augment:

            return self._forward_augment(x)  # augmented inference, None

        return self._forward_once(x, profile, visualize)  # single-scale inference, train

    def _forward_augment(self, x):

        img_size = x.shape[-2:]  # height, width

        s = [1, 0.83, 0.67]  # scales

        f = [None, 3, None]  # flips (2-ud, 3-lr)

        y = []  # outputs

        for si, fi in zip(s, f):

            xi = scale_img(x.flip(fi) if fi else x, si, gs=int(self.stride.max()))

            yi = self._forward_once(xi)[0]  # forward

            # cv2.imwrite(f'img_{si}.jpg', 255 * xi[0].cpu().numpy().transpose((1, 2, 0))[:, :, ::-1])  # save

            yi = self._descale_pred(yi, fi, si, img_size)

            y.append(yi)

        y = self._clip_augmented(y)  # clip augmented tails

        return torch.cat(y, 1), None  # augmented inference, train

    def _forward_once(self, x, profile=False, visualize=False):

        y, dt = [], []  # outputs

        for m in self.model:

            if m.f != -1:  # if not from previous layer

                x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f]  # from earlier layers

            if profile:

                self._profile_one_layer(m, x, dt)

            x = m(x)  # run

            y.append(x if m.i in self.save else None)  # save output

            if visualize:

                feature_visualization(x, m.type, m.i, save_dir=visualize)

        return x

for epoch in range(start_epoch, epochs):  # epoch

    model.train()

    if opt.image_weights:

        cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc  # class weights

        iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw)  

        dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n)  

    mloss = torch.zeros(3, device=device)  # mean losses

    if RANK != -1:

        train_loader.sampler.set_epoch(epoch)

    pbar = enumerate(train_loader)

    LOGGER.info(('\n' + '%10s' * 7) % ('Epoch', 'gpu_mem', 'box', 'obj', 'cls', 'labels', 'img_size'))

    if RANK in [-1, 0]:

        pbar = tqdm(pbar, total=nb, ncols=NCOLS, bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}')  # progress bar

    optimizer.zero_grad()

    for i, (imgs, targets, paths, _) in pbar:  # batch

        ni = i + nb * epoch  # number integrated batches (since train start)

        imgs = imgs.to(device, non_blocking=True).float() / 255   

        if ni <= nw:

            xi = [0, nw]  # x interp

            accumulate = max(1, np.interp(ni, xi, [1, nbs / batch_size]).round())

            for j, x in enumerate(optimizer.param_groups):

                x['lr'] = np.interp(ni, xi, [hyp['warmup_bias_lr'] if j == 2 else 0.0, x['initial_lr'] * lf(epoch)])

                if 'momentum' in x:

                    x['momentum'] = np.interp(ni, xi, [hyp['warmup_momentum'], hyp['momentum']])

        if opt.multi_scale:

            sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs  # size

            sf = sz / max(imgs.shape[2:])  # scale factor

            if sf != 1:

                ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]]                  imgs = nn.functional.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)

        # Forward

        with amp.autocast(enabled=cuda):

            pred = model(imgs)  # forward

            loss, loss_items = compute_loss(pred, targets.to(device))             if RANK != -1:

                loss *= WORLD_SIZE  # gradient averaged between devices in DDP mode

            if opt.quad:

                loss *= 4.

        # Backward

        scaler.scale(loss).backward()

        # Optimize

        if ni - last_opt_step >= accumulate:

            scaler.step(optimizer)  # optimizer.step

            scaler.update()

            optimizer.zero_grad()

            if ema:

                ema.update(model)

            last_opt_step = ni

        # Log

        if RANK in [-1, 0]:

            mloss = (mloss * i + loss_items) / (i + 1)  # update mean losses

            mem = f'{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G'  # (GB)

            pbar.set_description(('%10s' * 2 + '%10.4g' * 5) % (

                f'{epoch}/{epochs - 1}', mem, *mloss, targets.shape[0], imgs.shape[-1]))

            callbacks.run('on_train_batch_end', ni, model, imgs, targets, paths, plots, opt.sync_bn)

    # Scheduler

    lr = [x['lr'] for x in optimizer.param_groups]  # for loggers

    scheduler.step()

def run(weights=ROOT / 'yolov5s.pt',  # model.pt path(s)

        source=ROOT / 'data/images',  # file/dir/URL/glob, 0 for webcam

        imgsz=640,  # inference size (pixels)

        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_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

        ):

    source = str(source)

    save_img = not nosave and not source.endswith('.txt')  # save inference images

    is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)

    is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))

    webcam = source.isnumeric() or source.endswith('.txt') or (is_url and not is_file)

    if is_url and is_file:

        source = check_file(source)  # download

    # Directories

    save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)      (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir

    # Load model

    device = select_device(device)

    model = DetectMultiBackend(weights, device=device, dnn=dnn)

    stride, names, pt, jit, onnx = model.stride, model.names, model.pt, model.jit, model.onnx

    imgsz = check_img_size(imgsz, s=stride)  # check image size

    # Half

    half &= pt and device.type != 'cpu'  # half precision only supported by PyTorch on CUDA

    if pt:

        model.model.half() if half else model.model.float()

    # Dataloader

    if webcam:

        view_img = check_imshow()

        cudnn.benchmark = True  # set True to speed up constant image size inference

        dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt and not jit)

        bs = len(dataset)  # batch_size

    else:

        dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt and not jit)

        bs = 1  # batch_size

    vid_path, vid_writer = [None] * bs, [None] * bs

    # Run inference

    if pt and device.type != 'cpu':

        model(torch.zeros(1, 3, *imgsz).to(device).type_as(next(model.model.parameters())))  # warmup

    dt, seen = [0.0, 0.0, 0.0], 0

    for path, im, im0s, vid_cap, s in dataset:

        t1 = time_sync()

        im = torch.from_numpy(im).to(device)

        im = im.half() if half else im.float()  # uint8 to fp16/32

        im /= 255  # 0 - 255 to 0.0 - 1.0

        if len(im.shape) == 3:

            im = im[None]  # expand for batch dim

        t2 = time_sync()

        dt[0] += t2 - t1

        # Inference

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

        pred = model(im, augment=augment, visualize=visualize)

        t3 = time_sync()

        dt[1] += t3 - t2

        # NMS

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

        dt[2] += time_sync() - t3

        # Process predictions

        for i, det in enumerate(pred):  # per image

            seen += 1

            if webcam:  # batch_size >= 1

                p, im0, frame = path[i], im0s[i].copy(), dataset.count

                s += f'{i}: '

            else:

                p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)

            p = Path(p)  # to Path

            save_path = str(save_dir / p.name)  # im.jpg

            txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')  # im.txt

            s += '%gx%g ' % im.shape[2:]  # print string

            gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh

            imc = im0.copy() if save_crop else im0  # for save_crop

            annotator = Annotator(im0, line_width=line_thickness, example=str(names))

            if len(det):

                det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()

                # Print results

                for c in det[:, -1].unique():

                    n = (det[:, -1] == c).sum()  # detections per class

                    s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string

                # Write results

                for *xyxy, conf, cls in reversed(det):

                    if save_txt:  # Write to file

                        xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh

                        line = (cls, *xywh, conf) if save_conf else (cls, *xywh)

                        with open(txt_path + '.txt', 'a') as f:

                            f.write(('%g ' * len(line)).rstrip() % line + '\n')

                    if save_img or save_crop or view_img:  # Add bbox to image

                        c = int(cls)  # integer class

                        label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')

                        annotator.box_label(xyxy, label, color=colors(c, True))

                        if save_crop:

                            save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)

            # Print time (inference-only)

            LOGGER.info(f'{s}Done. ({t3 - t2:.3f}s)')

            # Stream results

            im0 = annotator.result()

            if view_img:

                cv2.imshow(str(p), im0)

               cv2.waitKey(1)  # 1 millisecond

            # Save results (image with detections)

            if save_img:

                if dataset.mode == 'image':

                    cv2.imwrite(save_path, im0)

                else:  # 'video' or 'stream'

                    if vid_path[i] != save_path:  # new video

                        vid_path[i] = save_path

                        if isinstance(vid_writer[i], cv2.VideoWriter):

                            vid_writer[i].release()  # release previous video writer

                        if vid_cap:  # video

                            fps = vid_cap.get(cv2.CAP_PROP_FPS)

                            w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))

                            h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

                        else:  # stream

                            fps, w, h = 30, im0.shape[1], im0.shape[0]

                            save_path += '.mp4'

                        vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))

                    vid_writer[i].write(im0)

    # Print results

    t = tuple(x / seen * 1E3 for x in dt)  # speeds per image

    LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t)

    if save_txt or save_img:

        s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''

        LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")

    if update:

        strip_optimizer(weights)  # update model (to fix SourceChangeWarning)