├── .idea
├── .gitignore
├── DRENet.iml
├── deployment.xml
├── inspectionProfiles
│ ├── Project_Default.xml
│ └── profiles_settings.xml
├── misc.xml
├── modules.xml
└── vcs.xml
├── DegradeGenerate.py
├── LICENSE
├── NetworkStructure.png
├── README.md
├── data
├── hyp.finetune.yaml
├── hyp.scratch.yaml
└── ship.yaml
├── detect.py
├── hubconf.py
├── models
├── Ablation_Only_CRMA.yaml
├── Ablation_Only_CRMA_Remove_PAN.yaml
├── Ablation_Only_CRMA_and_Scale.yaml
├── Ablation_Only_Detector.yaml
├── Ablation_Only_Enhancer.yaml
├── Ablation_Only_Scale.yaml
├── DRENet.yaml
├── __init__.py
├── common.py
├── experimental.py
├── export.py
├── yolo.py
└── yolov5s.yaml
├── requirements.txt
├── test.py
├── train.py
└── utils
├── __init__.py
├── activations.py
├── autoanchor.py
├── datasets.py
├── general.py
├── google_utils.py
├── loss.py
├── metrics.py
├── plots.py
└── torch_utils.py
/.idea/.gitignore:
--------------------------------------------------------------------------------
1 | # Default ignored files
2 | /shelf/
3 | /workspace.xml
4 | # Datasource local storage ignored files
5 | /dataSources/
6 | /dataSources.local.xml
7 | # Editor-based HTTP Client requests
8 | /httpRequests/
9 |
--------------------------------------------------------------------------------
/.idea/DRENet.iml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
--------------------------------------------------------------------------------
/.idea/deployment.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
9 |
10 |
11 |
12 |
13 |
14 |
15 |
16 |
17 |
18 |
19 |
20 |
21 |
22 |
23 |
24 |
25 |
26 |
27 |
28 |
29 |
30 |
31 |
32 |
33 |
34 |
35 |
36 |
37 |
38 |
39 |
40 |
41 |
42 |
--------------------------------------------------------------------------------
/.idea/inspectionProfiles/Project_Default.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
15 |
16 |
17 |
--------------------------------------------------------------------------------
/.idea/inspectionProfiles/profiles_settings.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
--------------------------------------------------------------------------------
/.idea/misc.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
--------------------------------------------------------------------------------
/.idea/modules.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 |
8 |
--------------------------------------------------------------------------------
/.idea/vcs.xml:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
--------------------------------------------------------------------------------
/DegradeGenerate.py:
--------------------------------------------------------------------------------
1 | import warnings
2 | import skimage
3 | from skimage import io
4 | import os
5 | from glob import glob
6 | from tqdm import tqdm
7 | import numpy as np
8 | import cv2
9 | from numba import jit
10 | import concurrent.futures
11 |
12 | warnings.filterwarnings('ignore')
13 |
14 | # @jit(nopython=True)
15 | def customBlur(imgFile):
16 | size = 512
17 | targetPath = r'.\train\degrade' # save path
18 | labelFile = imgFile.replace('png', 'txt').replace('images', 'labels')
19 | if os.path.exists(labelFile):
20 | img = io.imread(imgFile)
21 | label = np.loadtxt(labelFile, ndmin=2)
22 | dst = img.copy()
23 | centers = label.copy()
24 | for j in range(len(centers)):
25 | centers[j, 1] = label[j, 1] * size
26 | centers[j, 2] = label[j, 2] * size
27 | for i in range(img.shape[0]):
28 | for j in range(img.shape[1]):
29 | minDis = 130 * 130
30 | for center in centers:
31 | distance = (i - center[2]) ** 2 + (j - center[1]) ** 2
32 | if distance < minDis:
33 | minDis = distance
34 | # boxSize = (int(0.05* (minDis ** 0.5))) // 2 + 1
35 | boxSize = (int(1.03** (minDis ** 0.5))) // 2 # You can change the degradation configuration here
36 | # boxSize = int(((minDis**0.5)/5+1))//2
37 | dst[i, j,0] = img[max(i - boxSize, 0):min(i + boxSize + 1, size), max(j - boxSize, 0):min(j + boxSize + 1, size),0].mean()
38 | dst[i, j, 1] = img[max(i - boxSize, 0):min(i + boxSize + 1, size),
39 | max(j - boxSize, 0):min(j + boxSize + 1, size), 1].mean()
40 | dst[i, j, 2] = img[max(i - boxSize, 0):min(i + boxSize + 1, size),
41 | max(j - boxSize, 0):min(j + boxSize + 1, size), 2].mean()
42 | io.imsave(os.path.join(targetPath, os.path.basename(imgFile)), dst)
43 | else:
44 | img = cv2.imread(imgFile)
45 | dst = cv2.blur(img, (20, 20))
46 | cv2.imwrite(os.path.join(targetPath, os.path.basename(imgFile)), dst)
47 |
48 | if __name__=='__main__':
49 | sourcePath = r'.\train\images' # source path
50 | with concurrent.futures.ProcessPoolExecutor(1) as executor: # set 1 to other number for speedup
51 | imgfiles=glob(os.path.join(sourcePath, '*.png'))
52 | for i in tqdm(zip(imgfiles, executor.map(customBlur, imgfiles)),total=1):
53 | pass
54 |
--------------------------------------------------------------------------------
/NetworkStructure.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/WindVChen/DRENet/a187dbe0f623b521a62c6176c7cafaa7322f5f66/NetworkStructure.png
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # DRENet: Fast and Accurate Tiny Ship detection Method
2 |
3 | 
4 | 
5 | [](#License)
6 |
7 | ## Share us a :star: if this repo does help
8 |
9 | This repo is the official implementation of the DRENet in "A Degraded Reconstruction Enhancement-based Method for Tiny Ship Detection in Remote Sensing Images with A New Large-scale Dataset". The paper can be accessed in [[IEEE](https://ieeexplore.ieee.org/document/9791363) | [Lab Server](http://levir.buaa.edu.cn/publications/DRENet.pdf) | [ResearchGate](https://www.researchgate.net/profile/Keyan-Chen-6/publication/361178478_A_Degraded_Reconstruction_Enhancement-based_Method_for_Tiny_Ship_Detection_in_Remote_Sensing_Images_with_A_New_Large-scale_Dataset/links/62f47d69b8dc8b4403d4ce5e/A-Degraded-Reconstruction-Enhancement-Based-Method-for-Tiny-Ship-Detection-in-Remote-Sensing-Images-With-a-New-Large-Scale-Dataset.pdf)]. ***(Accepted by TGRS 2022)***
10 |
11 | If you encounter any question, please feel free to contact us. You can create an issue or just send email to me windvchen@gmail.com. Also welcome for any idea exchange and discussion.
12 |
13 | ## Updates
14 |
15 | ***05/30/2023***
16 |
17 | Thanks for the discussions in [#10](https://github.com/WindVChen/DRENet/issues/10). The current code now supports processing different resolutions without manually modifying the yaml configurations. To achieve this, you can comment out L139-141 in [common.py](./models/common.py). (Note that this feature has not well been tested, if there is a drop in performance, you can get back to the manual modification mentioned in [#4](https://github.com/WindVChen/DRENet/issues/4)).
18 |
19 | ***11/19/2022***
20 |
21 | We have improved the previous codes, and the current repository support directly detect the images without generating degraded images, which correspond to the file **detect.py**.
22 |
23 | ***06/10/2022***
24 |
25 | The code cleanup is finished and the complete codes are provided, also the weights of our model on LEVIR-Ship dataset.
26 |
27 | ***06/06/2022***
28 |
29 | We will finish the code cleanup within a week, and make both the code and dataset fully public. Please be patient.
30 |
31 |
32 | ## Table of Contents
33 |
34 | - [Introduction](#Introduction)
35 | - [Results and Trained Model](#Results-and-Trained-Model)
36 | - [Preliminaries](#Preliminaries)
37 | - [Environments](#Environments)
38 | - [Run Details](#Run-Details)
39 | - [Train Process](#Train-Process)
40 | - [Valid Process](#Valid-Process)
41 | - [Detect Process](#Detect-Process)
42 | - [Visualization](#Visualization)
43 | - [Citation](#Citation)
44 | - [License](#License)
45 |
46 | ## Introduction
47 | 
48 |
49 | We focus on **tiny ship detection** task in **medium-resolution (MR, about 16m/pixel)** remote sensing (RS) images . Compared with high-resolution (HR) RS image, an MR image covers a much wider area, thus facilitating quick ship detection. This direction has **great research significance**, and can greatly benefit the **rapid ship detection** under massive RS images.
50 |
51 | For the task, we propose an effective **Degraded Reconstruction Enhancement Network (DRENet)**, where a degraded reconstruction enhancer is designed to learn to regress an object-aware blurred version of the input image. Our method achieves both **great effectiveness and efficiency**, and outperforms many recent methods.
52 |
53 | ## Results and Trained Model
54 | ### Models trained on LEVIR-Ship dataset
55 | | Methods | Params(M) |FLOPs(G) | AP | FPS |
56 | |:---|:---:|:---:|:---:| :---:|
57 | | YOLOv3 | 61.52 | 99.2 | 69.9 | 61 |
58 | | YOLOv5s | 7.05 | 10.4 | 75.6 [[Google Drive](https://drive.google.com/file/d/10AQA_ynjvmVD8XSiOhM_9A64EDm5lxfS/view?usp=sharing)
| [Baidu Pan](https://pan.baidu.com/s/1AffKx_gChABQiicJv2zjtg) (code:ogdm)] | **95** |
59 | | Retinanet | 36.33 | 104.4 | 74.9 | 12 |
60 | | SSD | 24.39 | 175.2 | 52.6 | 25 |
61 | | FasterRCNN | 136.70 | 299.2 | 70.8 | 10 |
62 | | EfficientDet-D0 | **3.84** | **4.6** | 71.3 | 32 |
63 | | EfficientDet-D2 | 8.01 | 20.0 | 80.9 | 21 |
64 | | FCOS | 5.92 | 51.8 | 75.5 | 37 |
65 | | CenterNet | 191.24 | 584.6 | 77.7 | 25 |
66 | | HSFNet | 157.59 | 538.1 | 73.6 | 7 |
67 | | ImYOLOv3 | 62.86 | 101.9 | 72.6 | 51 |
68 | | MaskRCNN+DFR+RFE | 24.99 | 237.8 | 76.2 | 6 |
69 | | **DRENet** | 4.79 | 8.3 | **82.4** [[Google Drive](https://drive.google.com/file/d/1ApAejwSNYQDvROM1yRtltQOGdYAwYyF3/view?usp=sharing)
| [Baidu Pan](https://pan.baidu.com/s/1tBxhGOhmxc-L5ioHSqSjEQ) (code:x710)] | 85|
70 |
71 |
72 | ## Preliminaries
73 | Please at first download dataset [LEVIR-Ship](https://github.com/WindVChen/LEVIR-Ship), then prepare the dataset as the following structure:
74 | ```
75 | ├── train
76 | ├── images
77 | ├── img_1.png
78 | ├── img_2.png
79 | ├── ...
80 | ├── degrade
81 | # images processed by Selective Degradation (refer to our paper for detals)
82 | ├── degraded_img_1.png
83 | ├── degraded_img_2.png
84 | ├── ...
85 | ├── labels
86 | ├── label_1.txt
87 | ├── label_2.txt
88 | ├── ...
89 | ├── val
90 | ├── test
91 | ```
92 | Note that apart from the images and labels in LEVIR-Ship dataset, you should also generate the **degraded images**, which are the supervision of the enhancer (see details in our paper). Here, we provide [DegradeGenerate.py](DegradeGenerate.py) to easily generate the degraded images.
93 |
94 | After preparing the dataset as above, change the paths in [ship.yaml](data/ship.yaml).
95 |
96 | ***(The partitioned dataset, including the degraded images, can all be accessed [here](https://github.com/WindVChen/LEVIR-Ship))***
97 |
98 | ## Environments
99 |
100 | - Windows/Linux Support
101 | - python 3.8
102 | - pytorch 1.9.0
103 | - torchvision
104 | - wandb (Suggested, a good tool to visualize the training process. If not want to use it, you should comment out the related codes.)
105 | - ...... (See more details in [requirements.txt](requirements.txt))
106 |
107 | *(The code is constructed based on YOLOv5s, for more details about YOLOv5, please refer to their repo [here](https://github.com/ultralytics/yolov5).)*
108 |
109 |
110 | ## Run Details
111 | ### Train Process
112 | To train our `DRENet`, run:
113 | ```
114 | python train.py --cfg "./models/DRENet.yaml" --epochs 1000 --workers 8 --batch-size 16 --device 0 --project "./LEVIR-Ship" --data "./data/ship.yaml"
115 | ```
116 | **Parameters Description**
117 | - `cfg`: You can change it to use different network structures. More structure configurations can be found in [models](models) directory, where we provide the baseline YOLOv5s, and the ablation structures of DRENet. You can try them if you are interested
118 | - `epochs`: A longer training time is suggested, and 1,000 epochs are enough.
119 | - `project`: The path where you want to save your experiments. Also the name of the project in wandb.
120 |
121 | **Others**
122 |
123 | The current codes use **fixed weight balance**, which can also achieve a good result.
124 |
125 | If you want to make use of **automatic weight balance**, please search the key word `weightOptimizer` in [train.py](train.py) and uncomment the code lines, also the code lines with the key word `ForAuto` in [loss.py](utils/loss.py) be uncommented and the other lines be commented out.
126 |
127 | ### Valid Process
128 |
129 | To evaluate our `DRENet`, you should first train the network or download [our provided weights](#Models-trained-on-LEVIR-Ship-dataset), then run:
130 | ```
131 | python test.py --weights "./DRENet.pt" --project "runs/test" --device 0 --batch-size 16 --data "./data/ship.yaml"
132 | ```
133 | You can set how many detected results to plot by changing the value of `plot_batch_num` in [test.py](test.py). Also ensure that you have changed the val path in [ship.yaml](data/ship.yaml) into your test path.
134 |
135 | Please ensure that there are corresponding degraded images in the **degrade** folder. (See [#issue 4](https://github.com/WindVChen/DRENet/issues/4) for more details.)
136 |
137 | ### Detect Process
138 | To directly output the detect results **without the need of the degraded images**, please run the following command:
139 | ```
140 | python detect.py --weights "./DRENet.pt" --source "images/" --device 0
141 | ```
142 | where "--source" is the path that the images need detection in.
143 |
144 |
145 | ## Citation
146 | If you find this paper useful in your research, please consider citing:
147 | ```
148 | @ARTICLE{9791363,
149 | author={Chen, Jianqi and Chen, Keyan and Chen, Hao and Zou, Zhengxia and Shi, Zhenwei},
150 | journal={IEEE Transactions on Geoscience and Remote Sensing},
151 | title={A Degraded Reconstruction Enhancement-based Method for Tiny Ship Detection in Remote Sensing Images with A New Large-scale Dataset},
152 | year={2022},
153 | volume={60},
154 | number={},
155 | pages={1-14},
156 | doi={10.1109/TGRS.2022.3180894}}
157 | ```
158 |
159 |
160 | ## License
161 | This project is licensed under the GPL-3.0 License. See [LICENSE](LICENSE) for details
162 |
--------------------------------------------------------------------------------
/data/hyp.finetune.yaml:
--------------------------------------------------------------------------------
1 | # Hyperparameters for VOC finetuning
2 | # python train.py --batch 64 --weights yolov5m.pt --data voc.yaml --img 512 --epochs 50
3 | # See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials
4 |
5 |
6 | # Hyperparameter Evolution Results
7 | # Generations: 306
8 | # P R mAP.5 mAP.5:.95 box obj cls
9 | # Metrics: 0.6 0.936 0.896 0.684 0.0115 0.00805 0.00146
10 |
11 | lr0: 0.0032
12 | lrf: 0.12
13 | momentum: 0.843
14 | weight_decay: 0.00036
15 | warmup_epochs: 2.0
16 | warmup_momentum: 0.5
17 | warmup_bias_lr: 0.05
18 | box: 0.0296
19 | cls: 0.243
20 | cls_pw: 0.631
21 | obj: 0.301
22 | obj_pw: 0.911
23 | iou_t: 0.2
24 | anchor_t: 2.91
25 | # anchors: 3.63
26 | fl_gamma: 0.0
27 | hsv_h: 0.0138
28 | hsv_s: 0.664
29 | hsv_v: 0.464
30 | degrees: 0.373
31 | translate: 0.245
32 | scale: 0.898
33 | shear: 0.602
34 | perspective: 0.0
35 | flipud: 0.00856
36 | fliplr: 0.5
37 | mosaic: 1.0
38 | mixup: 0.243
39 |
--------------------------------------------------------------------------------
/data/hyp.scratch.yaml:
--------------------------------------------------------------------------------
1 | # Hyperparameters for COCO training from scratch
2 | # python train.py --batch 40 --cfg yolov5m.yaml --weights '' --data coco.yaml --img 640 --epochs 300
3 | # See tutorials for hyperparameter evolution https://github.com/ultralytics/yolov5#tutorials
4 |
5 |
6 | lr0: 0.01 # initial learning rate (SGD=1E-2, Adam=1E-3)
7 | lrf: 0.2 # final OneCycleLR learning rate (lr0 * lrf)
8 | momentum: 0.937 # SGD momentum/Adam beta1
9 | weight_decay: 0.0005 # optimizer weight decay 5e-4
10 | warmup_epochs: 3.0 # warmup epochs (fractions ok)
11 | warmup_momentum: 0.8 # warmup initial momentum
12 | warmup_bias_lr: 0.1 # warmup initial bias lr
13 | box: 0.05 # box loss gain
14 | cls: 0.5 # cls loss gain
15 | cls_pw: 1.0 # cls BCELoss positive_weight
16 | obj: 1.0 # obj loss gain (scale with pixels)
17 | obj_pw: 1.0 # obj BCELoss positive_weight
18 | iou_t: 0.20 # IoU training threshold
19 | anchor_t: 4.0 # anchor-multiple threshold
20 | # anchors: 3 # anchors per output layer (0 to ignore)
21 | fl_gamma: 0.0 # focal loss gamma (efficientDet default gamma=1.5)
22 | hsv_h: 0.015 # image HSV-Hue augmentation (fraction)
23 | hsv_s: 0.7 # image HSV-Saturation augmentation (fraction)
24 | hsv_v: 0.4 # image HSV-Value augmentation (fraction)
25 | degrees: 0.0 # image rotation (+/- deg)
26 | translate: 0.1 # image translation (+/- fraction)
27 | scale: 0.5 # image scale (+/- gain)
28 | shear: 0.0 # image shear (+/- deg)
29 | perspective: 0.0 # image perspective (+/- fraction), range 0-0.001
30 | flipud: 0.0 # image flip up-down (probability)
31 | fliplr: 0.5 # image flip left-right (probability)
32 | mosaic: 1.0 # image mosaic (probability)
33 | mixup: 0.0 # image mixup (probability)
34 |
--------------------------------------------------------------------------------
/data/ship.yaml:
--------------------------------------------------------------------------------
1 | train: ../../ForDRENet/finalDataSet/yolo/train/images/
2 | val: ../../ForDRENet/finalDataSet/yolo/val/images/ # when test, change "val" to "test"
3 |
4 | # number of classes
5 | nc: 1
6 |
7 | # class names
8 | names: [ 'ship' ]
9 |
--------------------------------------------------------------------------------
/detect.py:
--------------------------------------------------------------------------------
1 | import argparse
2 | import time
3 | from pathlib import Path
4 |
5 | import cv2
6 | import torch
7 | import torch.backends.cudnn as cudnn
8 | from numpy import random
9 |
10 | from models.experimental import attempt_load
11 | from utils.datasets import LoadStreams, LoadImages
12 | from utils.general import check_img_size, check_requirements, check_imshow, non_max_suppression, apply_classifier, \
13 | scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path
14 | from utils.plots import plot_one_box
15 | from utils.torch_utils import select_device, load_classifier, time_synchronized
16 |
17 |
18 | def detect(save_img=False):
19 | source, weights, view_img, save_txt, imgsz = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
20 | webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith(
21 | ('rtsp://', 'rtmp://', 'http://'))
22 |
23 | # Directories
24 | save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) # increment run
25 | (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
26 |
27 | # Initialize
28 | set_logging()
29 | device = select_device(opt.device)
30 | half = device.type != 'cpu' # half precision only supported on CUDA
31 |
32 | # Load model
33 | model = attempt_load(weights, map_location=device) # load FP32 model
34 | stride = int(model.stride.max()) # model stride
35 | imgsz = check_img_size(imgsz, s=stride) # check img_size
36 | if half:
37 | model.half() # to FP16
38 |
39 | # Second-stage classifier
40 | classify = False
41 | if classify:
42 | modelc = load_classifier(name='resnet101', n=2) # initialize
43 | modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']).to(device).eval()
44 |
45 | # Set Dataloader
46 | vid_path, vid_writer = None, None
47 | if webcam:
48 | view_img = check_imshow()
49 | cudnn.benchmark = True # set True to speed up constant image size inference
50 | dataset = LoadStreams(source, img_size=imgsz)
51 | else:
52 | save_img = True
53 | dataset = LoadImages(source, img_size=imgsz)
54 |
55 | # Get names and colors
56 | names = model.module.names if hasattr(model, 'module') else model.names
57 | colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]
58 |
59 | # Run inference
60 | if device.type != 'cpu':
61 | model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once
62 | t0 = time.time()
63 | for path, img, im0s, vid_cap in dataset:
64 | img = torch.from_numpy(img).to(device)
65 | img = img.half() if half else img.float() # uint8 to fp16/32
66 | img /= 255.0 # 0 - 255 to 0.0 - 1.0
67 | if img.ndimension() == 3:
68 | img = img.unsqueeze(0)
69 |
70 | # Inference
71 | t1 = time_synchronized()
72 | pred = model(img, augment=opt.augment)[0][0]
73 |
74 | # Apply NMS
75 | pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
76 | t2 = time_synchronized()
77 |
78 | # Apply Classifier
79 | if classify:
80 | pred = apply_classifier(pred, modelc, img, im0s)
81 |
82 | # Process detections
83 | for i, det in enumerate(pred): # detections per image
84 | if webcam: # batch_size >= 1
85 | p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(), dataset.count
86 | else:
87 | p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0)
88 |
89 | p = Path(p) # to Path
90 | save_path = str(save_dir / p.name) # img.jpg
91 | txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # img.txt
92 | s += '%gx%g ' % img.shape[2:] # print string
93 | gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
94 | if len(det):
95 | # Rescale boxes from img_size to im0 size
96 | det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
97 |
98 | # Print results
99 | for c in det[:, -1].unique():
100 | n = (det[:, -1] == c).sum() # detections per class
101 | s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string
102 |
103 | # Write results
104 | for *xyxy, conf, cls in reversed(det):
105 | if save_txt: # Write to file
106 | xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
107 | line = (cls, *xywh, conf) if opt.save_conf else (cls, *xywh) # label format
108 | with open(txt_path + '.txt', 'a') as f:
109 | f.write(('%g ' * len(line)).rstrip() % line + '\n')
110 |
111 | if save_img or view_img: # Add bbox to image
112 | label = f'{names[int(cls)]} {conf:.2f}'
113 | plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=1)
114 |
115 | # Print time (inference + NMS)
116 | print(f'{s}Done. ({t2 - t1:.3f}s)')
117 |
118 | # Stream results
119 | if view_img:
120 | cv2.imshow(str(p), im0)
121 | cv2.waitKey(1) # 1 millisecond
122 |
123 | # Save results (image with detections)
124 | if save_img:
125 | if dataset.mode == 'image':
126 | cv2.imwrite(save_path, im0)
127 | else: # 'video'
128 | if vid_path != save_path: # new video
129 | vid_path = save_path
130 | if isinstance(vid_writer, cv2.VideoWriter):
131 | vid_writer.release() # release previous video writer
132 |
133 | fourcc = 'mp4v' # output video codec
134 | fps = vid_cap.get(cv2.CAP_PROP_FPS)
135 | w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
136 | h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
137 | vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*fourcc), fps, (w, h))
138 | vid_writer.write(im0)
139 |
140 | if save_txt or save_img:
141 | s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
142 | print(f"Results saved to {save_dir}{s}")
143 |
144 | print(f'Done. ({time.time() - t0:.3f}s)')
145 |
146 |
147 | if __name__ == '__main__':
148 | parser = argparse.ArgumentParser()
149 | parser.add_argument('--weights', nargs='+', type=str, default=r'.\DRENet.pt', help='model.pt path(s)')
150 | parser.add_argument('--source', type=str, default='images/', help='source') # file/folder, 0 for webcam
151 | parser.add_argument('--img-size', type=int, default=512, help='inference size (pixels)')
152 | parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
153 | parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
154 | parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
155 | parser.add_argument('--view-img', action='store_true', help='display results')
156 | parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
157 | parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
158 | parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
159 | parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
160 | parser.add_argument('--augment', action='store_true', help='augmented inference')
161 | parser.add_argument('--update', action='store_true', help='update all models')
162 | parser.add_argument('--project', default='runs/detect', help='save results to project/name')
163 | parser.add_argument('--name', default='exp', help='save results to project/name')
164 | parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
165 | opt = parser.parse_args()
166 | print(opt)
167 | check_requirements()
168 |
169 | with torch.no_grad():
170 | if opt.update: # update all models (to fix SourceChangeWarning)
171 | for opt.weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']:
172 | detect()
173 | strip_optimizer(opt.weights)
174 | else:
175 | detect()
176 |
--------------------------------------------------------------------------------
/hubconf.py:
--------------------------------------------------------------------------------
1 | """File for accessing YOLOv5 via PyTorch Hub https://pytorch.org/hub/
2 |
3 | Usage:
4 | import torch
5 | model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True, channels=3, classes=80)
6 | """
7 |
8 | from pathlib import Path
9 |
10 | import torch
11 |
12 | from models.yolo import Model
13 | from utils.general import set_logging
14 | from utils.google_utils import attempt_download
15 |
16 | dependencies = ['torch', 'yaml']
17 | set_logging()
18 |
19 |
20 | def create(name, pretrained, channels, classes, autoshape):
21 | """Creates a specified YOLOv5 model
22 |
23 | Arguments:
24 | name (str): name of model, i.e. 'yolov5s'
25 | pretrained (bool): load pretrained weights into the model
26 | channels (int): number of input channels
27 | classes (int): number of model classes
28 |
29 | Returns:
30 | pytorch model
31 | """
32 | config = Path(__file__).parent / 'models' / f'{name}.yaml' # model.yaml path
33 | try:
34 | model = Model(config, channels, classes)
35 | if pretrained:
36 | fname = f'{name}.pt' # checkpoint filename
37 | attempt_download(fname) # download if not found locally
38 | ckpt = torch.load(fname, map_location=torch.device('cpu')) # load
39 | state_dict = ckpt['model'].float().state_dict() # to FP32
40 | state_dict = {k: v for k, v in state_dict.items() if model.state_dict()[k].shape == v.shape} # filter
41 | model.load_state_dict(state_dict, strict=False) # load
42 | if len(ckpt['model'].names) == classes:
43 | model.names = ckpt['model'].names # set class names attribute
44 | if autoshape:
45 | model = model.autoshape() # for file/URI/PIL/cv2/np inputs and NMS
46 | return model
47 |
48 | except Exception as e:
49 | help_url = 'https://github.com/ultralytics/yolov5/issues/36'
50 | s = 'Cache maybe be out of date, try force_reload=True. See %s for help.' % help_url
51 | raise Exception(s) from e
52 |
53 |
54 | def yolov5s(pretrained=False, channels=3, classes=80, autoshape=True):
55 | """YOLOv5-small model from https://github.com/ultralytics/yolov5
56 |
57 | Arguments:
58 | pretrained (bool): load pretrained weights into the model, default=False
59 | channels (int): number of input channels, default=3
60 | classes (int): number of model classes, default=80
61 |
62 | Returns:
63 | pytorch model
64 | """
65 | return create('yolov5s', pretrained, channels, classes, autoshape)
66 |
67 |
68 | def yolov5m(pretrained=False, channels=3, classes=80, autoshape=True):
69 | """YOLOv5-medium model from https://github.com/ultralytics/yolov5
70 |
71 | Arguments:
72 | pretrained (bool): load pretrained weights into the model, default=False
73 | channels (int): number of input channels, default=3
74 | classes (int): number of model classes, default=80
75 |
76 | Returns:
77 | pytorch model
78 | """
79 | return create('yolov5m', pretrained, channels, classes, autoshape)
80 |
81 |
82 | def yolov5l(pretrained=False, channels=3, classes=80, autoshape=True):
83 | """YOLOv5-large model from https://github.com/ultralytics/yolov5
84 |
85 | Arguments:
86 | pretrained (bool): load pretrained weights into the model, default=False
87 | channels (int): number of input channels, default=3
88 | classes (int): number of model classes, default=80
89 |
90 | Returns:
91 | pytorch model
92 | """
93 | return create('yolov5l', pretrained, channels, classes, autoshape)
94 |
95 |
96 | def yolov5x(pretrained=False, channels=3, classes=80, autoshape=True):
97 | """YOLOv5-xlarge model from https://github.com/ultralytics/yolov5
98 |
99 | Arguments:
100 | pretrained (bool): load pretrained weights into the model, default=False
101 | channels (int): number of input channels, default=3
102 | classes (int): number of model classes, default=80
103 |
104 | Returns:
105 | pytorch model
106 | """
107 | return create('yolov5x', pretrained, channels, classes, autoshape)
108 |
109 |
110 | def custom(path_or_model='path/to/model.pt', autoshape=True):
111 | """YOLOv5-custom model from https://github.com/ultralytics/yolov5
112 |
113 | Arguments (3 options):
114 | path_or_model (str): 'path/to/model.pt'
115 | path_or_model (dict): torch.load('path/to/model.pt')
116 | path_or_model (nn.Module): torch.load('path/to/model.pt')['model']
117 |
118 | Returns:
119 | pytorch model
120 | """
121 | model = torch.load(path_or_model) if isinstance(path_or_model, str) else path_or_model # load checkpoint
122 | if isinstance(model, dict):
123 | model = model['ema' if model.get('ema') else 'model'] # load model
124 |
125 | hub_model = Model(model.yaml).to(next(model.parameters()).device) # create
126 | hub_model.load_state_dict(model.float().state_dict()) # load state_dict
127 | hub_model.names = model.names # class names
128 | return hub_model.autoshape() if autoshape else hub_model
129 |
130 |
131 | if __name__ == '__main__':
132 | model = create(name='yolov5s', pretrained=True, channels=3, classes=80, autoshape=True) # pretrained example
133 | # model = custom(path_or_model='path/to/model.pt') # custom example
134 |
135 | # Verify inference
136 | import numpy as np
137 | from PIL import Image
138 |
139 | imgs = [Image.open('data/images/bus.jpg'), # PIL
140 | 'data/images/zidane.jpg', # filename
141 | 'https://github.com/ultralytics/yolov5/raw/master/data/images/bus.jpg', # URI
142 | np.zeros((640, 480, 3))] # numpy
143 |
144 | results = model(imgs) # batched inference
145 | results.print()
146 | results.save()
147 |
--------------------------------------------------------------------------------
/models/Ablation_Only_CRMA.yaml:
--------------------------------------------------------------------------------
1 | # parameters
2 | nc: 80 # number of classes
3 | depth_multiple: 0.33 # model depth multiple
4 | width_multiple: 0.50 # layer channel multiple
5 |
6 | # anchors
7 | anchors:
8 | - [10,13, 16,30, 33,23] # P3/8
9 | - [30,61, 62,45, 59,119] # P4/16
10 | - [116,90, 156,198, 373,326] # P5/32
11 |
12 | # YOLOv5 backbone
13 | backbone:
14 | # [from, number, module, args]
15 | [[-1, 1, Focus, [64, 3]], # 0-P1/2
16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
17 | [-1, 3, C3, [128]],
18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
19 | [-1, 9, C3, [256]],
20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
21 | [-1, 9, C3, [512]],
22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
23 | [-1, 1, SPP, [1024, [5, 9, 13]]],
24 | [-1, 3, C3ResAtnMHSA, [1024, 16, False]], # 9
25 | ]
26 |
27 | # YOLOv5 head
28 | head:
29 | [[-1, 1, Conv, [512, 1, 1]],
30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']],
31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4
32 | [-1, 3, C3ResAtnMHSA, [512, 32, False]], # 13
33 |
34 | [-1, 1, Conv, [256, 1, 1]],
35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']],
36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3
37 | [-1, 3, C3ResAtnMHSA, [256, 64, False]], # 17 (P3/8-small)
38 |
39 | [-1, 1, Conv, [256, 3, 2]],
40 | [[-1, 14], 1, Concat, [1]], # cat head P4
41 | [-1, 3, C3ResAtnMHSA, [512, 32, False]], # 20 (P4/16-medium)
42 |
43 | [-1, 1, Conv, [512, 3, 2]],
44 | [[-1, 10], 1, Concat, [1]], # cat head P5
45 | [-1, 3, C3ResAtnMHSA, [1024, 16, False]], # 23 (P5/32-large)
46 |
47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
48 | ]
49 |
--------------------------------------------------------------------------------
/models/Ablation_Only_CRMA_Remove_PAN.yaml:
--------------------------------------------------------------------------------
1 | # parameters
2 | nc: 80 # number of classes
3 | depth_multiple: 0.33 # model depth multiple
4 | width_multiple: 0.50 # layer channel multiple
5 |
6 | # anchors
7 | anchors:
8 | - [10,13, 16,30, 33,23] # P3/8
9 | - [30,61, 62,45, 59,119] # P4/16
10 | - [116,90, 156,198, 373,326] # P5/32
11 |
12 | # YOLOv5 backbone
13 | backbone:
14 | # [from, number, module, args]
15 | [[-1, 1, Focus, [64, 3]], # 0-P1/2
16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
17 | [-1, 3, C3, [128]],
18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
19 | [-1, 9, C3, [256]],
20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
21 | [-1, 9, C3, [512]],
22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
23 | [-1, 1, SPP, [1024, [5, 9, 13]]],
24 | [-1, 3, C3ResAtnMHSA, [1024, 16, False]], # 9
25 | ]
26 |
27 | # YOLOv5 head
28 | head:
29 | [[-1, 1, Conv, [512, 1, 1]],
30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']],
31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4
32 | [-1, 3, C3ResAtnMHSA, [512, 32, False]], # 13
33 |
34 | [-1, 1, Conv, [256, 1, 1]],
35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']],
36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3
37 | [-1, 3, C3ResAtnMHSA, [256, 64, False]], # 17 (P3/8-small)
38 |
39 | [14, 3, C3ResAtnMHSA, [512, 32, False]], # 18 (P4/16-medium)
40 |
41 | [10, 3, C3ResAtnMHSA, [1024, 16, False]], # 19 (P5/32-large)
42 |
43 | [[17, 18, 19], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
44 | ]
45 |
--------------------------------------------------------------------------------
/models/Ablation_Only_CRMA_and_Scale.yaml:
--------------------------------------------------------------------------------
1 | # parameters
2 | nc: 80 # number of classes
3 | depth_multiple: 0.33 # model depth multiple
4 | width_multiple: 0.50 # layer channel multiple
5 |
6 | # anchors
7 | anchors:
8 | - [10,13, 16,30, 33,23] # P3/8
9 | - [30,61, 62,45, 59,119] # P4/16
10 | - [116,90, 156,198, 373,326] # P5/32
11 |
12 | # YOLOv5 backbone
13 | backbone:
14 | # [from, number, module, args]
15 | [[-1, 1, Focus, [64, 3]], # 0-P1/2
16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
17 | [-1, 3, C3, [128]],
18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
19 | [-1, 9, C3, [256]],
20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
21 | [-1, 9, C3, [512]],
22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
23 | [-1, 1, SPP, [1024, [5, 9, 13]]],
24 | [-1, 3, C3ResAtnMHSA, [1024, 16, False]], # 9
25 | ]
26 |
27 | # YOLOv5 head
28 | head:
29 | [[-1, 1, Conv, [512, 1, 1]],
30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']],
31 | [[-1, 6], 1, ConcatFusionFactor, [1]], # cat backbone P4
32 | [-1, 3, C3ResAtnMHSA, [512, 32, False]], # 13
33 |
34 | [-1, 1, Conv, [256, 1, 1]],
35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']],
36 | [[-1, 4], 1, ConcatFusionFactor, [1]], # cat backbone P3
37 | [-1, 3, C3ResAtnMHSA, [256, 64, False]], # 17 (P3/8-small)
38 |
39 | [-1, 1, Conv, [256, 3, 2]],
40 | [[-1, 14], 1, ConcatFusionFactor, [1]], # cat head P4
41 | [-1, 3, C3ResAtnMHSA, [512, 32, False]], # 20 (P4/16-medium)
42 |
43 | [-1, 1, Conv, [512, 3, 2]],
44 | [[-1, 10], 1, ConcatFusionFactor, [1]], # cat head P5
45 | [-1, 3, C3ResAtnMHSA, [1024, 16, False]], # 23 (P5/32-large)
46 |
47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
48 | ]
49 |
--------------------------------------------------------------------------------
/models/Ablation_Only_Detector.yaml:
--------------------------------------------------------------------------------
1 | # parameters
2 | nc: 80 # number of classes
3 | depth_multiple: 0.33 # model depth multiple
4 | width_multiple: 0.50 # layer channel multiple
5 |
6 | # anchors
7 | anchors:
8 | - [10,13, 16,30, 33,23] # P3/8
9 | - [30,61, 62,45, 59,119] # P4/16
10 | - [116,90, 156,198, 373,326] # P5/32
11 |
12 | # YOLOv5 backbone
13 | backbone:
14 | # [from, number, module, args]
15 | [[-1, 1, Focus, [64, 3]], # 0-P1/2
16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
17 | [-1, 3, C3, [128]],
18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
19 | [-1, 9, C3, [256]],
20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
21 | [-1, 9, C3, [512]],
22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
23 | [-1, 1, SPP, [1024, [5, 9, 13]]],
24 | [-1, 3, C3ResAtnMHSA, [1024, 16, False]], # 9
25 | ]
26 |
27 | # YOLOv5 head
28 | head:
29 | [[-1, 1, Conv, [512, 1, 1]],
30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']],
31 | [[-1, 6], 1, ConcatFusionFactor, [1]], # cat backbone P4
32 | [-1, 3, C3ResAtnMHSA, [512, 32, False]], # 13
33 |
34 | [-1, 1, Conv, [256, 1, 1]],
35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']],
36 | [[-1, 4], 1, ConcatFusionFactor, [1]], # cat backbone P3
37 | [-1, 3, C3ResAtnMHSA, [256, 64, False]], # 17 (P3/8-small)
38 |
39 | [14, 3, C3ResAtnMHSA, [512, 32, False]], # 18 (P4/16-medium)
40 |
41 | [10, 3, C3ResAtnMHSA, [1024, 16, False]], # 19 (P5/32-large)
42 |
43 | [[17, 18, 19], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
44 | ]
45 |
--------------------------------------------------------------------------------
/models/Ablation_Only_Enhancer.yaml:
--------------------------------------------------------------------------------
1 | # parameters
2 | nc: 80 # number of classes
3 | depth_multiple: 0.33 # model depth multiple
4 | width_multiple: 0.50 # layer channel multiple
5 |
6 | # anchors
7 | anchors:
8 | - [10,13, 16,30, 33,23] # P3/8
9 | - [30,61, 62,45, 59,119] # P4/16
10 | - [116,90, 156,198, 373,326] # P5/32
11 |
12 | # YOLOv5 backbone
13 | backbone:
14 | # [from, number, module, args]
15 | [[-1, 1, Focus, [64, 3]], # 0-P1/2
16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
17 | [-1, 3, C3, [128]],
18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
19 | [-1, 9, C3, [256]],
20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
21 | [-1, 9, C3, [512]],
22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
23 | [-1, 1, SPP, [1024, [5, 9, 13]]],
24 | [-1, 3, C3, [1024, False]], # 9
25 | ]
26 |
27 | # YOLOv5 head
28 | head:
29 | [[-1, 1, Conv, [512, 1, 1]],
30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']],
31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4
32 | [-1, 3, C3, [512, False]], # 13
33 |
34 | [-1, 1, Conv, [256, 1, 1]],
35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']],
36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3
37 | [-1, 3, C3, [256, False]], # 17 (P3/8-small)
38 |
39 | [-1, 1, Conv, [256, 3, 2]],
40 | [[-1, 14], 1, Concat, [1]], # cat head P4
41 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium)
42 |
43 | [-1, 1, Conv, [512, 3, 2]],
44 | [[-1, 10], 1, Concat, [1]], # cat head P5
45 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large)
46 |
47 | [4, 1, RCAN, []],
48 |
49 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
50 | ]
51 |
--------------------------------------------------------------------------------
/models/Ablation_Only_Scale.yaml:
--------------------------------------------------------------------------------
1 | # parameters
2 | nc: 80 # number of classes
3 | depth_multiple: 0.33 # model depth multiple
4 | width_multiple: 0.50 # layer channel multiple
5 |
6 | # anchors
7 | anchors:
8 | - [10,13, 16,30, 33,23] # P3/8
9 | - [30,61, 62,45, 59,119] # P4/16
10 | - [116,90, 156,198, 373,326] # P5/32
11 |
12 | # YOLOv5 backbone
13 | backbone:
14 | # [from, number, module, args]
15 | [[-1, 1, Focus, [64, 3]], # 0-P1/2
16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
17 | [-1, 3, C3, [128]],
18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
19 | [-1, 9, C3, [256]],
20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
21 | [-1, 9, C3, [512]],
22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
23 | [-1, 1, SPP, [1024, [5, 9, 13]]],
24 | [-1, 3, C3, [1024, False]], # 9
25 | ]
26 |
27 | # YOLOv5 head
28 | head:
29 | [[-1, 1, Conv, [512, 1, 1]],
30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']],
31 | [[-1, 6], 1, ConcatFusionFactor, [1]], # cat backbone P4
32 | [-1, 3, C3, [512, False]], # 13
33 |
34 | [-1, 1, Conv, [256, 1, 1]],
35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']],
36 | [[-1, 4], 1, ConcatFusionFactor, [1]], # cat backbone P3
37 | [-1, 3, C3, [256, False]], # 17 (P3/8-small)
38 |
39 | [-1, 1, Conv, [256, 3, 2]],
40 | [[-1, 14], 1, ConcatFusionFactor, [1]], # cat head P4
41 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium)
42 |
43 | [-1, 1, Conv, [512, 3, 2]],
44 | [[-1, 10], 1, ConcatFusionFactor, [1]], # cat head P5
45 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large)
46 |
47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
48 | ]
49 |
--------------------------------------------------------------------------------
/models/DRENet.yaml:
--------------------------------------------------------------------------------
1 | # parameters
2 | nc: 80 # number of classes
3 | depth_multiple: 0.33 # model depth multiple
4 | width_multiple: 0.50 # layer channel multiple
5 |
6 | # anchors
7 | anchors:
8 | - [10,13, 16,30, 33,23] # P3/8
9 | - [30,61, 62,45, 59,119] # P4/16
10 | - [116,90, 156,198, 373,326] # P5/32
11 |
12 | # YOLOv5 backbone
13 | backbone:
14 | # [from, number, module, args]
15 | [[-1, 1, Focus, [64, 3]], # 0-P1/2
16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
17 | [-1, 3, C3, [128]],
18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
19 | [-1, 9, C3, [256]],
20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
21 | [-1, 9, C3, [512]],
22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
23 | [-1, 1, SPP, [1024, [5, 9, 13]]],
24 | [-1, 3, C3ResAtnMHSA, [1024, 16, False]], # 9
25 | ]
26 |
27 | # YOLOv5 head
28 | head:
29 | [[-1, 1, Conv, [512, 1, 1]],
30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']],
31 | [[-1, 6], 1, ConcatFusionFactor, [1]], # cat backbone P4
32 | [-1, 3, C3ResAtnMHSA, [512, 32, False]], # 13
33 |
34 | [-1, 1, Conv, [256, 1, 1]],
35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']],
36 | [[-1, 4], 1, ConcatFusionFactor, [1]], # cat backbone P3
37 | [-1, 3, C3ResAtnMHSA, [256, 64, False]], # 17 (P3/8-small)
38 |
39 | [14, 3, C3ResAtnMHSA, [512, 32, False]], # 18 (P4/16-medium)
40 |
41 | [10, 3, C3ResAtnMHSA, [1024, 16, False]], # 19 (P5/32-large)
42 |
43 | [4, 1, RCAN, []],
44 | [[17, 18, 19], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
45 | ]
46 |
--------------------------------------------------------------------------------
/models/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/WindVChen/DRENet/a187dbe0f623b521a62c6176c7cafaa7322f5f66/models/__init__.py
--------------------------------------------------------------------------------
/models/common.py:
--------------------------------------------------------------------------------
1 | # This file contains modules common to various models
2 |
3 | import math
4 |
5 | import numpy as np
6 | import requests
7 | import torch
8 | import torch.nn as nn
9 | from PIL import Image
10 |
11 | from utils.datasets import letterbox
12 | from utils.general import non_max_suppression, make_divisible, scale_coords, xyxy2xywh
13 | from utils.plots import color_list
14 |
15 |
16 | def autopad(k, p=None): # kernel, padding
17 | # Pad to 'same'
18 | if p is None:
19 | p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad
20 | return p
21 |
22 |
23 | def DWConv(c1, c2, k=1, s=1, act=True):
24 | # Depthwise convolution
25 | return Conv(c1, c2, k, s, g=math.gcd(c1, c2), act=act)
26 |
27 | class Conv(nn.Module):
28 | # Standard convolution
29 | def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True, bias=False): # ch_in, ch_out, kernel, stride, padding, groups
30 | super(Conv, self).__init__()
31 | self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False)
32 | self.bn = nn.BatchNorm2d(c2)
33 | self.act = nn.SiLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity())
34 |
35 | def forward(self, x):
36 | return self.act(self.bn(self.conv(x)))
37 |
38 | def fuseforward(self, x):
39 | return self.act(self.conv(x))
40 |
41 | class inv(nn.Module):
42 | # Standard convolution
43 | def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True, bias=False): # ch_in, ch_out, kernel, stride, padding, groups
44 | super(inv, self).__init__()
45 | self.INV = False
46 | self.inChannel = c1
47 | if self.inChannel<4 or self.inChannel<16 or not self.INV:
48 | self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=bias)
49 | self.bn = nn.BatchNorm2d(c2)
50 | self.act = nn.SiLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity())
51 | else:
52 | kernel_size = k
53 | stride = s
54 | channels = c1
55 | channelsOut = c2
56 | self.kernel_size = k
57 | self.stride = s
58 | self.channels = c1
59 | reduction_ratio = 4
60 | self.group_channels = 16
61 | self.groups = self.channels // self.group_channels
62 | self.conv1 = nn.Sequential(nn.Conv2d(channels, channels // reduction_ratio, 1, groups=g, bias=bias),
63 | nn.BatchNorm2d(channels // reduction_ratio),
64 | nn.ReLU())
65 | self.conv2 = nn.Conv2d(channels // reduction_ratio, kernel_size ** 2 * self.groups, 1, groups=g, bias=bias)
66 | if stride > 1:
67 | self.avgpool = nn.AvgPool2d(stride, stride)
68 | self.unfold = nn.Unfold(kernel_size, 1, (kernel_size - 1) // 2, stride)
69 | self.bn = nn.BatchNorm2d(channels)
70 | self.act = nn.SiLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity())
71 | self.conv3 = nn.Conv2d(channels, channelsOut, 1, groups=g, bias=bias)
72 | self.bn2 = nn.BatchNorm2d(channelsOut)
73 | self.act2 = nn.SiLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity())
74 |
75 | def forward(self, x):
76 | if self.inChannel < 4 or self.inChannel < 16 or not self.INV:
77 | return self.act(self.bn(self.conv(x)))
78 | else:
79 | weight = self.conv2(self.conv1(x if self.stride == 1 else self.avgpool(x)))
80 | b, c, h, w = weight.shape
81 | weight = weight.view(b, self.groups, self.kernel_size ** 2, h, w).unsqueeze(2)
82 | out = self.unfold(x).view(b, self.groups, self.group_channels, self.kernel_size ** 2, h, w)
83 | out = (weight * out).sum(dim=3).view(b, self.channels, h, w)
84 | out = self.act(self.bn(out))
85 | return self.act2(self.bn2(self.conv3(out)))
86 | def fuseforward(self, x):
87 | if self.inChannel < 4 or self.inChannel < 16 or not self.INV:
88 | return self.act(self.conv(x))
89 | else:
90 | return self.act2(self.conv3(x))
91 |
92 |
93 | class Bottleneck(nn.Module):
94 | # Standard bottleneck
95 | def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): # ch_in, ch_out, shortcut, groups, expansion
96 | super(Bottleneck, self).__init__()
97 | c_ = int(c2 * e) # hidden channels
98 | self.cv1 = Conv(c1, c_, 1, 1)
99 | self.cv2 = Conv(c_, c2, 3, 1, g=g)
100 | self.add = shortcut and c1 == c2
101 |
102 | def forward(self, x):
103 | return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
104 |
105 | class BottleneckResAtnMHSA(nn.Module):
106 | # Standard bottleneck
107 | def __init__(self, n_dims, size, shortcut=True): # ch_in, ch_out, shortcut, groups, expansion
108 | super(BottleneckResAtnMHSA, self).__init__()
109 |
110 | height=size
111 | width=size
112 | self.cv1 = Conv(n_dims, n_dims//2, 1, 1)
113 | self.cv2 = Conv(n_dims//2, n_dims, 1, 1)
114 | '''MHSA PARAGRAMS'''
115 | self.query = nn.Conv2d(n_dims//2, n_dims//2, kernel_size=1)
116 | self.key = nn.Conv2d(n_dims//2, n_dims//2, kernel_size=1)
117 | self.value = nn.Conv2d(n_dims//2, n_dims//2, kernel_size=1)
118 |
119 | self.rel_h = nn.Parameter(torch.randn([1, n_dims//2, height, 1]), requires_grad=True)
120 | self.rel_w = nn.Parameter(torch.randn([1, n_dims//2, 1, width]), requires_grad=True)
121 |
122 | self.softmax = nn.Softmax(dim=-1)
123 | self.add = shortcut
124 |
125 | def forward(self, x):
126 | x1=self.cv1(x)
127 | n_batch, C, width, height = x1.size()
128 | q = self.query(x1).view(n_batch, C, -1)
129 | k = self.key(x1).view(n_batch, C, -1)
130 | v = self.value(x1).view(n_batch, C, -1)
131 |
132 | content_content = torch.bmm(q.permute(0, 2, 1), k)
133 |
134 | content_position = (self.rel_h + self.rel_w).view(1, C, -1).permute(0, 2, 1)
135 |
136 | # If you want to use resolution-agnostic positional encoding, you can uncomment the following lines.
137 | # See details in https://github.com/WindVChen/DRENet/issues/10.
138 | # Note that the performance of this resolution-agnostic positional encoding is not tested.
139 | # content_position = (self.rel_h + self.rel_w)
140 | # content_position = nn.functional.interpolate(content_position, (int(content_content.shape[-1]**0.5), int(content_content.shape[-1]**0.5)), mode='bilinear')
141 | # content_position = content_position.view(1, C, -1).permute(0, 2, 1)
142 |
143 | content_position = torch.matmul(content_position, q)
144 |
145 | energy = content_content + content_position
146 | attention = self.softmax(energy)
147 |
148 | out = torch.bmm(v, attention.permute(0, 2, 1))
149 | out = out.view(n_batch, C, width, height)
150 |
151 |
152 | return x + self.cv2(out) if self.add else self.cv2(out)
153 |
154 |
155 | class BottleneckCSP(nn.Module):
156 | # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
157 | def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
158 | super(BottleneckCSP, self).__init__()
159 | c_ = int(c2 * e) # hidden channels
160 | self.cv1 = Conv(c1, c_, 1, 1)
161 | self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False)
162 | self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False)
163 | self.cv4 = Conv(2 * c_, c2, 1, 1)
164 | self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3)
165 | self.act = nn.LeakyReLU(0.1, inplace=True)
166 | self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
167 |
168 | def forward(self, x):
169 | y1 = self.cv3(self.m(self.cv1(x)))
170 | y2 = self.cv2(x)
171 | return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1))))
172 |
173 |
174 | class C3(nn.Module):
175 | # CSP Bottleneck with 3 convolutions
176 | def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
177 | super(C3, self).__init__()
178 | c_ = int(c2 * e) # hidden channels
179 | self.cv1 = Conv(c1, c_, 1, 1)
180 | self.cv2 = Conv(c1, c_, 1, 1)
181 | self.cv3 = Conv(2 * c_, c2, 1) # act=FReLU(c2)
182 | self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
183 | # self.m = nn.Sequential(*[CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)])
184 |
185 | def forward(self, x):
186 | return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), dim=1))
187 |
188 | class C3ResAtnMHSA(nn.Module):
189 | # CSP Bottleneck with 3 convolutions
190 | def __init__(self, c1, c2, n=1, size=14, shortcut=True, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion
191 | super(C3ResAtnMHSA, self).__init__()
192 | c_ = int(c2 * e) # hidden channels
193 | self.cv1 = Conv(c1, c_, 1, 1)
194 | self.cv2 = Conv(c1, c_, 1, 1)
195 | self.cv3 = Conv(2 * c_, c2, 1) # act=FReLU(c2)
196 | self.m = nn.Sequential(*[BottleneckResAtnMHSA(c_, size, shortcut=True) for _ in range(n)])
197 | # self.m = nn.Sequential(*[CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)])
198 |
199 | def forward(self, x):
200 | return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), dim=1))
201 |
202 |
203 | class SPP(nn.Module):
204 | # Spatial pyramid pooling layer used in YOLOv3-SPP
205 | def __init__(self, c1, c2, k=(5, 9, 13)):
206 | super(SPP, self).__init__()
207 | c_ = c1 // 2 # hidden channels
208 | self.cv1 = Conv(c1, c_, 1, 1)
209 | self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)
210 | self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])
211 |
212 | def forward(self, x):
213 | x = self.cv1(x)
214 | return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))
215 |
216 |
217 | class Focus(nn.Module):
218 | # Focus wh information into c-space
219 | def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups
220 | super(Focus, self).__init__()
221 | self.conv = Conv(c1 * 4, c2, k, s, p, g, act)
222 | # self.contract = Contract(gain=2)
223 |
224 | def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2)
225 | return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))
226 | # return self.conv(self.contract(x))
227 |
228 |
229 | class Contract(nn.Module):
230 | # Contract width-height into channels, i.e. x(1,64,80,80) to x(1,256,40,40)
231 | def __init__(self, gain=2):
232 | super().__init__()
233 | self.gain = gain
234 |
235 | def forward(self, x):
236 | N, C, H, W = x.size() # assert (H / s == 0) and (W / s == 0), 'Indivisible gain'
237 | s = self.gain
238 | x = x.view(N, C, H // s, s, W // s, s) # x(1,64,40,2,40,2)
239 | x = x.permute(0, 3, 5, 1, 2, 4).contiguous() # x(1,2,2,64,40,40)
240 | return x.view(N, C * s * s, H // s, W // s) # x(1,256,40,40)
241 |
242 |
243 | class Expand(nn.Module):
244 | # Expand channels into width-height, i.e. x(1,64,80,80) to x(1,16,160,160)
245 | def __init__(self, gain=2):
246 | super().__init__()
247 | self.gain = gain
248 |
249 | def forward(self, x):
250 | N, C, H, W = x.size() # assert C / s ** 2 == 0, 'Indivisible gain'
251 | s = self.gain
252 | x = x.view(N, s, s, C // s ** 2, H, W) # x(1,2,2,16,80,80)
253 | x = x.permute(0, 3, 4, 1, 5, 2).contiguous() # x(1,16,80,2,80,2)
254 | return x.view(N, C // s ** 2, H * s, W * s) # x(1,16,160,160)
255 |
256 |
257 | class Concat(nn.Module):
258 | # Concatenate a list of tensors along dimension
259 | def __init__(self, dimension=1, selectPos=None):
260 | super(Concat, self).__init__()
261 | self.d = dimension
262 | self.p=selectPos
263 |
264 | def forward(self, x):
265 | if isinstance(self.p, int):
266 | return torch.cat([x[0][self.p],x[1]], self.d)
267 | else:
268 | return torch.cat(x, self.d)
269 |
270 | class ConcatFusionFactor(nn.Module):
271 | # Concatenate a list of tensors along dimension
272 | def __init__(self, dimension=1):
273 | super(ConcatFusionFactor, self).__init__()
274 | self.d = dimension
275 | self.factor=torch.nn.Parameter(torch.FloatTensor([1]))
276 |
277 | def forward(self, x):
278 | x[0]=x[0]*self.factor
279 | return torch.cat(x, self.d)
280 |
281 |
282 | class NMS(nn.Module):
283 | # Non-Maximum Suppression (NMS) module
284 | conf = 0.25 # confidence threshold
285 | iou = 0.45 # IoU threshold
286 | classes = None # (optional list) filter by class
287 |
288 | def __init__(self):
289 | super(NMS, self).__init__()
290 |
291 | def forward(self, x):
292 | return non_max_suppression(x[0], conf_thres=self.conf, iou_thres=self.iou, classes=self.classes)
293 |
294 |
295 | class autoShape(nn.Module):
296 | # input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS
297 | img_size = 640 # inference size (pixels)
298 | conf = 0.25 # NMS confidence threshold
299 | iou = 0.45 # NMS IoU threshold
300 | classes = None # (optional list) filter by class
301 |
302 | def __init__(self, model):
303 | super(autoShape, self).__init__()
304 | self.model = model.eval()
305 |
306 | def autoshape(self):
307 | print('autoShape already enabled, skipping... ') # model already converted to model.autoshape()
308 | return self
309 |
310 | def forward(self, imgs, size=640, augment=False, profile=False):
311 | # Inference from various sources. For height=720, width=1280, RGB images example inputs are:
312 | # filename: imgs = 'data/samples/zidane.jpg'
313 | # URI: = 'https://github.com/ultralytics/yolov5/releases/download/v1.0/zidane.jpg'
314 | # OpenCV: = cv2.imread('image.jpg')[:,:,::-1] # HWC BGR to RGB x(720,1280,3)
315 | # PIL: = Image.open('image.jpg') # HWC x(720,1280,3)
316 | # numpy: = np.zeros((720,1280,3)) # HWC
317 | # torch: = torch.zeros(16,3,720,1280) # BCHW
318 | # multiple: = [Image.open('image1.jpg'), Image.open('image2.jpg'), ...] # list of images
319 |
320 | p = next(self.model.parameters()) # for device and type
321 | if isinstance(imgs, torch.Tensor): # torch
322 | return self.model(imgs.to(p.device).type_as(p), augment, profile) # inference
323 |
324 | # Pre-process
325 | n, imgs = (len(imgs), imgs) if isinstance(imgs, list) else (1, [imgs]) # number of images, list of images
326 | shape0, shape1 = [], [] # image and inference shapes
327 | for i, im in enumerate(imgs):
328 | if isinstance(im, str): # filename or uri
329 | im = Image.open(requests.get(im, stream=True).raw if im.startswith('http') else im) # open
330 | im = np.array(im) # to numpy
331 | if im.shape[0] < 5: # image in CHW
332 | im = im.transpose((1, 2, 0)) # reverse dataloader .transpose(2, 0, 1)
333 | im = im[:, :, :3] if im.ndim == 3 else np.tile(im[:, :, None], 3) # enforce 3ch input
334 | s = im.shape[:2] # HWC
335 | shape0.append(s) # image shape
336 | g = (size / max(s)) # gain
337 | shape1.append([y * g for y in s])
338 | imgs[i] = im # update
339 | shape1 = [make_divisible(x, int(self.stride.max())) for x in np.stack(shape1, 0).max(0)] # inference shape
340 | x = [letterbox(im, new_shape=shape1, auto=False)[0] for im in imgs] # pad
341 | x = np.stack(x, 0) if n > 1 else x[0][None] # stack
342 | x = np.ascontiguousarray(x.transpose((0, 3, 1, 2))) # BHWC to BCHW
343 | x = torch.from_numpy(x).to(p.device).type_as(p) / 255. # uint8 to fp16/32
344 |
345 | # Inference
346 | with torch.no_grad():
347 | y = self.model(x, augment, profile)[0] # forward
348 | y = non_max_suppression(y, conf_thres=self.conf, iou_thres=self.iou, classes=self.classes) # NMS
349 |
350 | # Post-process
351 | for i in range(n):
352 | scale_coords(shape1, y[i][:, :4], shape0[i])
353 |
354 | return Detections(imgs, y, self.names)
355 |
356 |
357 | class Detections:
358 | # detections class for YOLOv5 inference results
359 | def __init__(self, imgs, pred, names=None):
360 | super(Detections, self).__init__()
361 | d = pred[0].device # device
362 | gn = [torch.tensor([*[im.shape[i] for i in [1, 0, 1, 0]], 1., 1.], device=d) for im in imgs] # normalizations
363 | self.imgs = imgs # list of images as numpy arrays
364 | self.pred = pred # list of tensors pred[0] = (xyxy, conf, cls)
365 | self.names = names # class names
366 | self.xyxy = pred # xyxy pixels
367 | self.xywh = [xyxy2xywh(x) for x in pred] # xywh pixels
368 | self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)] # xyxy normalized
369 | self.xywhn = [x / g for x, g in zip(self.xywh, gn)] # xywh normalized
370 | self.n = len(self.pred)
371 |
372 | def display(self, pprint=False, show=False, save=False, render=False):
373 | colors = color_list()
374 | for i, (img, pred) in enumerate(zip(self.imgs, self.pred)):
375 | str = f'Image {i + 1}/{len(self.pred)}: {img.shape[0]}x{img.shape[1]} '
376 | if pred is not None:
377 | for c in pred[:, -1].unique():
378 | n = (pred[:, -1] == c).sum() # detections per class
379 | str += f'{n} {self.names[int(c)]}s, ' # add to string
380 | if show or save or render:
381 | img = Image.fromarray(img.astype(np.uint8)) if isinstance(img, np.ndarray) else img # from np
382 | for *box, conf, cls in pred: # xyxy, confidence, class
383 | # str += '%s %.2f, ' % (names[int(cls)], conf) # label
384 | ImageDraw.Draw(img).rectangle(box, width=4, outline=colors[int(cls) % 10]) # plot
385 | if pprint:
386 | print(str)
387 | if show:
388 | img.show(f'Image {i}') # show
389 | if save:
390 | f = f'results{i}.jpg'
391 | str += f"saved to '{f}'"
392 | img.save(f) # save
393 | if render:
394 | self.imgs[i] = np.asarray(img)
395 |
396 | def print(self):
397 | self.display(pprint=True) # print results
398 |
399 | def show(self):
400 | self.display(show=True) # show results
401 |
402 | def save(self):
403 | self.display(save=True) # save results
404 |
405 | def render(self):
406 | self.display(render=True) # render results
407 | return self.imgs
408 |
409 | def __len__(self):
410 | return self.n
411 |
412 | def tolist(self):
413 | # return a list of Detections objects, i.e. 'for result in results.tolist():'
414 | x = [Detections([self.imgs[i]], [self.pred[i]], self.names) for i in range(self.n)]
415 | for d in x:
416 | for k in ['imgs', 'pred', 'xyxy', 'xyxyn', 'xywh', 'xywhn']:
417 | setattr(d, k, getattr(d, k)[0]) # pop out of list
418 | return x
419 |
420 |
421 | class Classify(nn.Module):
422 | # Classification head, i.e. x(b,c1,20,20) to x(b,c2)
423 | def __init__(self, c1, c2, k=1, s=1, p=None, g=1): # ch_in, ch_out, kernel, stride, padding, groups
424 | super(Classify, self).__init__()
425 | self.aap = nn.AdaptiveAvgPool2d(1) # to x(b,c1,1,1)
426 | self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g) # to x(b,c2,1,1)
427 | self.flat = nn.Flatten()
428 |
429 | def forward(self, x):
430 | z = torch.cat([self.aap(y) for y in (x if isinstance(x, list) else [x])], 1) # cat if list
431 | return self.flat(self.conv(z)) # flatten to x(b,c2)
432 |
433 | class LocalReconstruct(nn.Module):
434 | def __init__(self, c1, c2):
435 | super(LocalReconstruct, self).__init__()
436 | self.reconstruct = nn.Sequential(
437 | Conv(c1, c2, 1, 1),
438 | Conv(c2, c2//4, 3, 1),
439 | Conv(c2//4, c2, 1, 1)
440 | )
441 |
442 | def forward(self, x):
443 | x0=self.reconstruct(x[0])
444 | x1=self.reconstruct(x[1])
445 | x2=self.reconstruct(x[2])
446 | return x0,x1,x2
447 |
448 | class SEAtn(nn.Module):
449 | def __init__(self, channel, reduction=16):
450 | super(SEAtn, self).__init__()
451 | self.avg_pool = nn.AdaptiveAvgPool2d(1)
452 | self.fc = nn.Sequential(
453 | nn.Linear(channel, channel // reduction, bias=False),
454 | nn.ReLU(inplace=True),
455 | nn.Linear(channel // reduction, channel, bias=False),
456 | nn.Sigmoid()
457 | )
458 |
459 | def forward(self, x):
460 | b, c, _, _ = x.size()
461 | y = self.avg_pool(x).view(b, c)
462 | y = self.fc(y).view(b, c, 1, 1)
463 | return y
464 |
465 | class AtnMut(nn.Module):
466 | def __init__(self, start, end):
467 | super(AtnMut, self).__init__()
468 | self.start=start
469 | self.end=end
470 |
471 | def forward(self, x):
472 | atn= x[0][:, self.start:self.end, :, :]
473 | obj= x[1]
474 | out= obj*atn.expand_as(obj)
475 | return out
476 |
477 | class MHSA(nn.Module):
478 | def __init__(self, n_dims, size):
479 | super(MHSA, self).__init__()
480 |
481 | height = size
482 | width = size
483 | self.query = nn.Conv2d(n_dims, n_dims, kernel_size=1)
484 | self.key = nn.Conv2d(n_dims, n_dims, kernel_size=1)
485 | self.value = nn.Conv2d(n_dims, n_dims, kernel_size=1)
486 |
487 | self.rel_h = nn.Parameter(torch.randn([1, n_dims, height, 1]), requires_grad=True)
488 | self.rel_w = nn.Parameter(torch.randn([1, n_dims, 1, width]), requires_grad=True)
489 |
490 | self.softmax = nn.Softmax(dim=-1)
491 |
492 | def forward(self, x):
493 | n_batch, C, width, height = x.size()
494 | q = self.query(x).view(n_batch, C, -1)
495 | k = self.key(x).view(n_batch, C, -1)
496 | v = self.value(x).view(n_batch, C, -1)
497 |
498 | content_content = torch.bmm(q.permute(0, 2, 1), k)
499 |
500 | content_position = (self.rel_h + self.rel_w).view(1, C, -1).permute(0, 2, 1)
501 | content_position = torch.matmul(content_position, q)
502 |
503 | energy = content_content + content_position
504 | attention = self.softmax(energy)
505 |
506 | out = torch.bmm(v, attention.permute(0, 2, 1))
507 | out = out.view(n_batch, C, width, height)
508 |
509 | return out
510 |
511 | ## Residual Channel Attention Network (RCAN)
512 | class RCAN(nn.Module):
513 | def __init__(self, c1, conv=Conv):
514 | super(RCAN, self).__init__()
515 |
516 | n_resgroups = 1
517 | n_resblocks = 1
518 | n_feats = c1
519 | kernel_size = 3
520 | reduction = 16
521 | scale = 2
522 | act = nn.SiLU()
523 |
524 | # define body module
525 | modules_body = [
526 | ResidualGroup(
527 | conv, n_feats, kernel_size, reduction, act=act, res_scale=1, n_resblocks=n_resblocks) \
528 | for _ in range(n_resgroups)]
529 |
530 | modules_body.append(conv(n_feats, n_feats, kernel_size))
531 |
532 | # define tail module
533 | modules_tail = [
534 | Upsampler(conv, scale, n_feats, act=False),
535 | conv(n_feats, 3, kernel_size)]
536 |
537 | self.body = nn.Sequential(*modules_body)
538 | self.tail = nn.Sequential(*modules_tail)
539 |
540 | def forward(self, x):
541 | res = self.body(x)
542 | res += x
543 |
544 | x = self.tail(res)
545 |
546 | return x
547 |
548 |
549 | class CALayer(nn.Module):
550 | def __init__(self, channel, reduction=16):
551 | super(CALayer, self).__init__()
552 | # global average pooling: feature --> point
553 | self.avg_pool = nn.AdaptiveAvgPool2d(1)
554 | # feature channel downscale and upscale --> channel weight
555 | self.conv_du = nn.Sequential(
556 | nn.Conv2d(channel, channel // reduction, 1, padding=0, bias=True),
557 | nn.ReLU(inplace=True),
558 | nn.Conv2d(channel // reduction, channel, 1, padding=0, bias=True),
559 | nn.Sigmoid()
560 | )
561 |
562 | def forward(self, x):
563 | y = self.avg_pool(x)
564 | y = self.conv_du(y)
565 | return x * y
566 |
567 |
568 | ## Residual Channel Attention Block (RCAB)
569 | class RCAB(nn.Module):
570 | def __init__(
571 | self, conv, n_feat, kernel_size, reduction,
572 | bias=True, bn=False, act=nn.ReLU(True), res_scale=1):
573 |
574 | super(RCAB, self).__init__()
575 | modules_body = []
576 | for i in range(2):
577 | modules_body.append(conv(n_feat, n_feat, kernel_size, bias=bias))
578 | if bn: modules_body.append(nn.BatchNorm2d(n_feat))
579 | if i == 0: modules_body.append(act)
580 | modules_body.append(CALayer(n_feat, reduction))
581 | self.body = nn.Sequential(*modules_body)
582 | self.res_scale = res_scale
583 |
584 | def forward(self, x):
585 | res = self.body(x)
586 | # res = self.body(x).mul(self.res_scale)
587 | res += x
588 | return res
589 |
590 |
591 | ## Residual Group (RG)
592 | class ResidualGroup(nn.Module):
593 | def __init__(self, conv, n_feat, kernel_size, reduction, act, res_scale, n_resblocks):
594 | super(ResidualGroup, self).__init__()
595 | modules_body = []
596 | modules_body = [
597 | RCAB(
598 | conv, n_feat, kernel_size, reduction, bias=True, bn=False, act=nn.ReLU(True), res_scale=1) \
599 | for _ in range(n_resblocks)]
600 | modules_body.append(conv(n_feat, n_feat, kernel_size))
601 | self.body = nn.Sequential(*modules_body)
602 |
603 | def forward(self, x):
604 | res = self.body(x)
605 | res += x
606 | return res
607 |
608 | class Upsampler(nn.Sequential):
609 | def __init__(self, conv, scale, n_feat, bn=False, act=False, bias=True):
610 |
611 | m = []
612 | for _ in range(int(math.log(scale, 2))):
613 | m.append(conv(n_feat, 4 * n_feat, 3, bias = bias))
614 | m.append(nn.PixelShuffle(2))
615 | if bn: m.append(nn.BatchNorm2d(n_feat))
616 | if act: m.append(act())
617 |
618 | super(Upsampler, self).__init__(*m)
--------------------------------------------------------------------------------
/models/experimental.py:
--------------------------------------------------------------------------------
1 | # This file contains experimental modules
2 |
3 | import numpy as np
4 | import torch
5 | import torch.nn as nn
6 |
7 | from models.common import Conv, DWConv, RCAN
8 | from utils.google_utils import attempt_download
9 |
10 |
11 | class CrossConv(nn.Module):
12 | # Cross Convolution Downsample
13 | def __init__(self, c1, c2, k=3, s=1, g=1, e=1.0, shortcut=False):
14 | # ch_in, ch_out, kernel, stride, groups, expansion, shortcut
15 | super(CrossConv, self).__init__()
16 | c_ = int(c2 * e) # hidden channels
17 | self.cv1 = Conv(c1, c_, (1, k), (1, s))
18 | self.cv2 = Conv(c_, c2, (k, 1), (s, 1), g=g)
19 | self.add = shortcut and c1 == c2
20 |
21 | def forward(self, x):
22 | return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
23 |
24 |
25 | class Sum(nn.Module):
26 | # Weighted sum of 2 or more layers https://arxiv.org/abs/1911.09070
27 | def __init__(self, n, weight=False): # n: number of inputs
28 | super(Sum, self).__init__()
29 | self.weight = weight # apply weights boolean
30 | self.iter = range(n - 1) # iter object
31 | if weight:
32 | self.w = nn.Parameter(-torch.arange(1., n) / 2, requires_grad=True) # layer weights
33 |
34 | def forward(self, x):
35 | y = x[0] # no weight
36 | if self.weight:
37 | w = torch.sigmoid(self.w) * 2
38 | for i in self.iter:
39 | y = y + x[i + 1] * w[i]
40 | else:
41 | for i in self.iter:
42 | y = y + x[i + 1]
43 | return y
44 |
45 |
46 | class GhostConv(nn.Module):
47 | # Ghost Convolution https://github.com/huawei-noah/ghostnet
48 | def __init__(self, c1, c2, k=1, s=1, g=1, act=True): # ch_in, ch_out, kernel, stride, groups
49 | super(GhostConv, self).__init__()
50 | c_ = c2 // 2 # hidden channels
51 | self.cv1 = Conv(c1, c_, k, s, None, g, act)
52 | self.cv2 = Conv(c_, c_, 5, 1, None, c_, act)
53 |
54 | def forward(self, x):
55 | y = self.cv1(x)
56 | return torch.cat([y, self.cv2(y)], 1)
57 |
58 |
59 | class GhostBottleneck(nn.Module):
60 | # Ghost Bottleneck https://github.com/huawei-noah/ghostnet
61 | def __init__(self, c1, c2, k=3, s=1): # ch_in, ch_out, kernel, stride
62 | super(GhostBottleneck, self).__init__()
63 | c_ = c2 // 2
64 | self.conv = nn.Sequential(GhostConv(c1, c_, 1, 1), # pw
65 | DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(), # dw
66 | GhostConv(c_, c2, 1, 1, act=False)) # pw-linear
67 | self.shortcut = nn.Sequential(DWConv(c1, c1, k, s, act=False),
68 | Conv(c1, c2, 1, 1, act=False)) if s == 2 else nn.Identity()
69 |
70 | def forward(self, x):
71 | return self.conv(x) + self.shortcut(x)
72 |
73 |
74 | class MixConv2d(nn.Module):
75 | # Mixed Depthwise Conv https://arxiv.org/abs/1907.09595
76 | def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True):
77 | super(MixConv2d, self).__init__()
78 | groups = len(k)
79 | if equal_ch: # equal c_ per group
80 | i = torch.linspace(0, groups - 1E-6, c2).floor() # c2 indices
81 | c_ = [(i == g).sum() for g in range(groups)] # intermediate channels
82 | else: # equal weight.numel() per group
83 | b = [c2] + [0] * groups
84 | a = np.eye(groups + 1, groups, k=-1)
85 | a -= np.roll(a, 1, axis=1)
86 | a *= np.array(k) ** 2
87 | a[0] = 1
88 | c_ = np.linalg.lstsq(a, b, rcond=None)[0].round() # solve for equal weight indices, ax = b
89 |
90 | self.m = nn.ModuleList([nn.Conv2d(c1, int(c_[g]), k[g], s, k[g] // 2, bias=False) for g in range(groups)])
91 | self.bn = nn.BatchNorm2d(c2)
92 | self.act = nn.LeakyReLU(0.1, inplace=True)
93 |
94 | def forward(self, x):
95 | return x + self.act(self.bn(torch.cat([m(x) for m in self.m], 1)))
96 |
97 |
98 | class Ensemble(nn.ModuleList):
99 | # Ensemble of models
100 | def __init__(self):
101 | super(Ensemble, self).__init__()
102 |
103 | def forward(self, x, augment=False):
104 | y = []
105 | for module in self:
106 | y.append(module(x, augment)[0])
107 | # y = torch.stack(y).max(0)[0] # max ensemble
108 | # y = torch.stack(y).mean(0) # mean ensemble
109 | y = torch.cat(y, 1) # nms ensemble
110 | return y, None # inference, train output
111 |
112 |
113 | def attempt_load(weights, map_location=None):
114 | # Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a
115 | model = Ensemble()
116 | for w in weights if isinstance(weights, list) else [weights]:
117 | attempt_download(w)
118 | ckpt = torch.load(w, map_location=map_location) # load
119 | model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().fuse().eval()) # FP32 model
120 |
121 | # Compatibility updates
122 | for m in model.modules():
123 | if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU]:
124 | m.inplace = True # pytorch 1.7.0 compatibility
125 | elif type(m) is Conv:
126 | m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility
127 |
128 | if len(model) == 1:
129 | return model[-1] # return model
130 | else:
131 | print('Ensemble created with %s\n' % weights)
132 | for k in ['names', 'stride']:
133 | setattr(model, k, getattr(model[-1], k))
134 | return model # return ensemble
135 |
--------------------------------------------------------------------------------
/models/export.py:
--------------------------------------------------------------------------------
1 | """Exports a YOLOv5 *.pt model to ONNX and TorchScript formats
2 |
3 | Usage:
4 | $ export PYTHONPATH="$PWD" && python models/export.py --weights ./weights/yolov5s.pt --img 640 --batch 1
5 | """
6 |
7 | import argparse
8 | import sys
9 | import time
10 |
11 | sys.path.append('./') # to run '$ python *.py' files in subdirectories
12 |
13 | import torch
14 | import torch.nn as nn
15 |
16 | import models
17 | from models.experimental import attempt_load
18 | from utils.activations import Hardswish, SiLU
19 | from utils.general import set_logging, check_img_size
20 |
21 | if __name__ == '__main__':
22 | parser = argparse.ArgumentParser()
23 | parser.add_argument('--weights', type=str, default='./yolov5s.pt', help='weights path') # from yolov5/models/
24 | parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size') # height, width
25 | parser.add_argument('--dynamic', action='store_true', help='dynamic ONNX axes')
26 | parser.add_argument('--batch-size', type=int, default=1, help='batch size')
27 | opt = parser.parse_args()
28 | opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand
29 | print(opt)
30 | set_logging()
31 | t = time.time()
32 |
33 | # Load PyTorch model
34 | model = attempt_load(opt.weights, map_location=torch.device('cpu')) # load FP32 model
35 | labels = model.names
36 |
37 | # Checks
38 | gs = int(max(model.stride)) # grid size (max stride)
39 | opt.img_size = [check_img_size(x, gs) for x in opt.img_size] # verify img_size are gs-multiples
40 |
41 | # Input
42 | img = torch.zeros(opt.batch_size, 3, *opt.img_size) # image size(1,3,320,192) iDetection
43 |
44 | # Update model
45 | for k, m in model.named_modules():
46 | m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility
47 | if isinstance(m, models.common.Conv): # assign export-friendly activations
48 | if isinstance(m.act, nn.Hardswish):
49 | m.act = Hardswish()
50 | elif isinstance(m.act, nn.SiLU):
51 | m.act = SiLU()
52 | # elif isinstance(m, models.yolo.Detect):
53 | # m.forward = m.forward_export # assign forward (optional)
54 | model.model[-1].export = True # set Detect() layer export=True
55 | y = model(img) # dry run
56 |
57 | # TorchScript export
58 | try:
59 | print('\nStarting TorchScript export with torch %s...' % torch.__version__)
60 | f = opt.weights.replace('.pt', '.torchscript.pt') # filename
61 | ts = torch.jit.trace(model, img)
62 | ts.save(f)
63 | print('TorchScript export success, saved as %s' % f)
64 | except Exception as e:
65 | print('TorchScript export failure: %s' % e)
66 |
67 | # ONNX export
68 | try:
69 | import onnx
70 |
71 | print('\nStarting ONNX export with onnx %s...' % onnx.__version__)
72 | f = opt.weights.replace('.pt', '.onnx') # filename
73 | torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'],
74 | output_names=['classes', 'boxes'] if y is None else ['output'],
75 | dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # size(1,3,640,640)
76 | 'output': {0: 'batch', 2: 'y', 3: 'x'}} if opt.dynamic else None)
77 |
78 | # Checks
79 | onnx_model = onnx.load(f) # load onnx model
80 | onnx.checker.check_model(onnx_model) # check onnx model
81 | # print(onnx.helper.printable_graph(onnx_model.graph)) # print a human readable model
82 | print('ONNX export success, saved as %s' % f)
83 | except Exception as e:
84 | print('ONNX export failure: %s' % e)
85 |
86 | # CoreML export
87 | try:
88 | import coremltools as ct
89 |
90 | print('\nStarting CoreML export with coremltools %s...' % ct.__version__)
91 | # convert model from torchscript and apply pixel scaling as per detect.py
92 | model = ct.convert(ts, inputs=[ct.ImageType(name='image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])])
93 | f = opt.weights.replace('.pt', '.mlmodel') # filename
94 | model.save(f)
95 | print('CoreML export success, saved as %s' % f)
96 | except Exception as e:
97 | print('CoreML export failure: %s' % e)
98 |
99 | # Finish
100 | print('\nExport complete (%.2fs). Visualize with https://github.com/lutzroeder/netron.' % (time.time() - t))
101 |
--------------------------------------------------------------------------------
/models/yolo.py:
--------------------------------------------------------------------------------
1 | import argparse
2 | import logging
3 | import sys
4 | from copy import deepcopy
5 | from pathlib import Path
6 |
7 | sys.path.append('./') # to run '$ python *.py' files in subdirectories
8 | logger = logging.getLogger(__name__)
9 |
10 | from models.common import *
11 | from models.experimental import MixConv2d, CrossConv
12 | from utils.autoanchor import check_anchor_order
13 | from utils.general import make_divisible, check_file, set_logging
14 | from utils.torch_utils import time_synchronized, fuse_conv_and_bn, model_info, scale_img, initialize_weights, \
15 | select_device, copy_attr
16 |
17 | try:
18 | import thop # for FLOPS computation
19 | except ImportError:
20 | thop = None
21 |
22 |
23 | class Detect(nn.Module):
24 | stride = None # strides computed during build
25 | export = False # onnx export
26 |
27 | def __init__(self, nc=80, anchors=(), ch=()): # detection layer
28 | super(Detect, self).__init__()
29 | self.nc = nc # number of classes
30 | self.no = nc + 5 # number of outputs per anchor
31 | self.nl = len(anchors) # number of detection layers
32 | self.na = len(anchors[0]) // 2 # number of anchors
33 | self.grid = [torch.zeros(1)] * self.nl # init grid
34 | a = torch.tensor(anchors).float().view(self.nl, -1, 2)
35 | self.register_buffer('anchors', a) # shape(nl,na,2)
36 | self.register_buffer('anchor_grid', a.clone().view(self.nl, 1, -1, 1, 1, 2)) # shape(nl,1,na,1,1,2)
37 | # self.m0=nn.ModuleList(MHSA(x, y, y) for x,y in zip(ch,[32,16,8]))
38 | # self.m1=nn.ModuleList(nn.BatchNorm2d(x) for x in ch)
39 | # self.m2 = nn.ModuleList(nn.Sigmoid() for _ in ch)
40 | self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) # output conv
41 |
42 | def forward(self, x):
43 | # x = x.copy() # for profiling
44 | z = [] # inference output
45 | self.training |= self.export
46 | for i in range(self.nl):
47 | # x[i] = self.m0[i](x[i]) + x[i]
48 | # x[i] = self.m1[i](x[i])
49 | # x[i] = self.m2[i](x[i])
50 | x[i] = self.m[i](x[i]) # conv
51 | bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85)
52 | x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
53 |
54 | if not self.training: # inference
55 | if self.grid[i].shape[2:4] != x[i].shape[2:4]:
56 | self.grid[i] = self._make_grid(nx, ny).to(x[i].device)
57 |
58 | y = x[i].sigmoid()
59 | y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i].to(x[i].device)) * self.stride[i] # xy
60 | y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
61 | z.append(y.view(bs, -1, self.no))
62 |
63 | return x if self.training else (torch.cat(z, 1), x)
64 |
65 | @staticmethod
66 | def _make_grid(nx=20, ny=20):
67 | yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)])
68 | return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float()
69 |
70 |
71 | class Model(nn.Module):
72 | def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, anchors=None): # model, input channels, number of classes
73 | super(Model, self).__init__()
74 |
75 | if isinstance(cfg, dict):
76 | self.yaml = cfg # model dict
77 | else: # is *.yaml
78 | import yaml # for torch hub
79 | self.yaml_file = Path(cfg).name
80 | with open(cfg) as f:
81 | self.yaml = yaml.load(f, Loader=yaml.SafeLoader) # model dict
82 |
83 | # Define model
84 | ch = self.yaml['ch'] = self.yaml.get('ch', ch) # input channels
85 | if nc and nc != self.yaml['nc']:
86 | logger.info(f"Overriding model.yaml nc={self.yaml['nc']} with nc={nc}")
87 | self.yaml['nc'] = nc # override yaml value
88 | if anchors:
89 | logger.info(f'Overriding model.yaml anchors with anchors={anchors}')
90 | self.yaml['anchors'] = round(anchors) # override yaml value
91 | self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) # model, savelist
92 | self.names = [str(i) for i in range(self.yaml['nc'])] # default names
93 | # print([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))])
94 |
95 | # Build strides, anchors
96 | m = self.model[-1] # Detect()
97 | if isinstance(m, Detect):
98 | s = 512 # 2x min stride
99 | m.stride = torch.tensor([s / x.shape[-2] for x in self.forward(torch.zeros(1, ch, s, s))[0]]) # forward
100 | m.anchors /= m.stride.view(-1, 1, 1)
101 | check_anchor_order(m)
102 | self.stride = m.stride
103 | self._initialize_biases() # only run once
104 | # print('Strides: %s' % m.stride.tolist())
105 |
106 | # Init weights, biases
107 | initialize_weights(self)
108 | self.info()
109 | logger.info('')
110 |
111 |
112 |
113 | def forward(self, x, augment=False, profile=False):
114 | if augment:
115 | img_size = x.shape[-2:] # height, width
116 | s = [1, 0.83, 0.67] # scales
117 | f = [None, 3, None] # flips (2-ud, 3-lr)
118 | y = [] # outputs
119 | for si, fi in zip(s, f):
120 | xi = scale_img(x.flip(fi) if fi else x, si, gs=int(self.stride.max()))
121 | yi = self.forward_once(xi)[0][0] # forward
122 | # cv2.imwrite('img%g.jpg' % s, 255 * xi[0].numpy().transpose((1, 2, 0))[:, :, ::-1]) # save
123 | yi[..., :4] /= si # de-scale
124 | if fi == 2:
125 | yi[..., 1] = img_size[0] - yi[..., 1] # de-flip ud
126 | elif fi == 3:
127 | yi[..., 0] = img_size[1] - yi[..., 0] # de-flip lr
128 | y.append(yi)
129 | return torch.cat(y, 1), None # augmented inference, train
130 | else:
131 | return self.forward_once(x, profile) # single-scale inference, train
132 |
133 | def forward_once(self, x, profile=False):
134 | multi = None
135 | y, dt = [], [] # outputs
136 | for m in self.model:
137 | if m.f != -1: # if not from previous layer
138 | 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
139 |
140 | if profile:
141 | o = thop.profile(m, inputs=(x,), verbose=False)[0] / 1E9 * 2 if thop else 0 # FLOPS
142 | t = time_synchronized()
143 | for _ in range(10):
144 | _ = m(x)
145 | dt.append((time_synchronized() - t) * 100)
146 | print('%10.1f%10.0f%10.1fms %-40s' % (o, m.np, dt[-1], m.type))
147 |
148 | x = m(x) # run
149 |
150 | if 'RCAN' in m.type:
151 | multi = x
152 |
153 | y.append(x if m.i in self.save else None) # save output
154 |
155 | if profile:
156 | print('%.1fms total' % sum(dt))
157 | return x, multi
158 |
159 | def _initialize_biases(self, cf=None): # initialize biases into Detect(), cf is class frequency
160 | # https://arxiv.org/abs/1708.02002 section 3.3
161 | # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1.
162 | m = self.model[-1] # Detect() module
163 | for mi, s in zip(m.m, m.stride): # from
164 | b = mi.bias.view(m.na, -1) # conv.bias(255) to (3,85)
165 | b.data[:, 4] += math.log(8 / (640 / s) ** 2) # obj (8 objects per 640 image)
166 | b.data[:, 5:] += math.log(0.6 / (m.nc - 0.99)) if cf is None else torch.log(cf / cf.sum()) # cls
167 | mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
168 |
169 | def _print_biases(self):
170 | m = self.model[-1] # Detect() module
171 | for mi in m.m: # from
172 | b = mi.bias.detach().view(m.na, -1).T # conv.bias(255) to (3,85)
173 | print(('%6g Conv2d.bias:' + '%10.3g' * 6) % (mi.weight.shape[1], *b[:5].mean(1).tolist(), b[5:].mean()))
174 |
175 | # def _print_weights(self):
176 | # for m in self.model.modules():
177 | # if type(m) is Bottleneck:
178 | # print('%10.3g' % (m.w.detach().sigmoid() * 2)) # shortcut weights
179 |
180 | def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers
181 | print('Fusing layers... ')
182 | for m in self.model.modules():
183 | if type(m) is Conv and hasattr(m, 'bn'):
184 | m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv
185 | delattr(m, 'bn') # remove batchnorm
186 | m.forward = m.fuseforward # update forward
187 | elif type(m) is RCAN:
188 | m.body = nn.Sequential()
189 | m.tail = nn.Sequential()
190 | self.info()
191 | return self
192 |
193 | def nms(self, mode=True): # add or remove NMS module
194 | present = type(self.model[-1]) is NMS # last layer is NMS
195 | if mode and not present:
196 | print('Adding NMS... ')
197 | m = NMS() # module
198 | m.f = -1 # from
199 | m.i = self.model[-1].i + 1 # index
200 | self.model.add_module(name='%s' % m.i, module=m) # add
201 | self.eval()
202 | elif not mode and present:
203 | print('Removing NMS... ')
204 | self.model = self.model[:-1] # remove
205 | return self
206 |
207 | def autoshape(self): # add autoShape module
208 | print('Adding autoShape... ')
209 | m = autoShape(self) # wrap model
210 | copy_attr(m, self, include=('yaml', 'nc', 'hyp', 'names', 'stride'), exclude=()) # copy attributes
211 | return m
212 |
213 | def info(self, verbose=False, img_size=512): # print model information
214 | model_info(self, verbose, img_size)
215 |
216 |
217 | def parse_model(d, ch): # model_dict, input_channels(3)
218 | logger.info('\n%3s%18s%3s%10s %-40s%-30s' % ('', 'from', 'n', 'params', 'module', 'arguments'))
219 | anchors, nc, gd, gw = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple']
220 | na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors # number of anchors
221 | no = na * (nc + 5) # number of outputs = anchors * (classes + 5)
222 |
223 | layers, save, c2 = [], [], ch[-1] # layers, savelist, ch out
224 | for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']): # from, number, module, args
225 | m = eval(m) if isinstance(m, str) else m # eval strings
226 | for j, a in enumerate(args):
227 | try:
228 | args[j] = eval(a) if isinstance(a, str) else a # eval strings
229 | except:
230 | pass
231 |
232 | n = max(round(n * gd), 1) if n > 1 else n # depth gain
233 | if m in [Conv, Bottleneck, SPP, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP, C3, C3ResAtnMHSA]:
234 | if f<0:
235 | c1, c2 = ch[f], args[0]
236 | else:
237 | c1, c2 = ch[f+1], args[0]
238 |
239 | # Normal
240 | # if i > 0 and args[0] != no: # channel expansion factor
241 | # ex = 1.75 # exponential (default 2.0)
242 | # e = math.log(c2 / ch[1]) / math.log(2)
243 | # c2 = int(ch[1] * ex ** e)
244 | # if m != Focus:
245 |
246 | c2 = make_divisible(c2 * gw, 8) if c2 != no else c2
247 |
248 | # Experimental
249 | # if i > 0 and args[0] != no: # channel expansion factor
250 | # ex = 1 + gw # exponential (default 2.0)
251 | # ch1 = 32 # ch[1]
252 | # e = math.log(c2 / ch1) / math.log(2) # level 1-n
253 | # c2 = int(ch1 * ex ** e)
254 | # if m != Focus:
255 | # c2 = make_divisible(c2, 8) if c2 != no else c2
256 |
257 | args = [c1, c2, *args[1:]]
258 | if m in [BottleneckCSP, C3, C3ResAtnMHSA]:
259 | args.insert(2, n)
260 | n = 1
261 | elif m in [nn.BatchNorm2d, SEAtn]:
262 | args = [ch[f]]
263 | c2=ch[f if f < 0 else f + 1]
264 | elif m in [Concat,ConcatFusionFactor]:
265 | c2 = sum([ch[x if x < 0 else x + 1] for x in f])
266 | elif m is Detect:
267 | args.append([ch[x + 1] for x in f])
268 | if isinstance(args[1], int): # number of anchors
269 | args[1] = [list(range(args[1] * 2))] * len(f)
270 | elif m is Contract:
271 | c2 = ch[f if f < 0 else f + 1] * args[0] ** 2
272 | elif m is Expand:
273 | c2 = ch[f if f < 0 else f + 1] // args[0] ** 2
274 | elif m is MHSA:
275 | args=[ch[f],*args[:]]
276 | c2=ch[f if f < 0 else f + 1]
277 | elif m is RCAN:
278 | args = [ch[f if f < 0 else f + 1]]
279 | c2 = ch[f if f < 0 else f + 1]
280 | elif m is AtnMut:
281 | args[0]=args[0]//2
282 | args[1] = args[1] // 2
283 | c2=ch[f[1] if f[1] < 0 else f[1] + 1]
284 | elif m is LocalReconstruct:
285 | args[0] = args[0] // 2
286 | args[1] = args[1] // 2
287 | c2 = ch[f[1] if f[1] < 0 else f[1] + 1]
288 | else:
289 | c2 = ch[f if f < 0 else f + 1]
290 |
291 | m_ = nn.Sequential(*[m(*args) for _ in range(n)]) if n > 1 else m(*args) # module
292 | t = str(m)[8:-2].replace('__main__.', '') # module type
293 | np = sum([x.numel() for x in m_.parameters()]) # number params
294 | m_.i, m_.f, m_.type, m_.np = i, f, t, np # attach index, 'from' index, type, number params
295 | logger.info('%3s%18s%3s%10.0f %-40s%-30s' % (i, f, n, np, t, args)) # print
296 | save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist
297 | layers.append(m_)
298 | ch.append(c2)
299 | return nn.Sequential(*layers), sorted(save)
300 |
301 |
302 | if __name__ == '__main__':
303 | parser = argparse.ArgumentParser()
304 | parser.add_argument('--cfg', type=str, default='yolov5s.yaml', help='model.yaml')
305 | parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
306 | opt = parser.parse_args()
307 | opt.cfg = check_file(opt.cfg) # check file
308 | set_logging()
309 | device = select_device(opt.device)
310 |
311 | # Create model
312 | model = Model(opt.cfg).to(device)
313 | model.train()
314 |
315 | # Profile
316 | # img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device)
317 | # y = model(img, profile=True)
318 |
319 | # Tensorboard
320 | # from torch.utils.tensorboard import SummaryWriter
321 | # tb_writer = SummaryWriter()
322 | # print("Run 'tensorboard --logdir=models/runs' to view tensorboard at http://localhost:6006/")
323 | # tb_writer.add_graph(model.model, img) # add model to tensorboard
324 | # tb_writer.add_image('test', img[0], dataformats='CWH') # add model to tensorboard
325 |
--------------------------------------------------------------------------------
/models/yolov5s.yaml:
--------------------------------------------------------------------------------
1 | # parameters
2 | nc: 80 # number of classes
3 | depth_multiple: 0.33 # model depth multiple
4 | width_multiple: 0.50 # layer channel multiple
5 |
6 | # anchors
7 | anchors:
8 | - [10,13, 16,30, 33,23] # P3/8
9 | - [30,61, 62,45, 59,119] # P4/16
10 | - [116,90, 156,198, 373,326] # P5/32
11 |
12 | # YOLOv5 backbone
13 | backbone:
14 | # [from, number, module, args]
15 | [[-1, 1, Focus, [64, 3]], # 0-P1/2
16 | [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
17 | [-1, 3, C3, [128]],
18 | [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
19 | [-1, 9, C3, [256]],
20 | [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
21 | [-1, 9, C3, [512]],
22 | [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
23 | [-1, 1, SPP, [1024, [5, 9, 13]]],
24 | [-1, 3, C3, [1024, False]], # 9
25 | ]
26 |
27 | # YOLOv5 head
28 | head:
29 | [[-1, 1, Conv, [512, 1, 1]],
30 | [-1, 1, nn.Upsample, [None, 2, 'nearest']],
31 | [[-1, 6], 1, Concat, [1]], # cat backbone P4
32 | [-1, 3, C3, [512, False]], # 13
33 |
34 | [-1, 1, Conv, [256, 1, 1]],
35 | [-1, 1, nn.Upsample, [None, 2, 'nearest']],
36 | [[-1, 4], 1, Concat, [1]], # cat backbone P3
37 | [-1, 3, C3, [256, False]], # 17 (P3/8-small)
38 |
39 | [-1, 1, Conv, [256, 3, 2]],
40 | [[-1, 14], 1, Concat, [1]], # cat head P4
41 | [-1, 3, C3, [512, False]], # 20 (P4/16-medium)
42 |
43 | [-1, 1, Conv, [512, 3, 2]],
44 | [[-1, 10], 1, Concat, [1]], # cat head P5
45 | [-1, 3, C3, [1024, False]], # 23 (P5/32-large)
46 |
47 | [[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
48 | ]
49 |
--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | # pip install -r requirements.txt
2 |
3 | # base ----------------------------------------
4 | Cython
5 | matplotlib>=3.2.2
6 | numpy>=1.18.5
7 | opencv-python>=4.1.2
8 | Pillow
9 | PyYAML>=5.3
10 | scipy>=1.4.1
11 | tensorboard>=2.2
12 | torch>=1.7.0
13 | torchvision>=0.8.1
14 | tqdm>=4.41.0
15 |
16 | # logging -------------------------------------
17 | wandb
18 |
19 | # plotting ------------------------------------
20 | seaborn>=0.11.0
21 | pandas
22 |
23 | # export --------------------------------------
24 | # coremltools==4.0
25 | # onnx>=1.8.0
26 | # scikit-learn==0.19.2 # for coreml quantization
27 |
28 | # extras --------------------------------------
29 | thop # FLOPS computation
30 | pycocotools>=2.0 # COCO mAP
31 |
--------------------------------------------------------------------------------
/test.py:
--------------------------------------------------------------------------------
1 | import argparse
2 | import json
3 | import os
4 | from pathlib import Path
5 | from threading import Thread
6 |
7 | import numpy as np
8 | import torch
9 | import yaml
10 | from tqdm import tqdm
11 |
12 | from models.experimental import attempt_load
13 | from utils.datasets import create_dataloader
14 | from utils.general import coco80_to_coco91_class, check_dataset, check_file, check_img_size, check_requirements, \
15 | box_iou, non_max_suppression, scale_coords, xyxy2xywh, xywh2xyxy, set_logging, increment_path, colorstr
16 | from utils.metrics import ap_per_class, ConfusionMatrix
17 | from utils.plots import plot_images, output_to_target, plot_study_txt, plot_images_each, plot_images_together
18 | from utils.torch_utils import select_device, time_synchronized
19 | import time
20 |
21 | def test(data,
22 | weights=None,
23 | batch_size=32,
24 | imgsz=640,
25 | conf_thres=0.001,
26 | iou_thres=0.6, # for NMS
27 | save_json=False,
28 | single_cls=False,
29 | augment=False,
30 | verbose=False,
31 | model=None,
32 | dataloader=None,
33 | save_dir=Path(''), # for saving images
34 | save_txt=False, # for auto-labelling
35 | save_hybrid=False, # for hybrid auto-labelling
36 | save_conf=False, # save auto-label confidences
37 | plots=True,
38 | log_imgs=0, # number of logged images
39 | compute_loss=None):
40 |
41 | plot_batch_num = 3 # How many batches you want to plot in test phase.
42 |
43 | # Initialize/load model and set device
44 | training = model is not None
45 | if training: # called by train.py
46 | device = next(model.parameters()).device # get model device
47 |
48 | else: # called directly
49 | set_logging()
50 | device = select_device(opt.device, batch_size=batch_size)
51 |
52 | # Directories
53 | save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) # increment run
54 | (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
55 |
56 | # Load model
57 | model = attempt_load(weights, map_location=device) # load FP32 model
58 | imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
59 |
60 | # Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99
61 | # if device.type != 'cpu' and torch.cuda.device_count() > 1:
62 | # model = nn.DataParallel(model)
63 |
64 | # Half
65 | half = device.type != 'cpu' # half precision only supported on CUDA
66 | if half:
67 | model.half()
68 |
69 | # Configure
70 | model.eval()
71 | is_coco = data.endswith('coco.yaml') # is COCO dataset
72 | with open(data) as f:
73 | data = yaml.load(f, Loader=yaml.SafeLoader) # model dict
74 | check_dataset(data) # check
75 | nc = 1 if single_cls else int(data['nc']) # number of classes
76 | iouv = torch.linspace(0.5, 0.95, 10).to(device) # iou vector for mAP@0.5:0.95
77 | niou = iouv.numel()
78 |
79 | # Logging
80 | log_imgs, wandb = min(log_imgs, 100), None # ceil
81 | try:
82 | import wandb # Weights & Biases
83 | except ImportError:
84 | log_imgs = 0
85 |
86 | # Dataloader
87 | if not training:
88 | img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
89 | _ = model(img.half() if half else img) if device.type != 'cpu' else None # run once
90 | path = data['test'] if opt.task == 'test' else data['val'] # path to val/test images
91 | dataloader = create_dataloader(path, imgsz, batch_size, model.stride.max(), opt, pad=0.0, rect=True,
92 | prefix=colorstr('test: ' if opt.task == 'test' else 'val: '))[0]
93 |
94 | seen = 0
95 | confusion_matrix = ConfusionMatrix(nc=nc)
96 | names = {k: v for k, v in enumerate(model.names if hasattr(model, 'names') else model.module.names)}
97 | coco91class = coco80_to_coco91_class()
98 | s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R', 'mAP@.5', 'mAP@.5:.95')
99 | p, r, f1, mp, mr, map50, map, t0, t1 = 0., 0., 0., 0., 0., 0., 0., 0., 0.
100 | loss = torch.zeros(4, device=device)
101 | jdict, stats, ap, ap_class, wandb_images = [], [], [], [], []
102 | alltime=0
103 | for batch_i, (img, dgimgs, targets, paths, shapes) in enumerate(tqdm(dataloader, desc=s)):
104 | img = img.to(device, non_blocking=True)
105 | img = img.half() if half else img.float() # uint8 to fp16/32
106 | img /= 255.0 # 0 - 255 to 0.0 - 1.0
107 | targets = targets.to(device)
108 | dgimgs = dgimgs.to(device)
109 | nb, _, height, width = img.shape # batch size, channels, height, width
110 |
111 | with torch.no_grad():
112 | # Run model
113 | t = time_synchronized()
114 | start = time.time()
115 | (out, train_out), pdg = model(img, augment=augment) # inference and training outputs
116 | alltime+=time.time()-start
117 | t0 += time_synchronized() - t
118 |
119 | # Compute loss
120 | if compute_loss:
121 | loss += compute_loss(([x.float() for x in train_out], pdg), dgimgs, targets)[1][:4] # box, obj, cls
122 |
123 | # Run NMS
124 | targets[:, 2:] *= torch.Tensor([width, height, width, height]).to(device) # to pixels
125 | lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else [] # for autolabelling
126 | t = time_synchronized()
127 | out = non_max_suppression(out, conf_thres=conf_thres, iou_thres=iou_thres, labels=lb, multi_label=True)
128 | t1 += time_synchronized() - t
129 |
130 | # Statistics per image
131 | for si, pred in enumerate(out):
132 | labels = targets[targets[:, 0] == si, 1:]
133 | nl = len(labels)
134 | tcls = labels[:, 0].tolist() if nl else [] # target class
135 | path = Path(paths[si])
136 | seen += 1
137 |
138 | if len(pred) == 0:
139 | if nl:
140 | stats.append((torch.zeros(0, niou, dtype=torch.bool), torch.Tensor(), torch.Tensor(), tcls))
141 | continue
142 |
143 | # Predictions
144 | predn = pred.clone()
145 | scale_coords(img[si].shape[1:], predn[:, :4], shapes[si][0], shapes[si][1]) # native-space pred
146 |
147 | # Append to text file
148 | if save_txt:
149 | gn = torch.tensor(shapes[si][0])[[1, 0, 1, 0]] # normalization gain whwh
150 | for *xyxy, conf, cls in predn.tolist():
151 | xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
152 | line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # label format
153 | with open(save_dir / 'labels' / (path.stem + '.txt'), 'a') as f:
154 | f.write(('%g ' * len(line)).rstrip() % line + '\n')
155 |
156 | # W&B logging
157 | if plots and len(wandb_images) < log_imgs:
158 | box_data = [{"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
159 | "class_id": int(cls),
160 | "box_caption": "%s %.3f" % (names[cls], conf),
161 | "scores": {"class_score": conf},
162 | "domain": "pixel"} for *xyxy, conf, cls in pred.tolist()]
163 | boxes = {"predictions": {"box_data": box_data, "class_labels": names}} # inference-space
164 | if img.shape[1]!=3:
165 | wandb_images.append(wandb.Image(img[si][:3,:,:], boxes=boxes, caption=path.name))
166 | else:
167 | wandb_images.append(wandb.Image(img[si], boxes=boxes, caption=path.name))
168 |
169 | # Append to pycocotools JSON dictionary
170 | if save_json:
171 | # [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ...
172 | image_id = int(path.stem) if path.stem.isnumeric() else path.stem
173 | box = xyxy2xywh(predn[:, :4]) # xywh
174 | box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner
175 | for p, b in zip(pred.tolist(), box.tolist()):
176 | jdict.append({'image_id': image_id,
177 | 'category_id': coco91class[int(p[5])] if is_coco else int(p[5]),
178 | 'bbox': [round(x, 3) for x in b],
179 | 'score': round(p[4], 5)})
180 |
181 | # Assign all predictions as incorrect
182 | correct = torch.zeros(pred.shape[0], niou, dtype=torch.bool, device=device)
183 | if nl:
184 | detected = [] # target indices
185 | tcls_tensor = labels[:, 0]
186 |
187 | # target boxes
188 | tbox = xywh2xyxy(labels[:, 1:5])
189 | scale_coords(img[si].shape[1:], tbox, shapes[si][0], shapes[si][1]) # native-space labels
190 | if plots:
191 | confusion_matrix.process_batch(predn, torch.cat((labels[:, 0:1], tbox), 1))
192 |
193 | # Per target class
194 | for cls in torch.unique(tcls_tensor):
195 | ti = (cls == tcls_tensor).nonzero(as_tuple=False).view(-1) # prediction indices
196 | pi = (cls == pred[:, 5]).nonzero(as_tuple=False).view(-1) # target indices
197 |
198 | # Search for detections
199 | if pi.shape[0]:
200 | # Prediction to target ious
201 | ious, i = box_iou(predn[pi, :4], tbox[ti]).max(1) # best ious, indices
202 |
203 | # Append detections
204 | detected_set = set()
205 | for j in (ious > iouv[0]).nonzero(as_tuple=False):
206 | d = ti[i[j]] # detected target
207 | if d.item() not in detected_set:
208 | detected_set.add(d.item())
209 | detected.append(d)
210 | correct[pi[j]] = ious[j] > iouv # iou_thres is 1xn
211 | if len(detected) == nl: # all targets already located in image
212 | break
213 |
214 | # Append statistics (correct, conf, pcls, tcls)
215 | stats.append((correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls))
216 | # Plot images
217 | if plots and batch_i < plot_batch_num:
218 | f = save_dir / f'test_batch{batch_i}_labels_pred.jpg' #altogether
219 | Thread(target=plot_images_together, args=(img, targets, output_to_target(out), paths, f, names), daemon=True).start()
220 | # f = save_dir / f'test_batch{batch_i}_labels.jpg' # labels
221 | # Thread(target=plot_images_each, args=(img, targets, paths, f, names), daemon=True).start()
222 | # f = save_dir / f'test_batch{batch_i}_pred.jpg' # predictions
223 | # if pdg!=None:
224 | # Thread(target=plot_images_each, args=(img, output_to_target(out), paths, f, names), daemon=True).start()
225 | # else:
226 | # Thread(target=plot_images_each, args=(img, output_to_target(out), paths, f, names), daemon=True).start()
227 | print(alltime)
228 |
229 | # Compute statistics
230 | stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy
231 | if len(stats) and stats[0].any():
232 | p, r, ap, f1, ap_class = ap_per_class(*stats, plot=plots, save_dir=save_dir, names=names)
233 | p, r, ap50, ap = p[:, 0], r[:, 0], ap[:, 0], ap.mean(1) # [P, R, AP@0.5, AP@0.5:0.95]
234 | mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
235 | nt = np.bincount(stats[3].astype(np.int64), minlength=nc) # number of targets per class
236 | else:
237 | nt = torch.zeros(1)
238 |
239 | # Print results
240 | pf = '%20s' + '%12.3g' * 6 # print format
241 | print(pf % ('all', seen, nt.sum(), mp, mr, map50, map))
242 |
243 | # Print results per class
244 | if (verbose or (nc < 50 and not training)) and nc > 1 and len(stats):
245 | for i, c in enumerate(ap_class):
246 | print(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i]))
247 |
248 | # Print speeds
249 | t = tuple(x / seen * 1E3 for x in (t0, t1, t0 + t1)) + (imgsz, imgsz, batch_size) # tuple
250 | if not training:
251 | print('Speed: %.1f/%.1f/%.1f ms inference/NMS/total per %gx%g image at batch-size %g' % t)
252 |
253 | # Plots
254 | if plots:
255 | confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))
256 | if wandb and wandb.run:
257 | wandb.log({"Images": wandb_images})
258 | wandb.log({"Validation": [wandb.Image(str(f), caption=f.name) for f in sorted(save_dir.glob('test*.jpg'))]})
259 |
260 | # Save JSON
261 | if save_json and len(jdict):
262 | w = Path(weights[0] if isinstance(weights, list) else weights).stem if weights is not None else '' # weights
263 | anno_json = '../coco/annotations/instances_val2017.json' # annotations json
264 | pred_json = str(save_dir / f"{w}_predictions.json") # predictions json
265 | print('\nEvaluating pycocotools mAP... saving %s...' % pred_json)
266 | with open(pred_json, 'w') as f:
267 | json.dump(jdict, f)
268 |
269 | try: # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
270 | from pycocotools.coco import COCO
271 | from pycocotools.cocoeval import COCOeval
272 |
273 | anno = COCO(anno_json) # init annotations api
274 | pred = anno.loadRes(pred_json) # init predictions api
275 | eval = COCOeval(anno, pred, 'bbox')
276 | if is_coco:
277 | eval.params.imgIds = [int(Path(x).stem) for x in dataloader.dataset.img_files] # image IDs to evaluate
278 | eval.evaluate()
279 | eval.accumulate()
280 | eval.summarize()
281 | map, map50 = eval.stats[:2] # update results (mAP@0.5:0.95, mAP@0.5)
282 | except Exception as e:
283 | print(f'pycocotools unable to run: {e}')
284 |
285 | # Return results
286 | if not training:
287 | s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
288 | print(f"Results saved to {save_dir}{s}")
289 | model.float() # for training
290 | maps = np.zeros(nc) + map
291 | for i, c in enumerate(ap_class):
292 | maps[c] = ap[i]
293 | return (mp, mr, map50, map, *(loss.cpu() / len(dataloader)).tolist()), maps, t
294 |
295 |
296 | if __name__ == '__main__':
297 | parser = argparse.ArgumentParser(prog='test.py')
298 | parser.add_argument('--weights', nargs='+', type=str, default=r'.\DRENet.pt', help='model.pt path(s)')
299 | parser.add_argument('--data', type=str, default='data/ship.yaml', help='*.data path')
300 | parser.add_argument('--batch-size', type=int, default=16, help='size of each image batch')
301 | parser.add_argument('--img-size', type=int, default=512, help='inference size (pixels)')
302 | parser.add_argument('--conf-thres', type=float, default=0.001, help='object confidence threshold')
303 | parser.add_argument('--iou-thres', type=float, default=0.6, help='IOU threshold for NMS')
304 | parser.add_argument('--task', default='val', help="'val', 'test', 'study'")
305 | parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
306 | parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset')
307 | parser.add_argument('--augment', action='store_true', help='augmented inference')
308 | parser.add_argument('--verbose', action='store_true', help='report mAP by class')
309 | parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
310 | parser.add_argument('--save-hybrid', action='store_true', help='save label+prediction hybrid results to *.txt')
311 | parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
312 | parser.add_argument('--save-json', action='store_true', help='save a cocoapi-compatible JSON results file')
313 | parser.add_argument('--project', default='runs/test', help='save to project/name')
314 | parser.add_argument('--name', default='test', help='save to project/name')
315 | parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
316 | opt = parser.parse_args()
317 | opt.save_json |= opt.data.endswith('coco.yaml')
318 | opt.data = check_file(opt.data) # check file
319 | print(opt)
320 | check_requirements()
321 |
322 | if opt.task in ['val', 'test']: # run normally
323 | test(opt.data,
324 | opt.weights,
325 | opt.batch_size,
326 | opt.img_size,
327 | opt.conf_thres,
328 | opt.iou_thres,
329 | opt.save_json,
330 | opt.single_cls,
331 | opt.augment,
332 | opt.verbose,
333 | save_txt=opt.save_txt | opt.save_hybrid,
334 | save_hybrid=opt.save_hybrid,
335 | save_conf=opt.save_conf,
336 | )
337 |
338 | elif opt.task == 'study': # run over a range of settings and save/plot
339 | for weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']:
340 | f = 'study_%s_%s.txt' % (Path(opt.data).stem, Path(weights).stem) # filename to save to
341 | x = list(range(320, 800, 64)) # x axis
342 | y = [] # y axis
343 | for i in x: # img-size
344 | print('\nRunning %s point %s...' % (f, i))
345 | r, _, t = test(opt.data, weights, opt.batch_size, i, opt.conf_thres, opt.iou_thres, opt.save_json,
346 | plots=False)
347 | y.append(r + t) # results and times
348 | np.savetxt(f, y, fmt='%10.4g') # save
349 | os.system('zip -r study.zip study_*.txt')
350 | plot_study_txt(f, x) # plot
351 |
--------------------------------------------------------------------------------
/utils/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/WindVChen/DRENet/a187dbe0f623b521a62c6176c7cafaa7322f5f66/utils/__init__.py
--------------------------------------------------------------------------------
/utils/activations.py:
--------------------------------------------------------------------------------
1 | # Activation functions
2 |
3 | import torch
4 | import torch.nn as nn
5 | import torch.nn.functional as F
6 |
7 |
8 | # SiLU https://arxiv.org/pdf/1606.08415.pdf ----------------------------------------------------------------------------
9 | class SiLU(nn.Module): # export-friendly version of nn.SiLU()
10 | @staticmethod
11 | def forward(x):
12 | return x * torch.sigmoid(x)
13 |
14 |
15 | class Hardswish(nn.Module): # export-friendly version of nn.Hardswish()
16 | @staticmethod
17 | def forward(x):
18 | # return x * F.hardsigmoid(x) # for torchscript and CoreML
19 | return x * F.hardtanh(x + 3, 0., 6.) / 6. # for torchscript, CoreML and ONNX
20 |
21 |
22 | class MemoryEfficientSwish(nn.Module):
23 | class F(torch.autograd.Function):
24 | @staticmethod
25 | def forward(ctx, x):
26 | ctx.save_for_backward(x)
27 | return x * torch.sigmoid(x)
28 |
29 | @staticmethod
30 | def backward(ctx, grad_output):
31 | x = ctx.saved_tensors[0]
32 | sx = torch.sigmoid(x)
33 | return grad_output * (sx * (1 + x * (1 - sx)))
34 |
35 | def forward(self, x):
36 | return self.F.apply(x)
37 |
38 |
39 | # Mish https://github.com/digantamisra98/Mish --------------------------------------------------------------------------
40 | class Mish(nn.Module):
41 | @staticmethod
42 | def forward(x):
43 | return x * F.softplus(x).tanh()
44 |
45 |
46 | class MemoryEfficientMish(nn.Module):
47 | class F(torch.autograd.Function):
48 | @staticmethod
49 | def forward(ctx, x):
50 | ctx.save_for_backward(x)
51 | return x.mul(torch.tanh(F.softplus(x))) # x * tanh(ln(1 + exp(x)))
52 |
53 | @staticmethod
54 | def backward(ctx, grad_output):
55 | x = ctx.saved_tensors[0]
56 | sx = torch.sigmoid(x)
57 | fx = F.softplus(x).tanh()
58 | return grad_output * (fx + x * sx * (1 - fx * fx))
59 |
60 | def forward(self, x):
61 | return self.F.apply(x)
62 |
63 |
64 | # FReLU https://arxiv.org/abs/2007.11824 -------------------------------------------------------------------------------
65 | class FReLU(nn.Module):
66 | def __init__(self, c1, k=3): # ch_in, kernel
67 | super().__init__()
68 | self.conv = nn.Conv2d(c1, c1, k, 1, 1, groups=c1, bias=False)
69 | self.bn = nn.BatchNorm2d(c1)
70 |
71 | def forward(self, x):
72 | return torch.max(x, self.bn(self.conv(x)))
73 |
--------------------------------------------------------------------------------
/utils/autoanchor.py:
--------------------------------------------------------------------------------
1 | # Auto-anchor utils
2 |
3 | import numpy as np
4 | import torch
5 | import yaml
6 | from scipy.cluster.vq import kmeans
7 | from tqdm import tqdm
8 |
9 | from utils.general import colorstr
10 |
11 |
12 | def check_anchor_order(m):
13 | # Check anchor order against stride order for YOLOv5 Detect() module m, and correct if necessary
14 | a = m.anchor_grid.prod(-1).view(-1) # anchor area
15 | da = a[-1] - a[0] # delta a
16 | ds = m.stride[-1] - m.stride[0] # delta s
17 | if da.sign() != ds.sign(): # same order
18 | print('Reversing anchor order')
19 | m.anchors[:] = m.anchors.flip(0)
20 | m.anchor_grid[:] = m.anchor_grid.flip(0)
21 |
22 |
23 | def check_anchors(dataset, model, thr=4.0, imgsz=640):
24 | # Check anchor fit to data, recompute if necessary
25 | prefix = colorstr('autoanchor: ')
26 | print(f'\n{prefix}Analyzing anchors... ', end='')
27 | m = model.module.model[-1] if hasattr(model, 'module') else model.model[-1] # Detect()
28 | shapes = imgsz * dataset.shapes / dataset.shapes.max(1, keepdims=True)
29 | scale = np.random.uniform(0.9, 1.1, size=(shapes.shape[0], 1)) # augment scale
30 | if shapes.shape[1]!=2:
31 | wh = torch.tensor(np.concatenate([l[:, 3:5] * s[:-1] for s, l in zip(shapes * scale, dataset.labels)])).float() # wh
32 | else:
33 | wh = torch.tensor(np.concatenate([l[:, 3:5] * s for s, l in zip(shapes * scale, dataset.labels)])).float()
34 |
35 | def metric(k): # compute metric
36 | r = wh[:, None] / k[None]
37 | x = torch.min(r, 1. / r).min(2)[0] # ratio metric
38 | best = x.max(1)[0] # best_x
39 | aat = (x > 1. / thr).float().sum(1).mean() # anchors above threshold
40 | bpr = (best > 1. / thr).float().mean() # best possible recall
41 | return bpr, aat
42 |
43 | anchors = m.anchor_grid.clone().cpu().view(-1, 2) # current anchors
44 | bpr, aat = metric(anchors)
45 | print(f'anchors/target = {aat:.2f}, Best Possible Recall (BPR) = {bpr:.4f}', end='')
46 | bpr=0.97
47 | if bpr < 0.98: # threshold to recompute
48 | print('. Attempting to improve anchors, please wait...')
49 | na = m.anchor_grid.numel() // 2 # number of anchors
50 | try:
51 | anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr, gen=1000, verbose=False)
52 | except Exception as e:
53 | print(f'{prefix}ERROR: {e}')
54 | new_bpr = metric(anchors)[0]
55 | if new_bpr > bpr: # replace anchors
56 | anchors = torch.tensor(anchors, device=m.anchors.device).type_as(m.anchors)
57 | m.anchor_grid[:] = anchors.clone().view_as(m.anchor_grid) # for inference
58 | m.anchors[:] = anchors.clone().view_as(m.anchors) / m.stride.to(m.anchors.device).view(-1, 1, 1) # loss
59 | check_anchor_order(m)
60 | print(f'{prefix}New anchors saved to model. Update model *.yaml to use these anchors in the future.')
61 | else:
62 | print(f'{prefix}Original anchors better than new anchors. Proceeding with original anchors.')
63 | print('') # newline
64 |
65 |
66 | def kmean_anchors(path='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=1000, verbose=True):
67 | """ Creates kmeans-evolved anchors from training dataset
68 |
69 | Arguments:
70 | path: path to dataset *.yaml, or a loaded dataset
71 | n: number of anchors
72 | img_size: image size used for training
73 | thr: anchor-label wh ratio threshold hyperparameter hyp['anchor_t'] used for training, default=4.0
74 | gen: generations to evolve anchors using genetic algorithm
75 | verbose: print all results
76 |
77 | Return:
78 | k: kmeans evolved anchors
79 |
80 | Usage:
81 | from utils.autoanchor import *; _ = kmean_anchors()
82 | """
83 | thr = 1. / thr
84 | prefix = colorstr('autoanchor: ')
85 |
86 | def metric(k, wh): # compute metrics
87 | r = wh[:, None] / k[None]
88 | x = torch.min(r, 1. / r).min(2)[0] # ratio metric
89 | # x = wh_iou(wh, torch.tensor(k)) # iou metric
90 | return x, x.max(1)[0] # x, best_x
91 |
92 | def anchor_fitness(k): # mutation fitness
93 | _, best = metric(torch.tensor(k, dtype=torch.float32), wh)
94 | return (best * (best > thr).float()).mean() # fitness
95 |
96 | def print_results(k):
97 | k = k[np.argsort(k.prod(1))] # sort small to large
98 | x, best = metric(k, wh0)
99 | bpr, aat = (best > thr).float().mean(), (x > thr).float().mean() * n # best possible recall, anch > thr
100 | print(f'{prefix}thr={thr:.2f}: {bpr:.4f} best possible recall, {aat:.2f} anchors past thr')
101 | print(f'{prefix}n={n}, img_size={img_size}, metric_all={x.mean():.3f}/{best.mean():.3f}-mean/best, '
102 | f'past_thr={x[x > thr].mean():.3f}-mean: ', end='')
103 | for i, x in enumerate(k):
104 | print('%i,%i' % (round(x[0]), round(x[1])), end=', ' if i < len(k) - 1 else '\n') # use in *.cfg
105 | return k
106 |
107 | if isinstance(path, str): # *.yaml file
108 | with open(path) as f:
109 | data_dict = yaml.load(f, Loader=yaml.SafeLoader) # model dict
110 | from utils.datasets import LoadImagesAndLabels
111 | dataset = LoadImagesAndLabels(data_dict['train'], augment=True, rect=True)
112 | else:
113 | dataset = path # dataset
114 |
115 | # Get label wh
116 | shapes = img_size * dataset.shapes / dataset.shapes.max(1, keepdims=True)
117 | wh0 = np.concatenate([l[:, 3:5] * s for s, l in zip(shapes, dataset.labels)]) # wh
118 |
119 | # Filter
120 | i = (wh0 < 3.0).any(1).sum()
121 | if i:
122 | print(f'{prefix}WARNING: Extremely small objects found. {i} of {len(wh0)} labels are < 3 pixels in size.')
123 | wh = wh0[(wh0 >= 2.0).any(1)] # filter > 2 pixels
124 | # wh = wh * (np.random.rand(wh.shape[0], 1) * 0.9 + 0.1) # multiply by random scale 0-1
125 |
126 | # Kmeans calculation
127 | print(f'{prefix}Running kmeans for {n} anchors on {len(wh)} points...')
128 | s = wh.std(0) # sigmas for whitening
129 | k, dist = kmeans(wh / s, n, iter=30) # points, mean distance
130 | assert len(k) == n, print(f'{prefix}ERROR: scipy.cluster.vq.kmeans requested {n} points but returned only {len(k)}')
131 | k *= s
132 | wh = torch.tensor(wh, dtype=torch.float32) # filtered
133 | wh0 = torch.tensor(wh0, dtype=torch.float32) # unfiltered
134 | k = print_results(k)
135 |
136 | # Plot
137 | # k, d = [None] * 20, [None] * 20
138 | # for i in tqdm(range(1, 21)):
139 | # k[i-1], d[i-1] = kmeans(wh / s, i) # points, mean distance
140 | # fig, ax = plt.subplots(1, 2, figsize=(14, 7), tight_layout=True)
141 | # ax = ax.ravel()
142 | # ax[0].plot(np.arange(1, 21), np.array(d) ** 2, marker='.')
143 | # fig, ax = plt.subplots(1, 2, figsize=(14, 7)) # plot wh
144 | # ax[0].hist(wh[wh[:, 0]<100, 0],400)
145 | # ax[1].hist(wh[wh[:, 1]<100, 1],400)
146 | # fig.savefig('wh.png', dpi=200)
147 |
148 | # Evolve
149 | npr = np.random
150 | f, sh, mp, s = anchor_fitness(k), k.shape, 0.9, 0.1 # fitness, generations, mutation prob, sigma
151 | pbar = tqdm(range(gen), desc=f'{prefix}Evolving anchors with Genetic Algorithm:') # progress bar
152 | for _ in pbar:
153 | v = np.ones(sh)
154 | while (v == 1).all(): # mutate until a change occurs (prevent duplicates)
155 | v = ((npr.random(sh) < mp) * npr.random() * npr.randn(*sh) * s + 1).clip(0.3, 3.0)
156 | kg = (k.copy() * v).clip(min=2.0)
157 | fg = anchor_fitness(kg)
158 | if fg > f:
159 | f, k = fg, kg.copy()
160 | pbar.desc = f'{prefix}Evolving anchors with Genetic Algorithm: fitness = {f:.4f}'
161 | if verbose:
162 | print_results(k)
163 |
164 | return print_results(k)
165 |
--------------------------------------------------------------------------------
/utils/general.py:
--------------------------------------------------------------------------------
1 | # General utils
2 |
3 | import glob
4 | import logging
5 | import math
6 | import os
7 | import platform
8 | import random
9 | import re
10 | import subprocess
11 | import time
12 | from pathlib import Path
13 |
14 | import cv2
15 | import numpy as np
16 | import torch
17 | import torchvision
18 | import yaml
19 |
20 | from utils.google_utils import gsutil_getsize
21 | from utils.metrics import fitness
22 | from utils.torch_utils import init_torch_seeds
23 |
24 | # Settings
25 | torch.set_printoptions(linewidth=320, precision=5, profile='long')
26 | np.set_printoptions(linewidth=320, formatter={'float_kind': '{:11.5g}'.format}) # format short g, %precision=5
27 | cv2.setNumThreads(0) # prevent OpenCV from multithreading (incompatible with PyTorch DataLoader)
28 | os.environ['NUMEXPR_MAX_THREADS'] = str(min(os.cpu_count(), 8)) # NumExpr max threads
29 |
30 |
31 | def set_logging(rank=-1):
32 | logging.basicConfig(
33 | format="%(message)s",
34 | level=logging.INFO if rank in [-1, 0] else logging.WARN)
35 |
36 |
37 | def init_seeds(seed=0):
38 | # Initialize random number generator (RNG) seeds
39 | random.seed(seed)
40 | np.random.seed(seed)
41 | init_torch_seeds(seed)
42 |
43 |
44 | def get_latest_run(search_dir='.'):
45 | # Return path to most recent 'last.pt' in /runs (i.e. to --resume from)
46 | last_list = glob.glob(f'{search_dir}/**/last*.pt', recursive=True)
47 | return max(last_list, key=os.path.getctime) if last_list else ''
48 |
49 |
50 | def isdocker():
51 | # Is environment a Docker container
52 | return Path('/workspace').exists() # or Path('/.dockerenv').exists()
53 |
54 |
55 | def check_online():
56 | # Check internet connectivity
57 | import socket
58 | try:
59 | socket.create_connection(("1.1.1.1", 443), 5) # check host accesability
60 | return True
61 | except OSError:
62 | return False
63 |
64 |
65 | def check_git_status():
66 | # Recommend 'git pull' if code is out of date
67 | print(colorstr('github: '), end='')
68 | try:
69 | assert Path('.git').exists(), 'skipping check (not a git repository)'
70 | assert not isdocker(), 'skipping check (Docker image)'
71 | assert check_online(), 'skipping check (offline)'
72 |
73 | cmd = 'git fetch && git config --get remote.origin.url'
74 | url = subprocess.check_output(cmd, shell=True).decode().strip().rstrip('.git') # github repo url
75 | branch = subprocess.check_output('git rev-parse --abbrev-ref HEAD', shell=True).decode().strip() # checked out
76 | n = int(subprocess.check_output(f'git rev-list {branch}..origin/master --count', shell=True)) # commits behind
77 | if n > 0:
78 | s = f"⚠️ WARNING: code is out of date by {n} commit{'s' * (n > 1)}. " \
79 | f"Use 'git pull' to update or 'git clone {url}' to download latest."
80 | else:
81 | s = f'up to date with {url} ✅'
82 | print(s.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else s)
83 | except Exception as e:
84 | print(e)
85 |
86 |
87 | def check_requirements(file='requirements.txt', exclude=()):
88 | # Check installed dependencies meet requirements
89 | import pkg_resources
90 | requirements = [f'{x.name}{x.specifier}' for x in pkg_resources.parse_requirements(Path(file).open())
91 | if x.name not in exclude]
92 | pkg_resources.require(requirements) # DistributionNotFound or VersionConflict exception if requirements not met
93 |
94 |
95 | def check_img_size(img_size, s=32):
96 | # Verify img_size is a multiple of stride s
97 | new_size = make_divisible(img_size, int(s)) # ceil gs-multiple
98 | if new_size != img_size:
99 | print('WARNING: --img-size %g must be multiple of max stride %g, updating to %g' % (img_size, s, new_size))
100 | return new_size
101 |
102 |
103 | def check_imshow():
104 | # Check if environment supports image displays
105 | try:
106 | assert not isdocker(), 'cv2.imshow() is disabled in Docker environments'
107 | cv2.imshow('test', np.zeros((1, 1, 3)))
108 | cv2.waitKey(1)
109 | cv2.destroyAllWindows()
110 | cv2.waitKey(1)
111 | return True
112 | except Exception as e:
113 | print(f'WARNING: Environment does not support cv2.imshow() or PIL Image.show() image displays\n{e}')
114 | return False
115 |
116 |
117 | def check_file(file):
118 | # Search for file if not found
119 | if os.path.isfile(file) or file == '':
120 | return file
121 | else:
122 | files = glob.glob('./**/' + file, recursive=True) # find file
123 | assert len(files), 'File Not Found: %s' % file # assert file was found
124 | assert len(files) == 1, "Multiple files match '%s', specify exact path: %s" % (file, files) # assert unique
125 | return files[0] # return file
126 |
127 |
128 | def check_dataset(dict):
129 | # Download dataset if not found locally
130 | val, s = dict.get('val'), dict.get('download')
131 | if val and len(val):
132 | val = [Path(x).resolve() for x in (val if isinstance(val, list) else [val])] # val path
133 | if not all(x.exists() for x in val):
134 | print('\nWARNING: Dataset not found, nonexistent paths: %s' % [str(x) for x in val if not x.exists()])
135 | if s and len(s): # download script
136 | print('Downloading %s ...' % s)
137 | if s.startswith('http') and s.endswith('.zip'): # URL
138 | f = Path(s).name # filename
139 | torch.hub.download_url_to_file(s, f)
140 | r = os.system('unzip -q %s -d ../ && rm %s' % (f, f)) # unzip
141 | else: # bash script
142 | r = os.system(s)
143 | print('Dataset autodownload %s\n' % ('success' if r == 0 else 'failure')) # analyze return value
144 | else:
145 | raise Exception('Dataset not found.')
146 |
147 |
148 | def make_divisible(x, divisor):
149 | # Returns x evenly divisible by divisor
150 | return math.ceil(x / divisor) * divisor
151 |
152 |
153 | def clean_str(s):
154 | # Cleans a string by replacing special characters with underscore _
155 | return re.sub(pattern="[|@#!¡·$€%&()=?¿^*;:,¨´><+]", repl="_", string=s)
156 |
157 |
158 | def one_cycle(y1=0.0, y2=1.0, steps=100):
159 | # lambda function for sinusoidal ramp from y1 to y2
160 | return lambda x: ((1 - math.cos(x * math.pi / steps)) / 2) * (y2 - y1) + y1
161 |
162 |
163 | def colorstr(*input):
164 | # Colors a string https://en.wikipedia.org/wiki/ANSI_escape_code, i.e. colorstr('blue', 'hello world')
165 | *args, string = input if len(input) > 1 else ('blue', 'bold', input[0]) # color arguments, string
166 | colors = {'black': '\033[30m', # basic colors
167 | 'red': '\033[31m',
168 | 'green': '\033[32m',
169 | 'yellow': '\033[33m',
170 | 'blue': '\033[34m',
171 | 'magenta': '\033[35m',
172 | 'cyan': '\033[36m',
173 | 'white': '\033[37m',
174 | 'bright_black': '\033[90m', # bright colors
175 | 'bright_red': '\033[91m',
176 | 'bright_green': '\033[92m',
177 | 'bright_yellow': '\033[93m',
178 | 'bright_blue': '\033[94m',
179 | 'bright_magenta': '\033[95m',
180 | 'bright_cyan': '\033[96m',
181 | 'bright_white': '\033[97m',
182 | 'end': '\033[0m', # misc
183 | 'bold': '\033[1m',
184 | 'underline': '\033[4m'}
185 | return ''.join(colors[x] for x in args) + f'{string}' + colors['end']
186 |
187 |
188 | def labels_to_class_weights(labels, nc=80):
189 | # Get class weights (inverse frequency) from training labels
190 | if labels[0] is None: # no labels loaded
191 | return torch.Tensor()
192 |
193 | labels = np.concatenate(labels, 0) # labels.shape = (866643, 5) for COCO
194 | classes = labels[:, 0].astype(np.int) # labels = [class xywh]
195 | weights = np.bincount(classes, minlength=nc) # occurrences per class
196 |
197 | # Prepend gridpoint count (for uCE training)
198 | # gpi = ((320 / 32 * np.array([1, 2, 4])) ** 2 * 3).sum() # gridpoints per image
199 | # weights = np.hstack([gpi * len(labels) - weights.sum() * 9, weights * 9]) ** 0.5 # prepend gridpoints to start
200 |
201 | weights[weights == 0] = 1 # replace empty bins with 1
202 | weights = 1 / weights # number of targets per class
203 | weights /= weights.sum() # normalize
204 | return torch.from_numpy(weights)
205 |
206 |
207 | def labels_to_image_weights(labels, nc=80, class_weights=np.ones(80)):
208 | # Produces image weights based on class_weights and image contents
209 | class_counts = np.array([np.bincount(x[:, 0].astype(np.int), minlength=nc) for x in labels])
210 | image_weights = (class_weights.reshape(1, nc) * class_counts).sum(1)
211 | # index = random.choices(range(n), weights=image_weights, k=1) # weight image sample
212 | return image_weights
213 |
214 |
215 | def coco80_to_coco91_class(): # converts 80-index (val2014) to 91-index (paper)
216 | # https://tech.amikelive.com/node-718/what-object-categories-labels-are-in-coco-dataset/
217 | # a = np.loadtxt('data/coco.names', dtype='str', delimiter='\n')
218 | # b = np.loadtxt('data/coco_paper.names', dtype='str', delimiter='\n')
219 | # x1 = [list(a[i] == b).index(True) + 1 for i in range(80)] # darknet to coco
220 | # x2 = [list(b[i] == a).index(True) if any(b[i] == a) else None for i in range(91)] # coco to darknet
221 | x = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34,
222 | 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63,
223 | 64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90]
224 | return x
225 |
226 |
227 | def xyxy2xywh(x):
228 | # Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right
229 | y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
230 | y[:, 0] = (x[:, 0] + x[:, 2]) / 2 # x center
231 | y[:, 1] = (x[:, 1] + x[:, 3]) / 2 # y center
232 | y[:, 2] = x[:, 2] - x[:, 0] # width
233 | y[:, 3] = x[:, 3] - x[:, 1] # height
234 | return y
235 |
236 |
237 | def xywh2xyxy(x):
238 | # Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
239 | y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
240 | y[:, 0] = x[:, 0] - x[:, 2] / 2 # top left x
241 | y[:, 1] = x[:, 1] - x[:, 3] / 2 # top left y
242 | y[:, 2] = x[:, 0] + x[:, 2] / 2 # bottom right x
243 | y[:, 3] = x[:, 1] + x[:, 3] / 2 # bottom right y
244 | return y
245 |
246 |
247 | def xywhn2xyxy(x, w=640, h=640, padw=0, padh=0):
248 | # Convert nx4 boxes from [x, y, w, h] normalized to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
249 | y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
250 | y[:, 0] = w * (x[:, 0] - x[:, 2] / 2) + padw # top left x
251 | y[:, 1] = h * (x[:, 1] - x[:, 3] / 2) + padh # top left y
252 | y[:, 2] = w * (x[:, 0] + x[:, 2] / 2) + padw # bottom right x
253 | y[:, 3] = h * (x[:, 1] + x[:, 3] / 2) + padh # bottom right y
254 | return y
255 |
256 |
257 | def xyn2xy(x, w=640, h=640, padw=0, padh=0):
258 | # Convert normalized segments into pixel segments, shape (n,2)
259 | y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
260 | y[:, 0] = w * x[:, 0] + padw # top left x
261 | y[:, 1] = h * x[:, 1] + padh # top left y
262 | return y
263 |
264 |
265 | def segment2box(segment, width=640, height=640):
266 | # Convert 1 segment label to 1 box label, applying inside-image constraint, i.e. (xy1, xy2, ...) to (xyxy)
267 | x, y = segment.T # segment xy
268 | inside = (x >= 0) & (y >= 0) & (x <= width) & (y <= height)
269 | x, y, = x[inside], y[inside]
270 | return np.array([x.min(), y.min(), x.max(), y.max()]) if any(x) else np.zeros((1, 4)) # cls, xyxy
271 |
272 |
273 | def segments2boxes(segments):
274 | # Convert segment labels to box labels, i.e. (cls, xy1, xy2, ...) to (cls, xywh)
275 | boxes = []
276 | for s in segments:
277 | x, y = s.T # segment xy
278 | boxes.append([x.min(), y.min(), x.max(), y.max()]) # cls, xyxy
279 | return xyxy2xywh(np.array(boxes)) # cls, xywh
280 |
281 |
282 | def resample_segments(segments, n=1000):
283 | # Up-sample an (n,2) segment
284 | for i, s in enumerate(segments):
285 | x = np.linspace(0, len(s) - 1, n)
286 | xp = np.arange(len(s))
287 | segments[i] = np.concatenate([np.interp(x, xp, s[:, i]) for i in range(2)]).reshape(2, -1).T # segment xy
288 | return segments
289 |
290 |
291 | def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None):
292 | # Rescale coords (xyxy) from img1_shape to img0_shape
293 | if ratio_pad is None: # calculate from img0_shape
294 | gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain = old / new
295 | pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding
296 | else:
297 | gain = ratio_pad[0][0]
298 | pad = ratio_pad[1]
299 |
300 | coords[:, [0, 2]] -= pad[0] # x padding
301 | coords[:, [1, 3]] -= pad[1] # y padding
302 | coords[:, :4] /= gain
303 | clip_coords(coords, img0_shape)
304 | return coords
305 |
306 |
307 | def clip_coords(boxes, img_shape):
308 | # Clip bounding xyxy bounding boxes to image shape (height, width)
309 | boxes[:, 0].clamp_(0, img_shape[1]) # x1
310 | boxes[:, 1].clamp_(0, img_shape[0]) # y1
311 | boxes[:, 2].clamp_(0, img_shape[1]) # x2
312 | boxes[:, 3].clamp_(0, img_shape[0]) # y2
313 |
314 |
315 | def bbox_iou(box1, box2, x1y1x2y2=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-9):
316 | # Returns the IoU of box1 to box2. box1 is 4, box2 is nx4
317 | box2 = box2.T
318 |
319 | # Get the coordinates of bounding boxes
320 | if x1y1x2y2: # x1, y1, x2, y2 = box1
321 | b1_x1, b1_y1, b1_x2, b1_y2 = box1[0], box1[1], box1[2], box1[3]
322 | b2_x1, b2_y1, b2_x2, b2_y2 = box2[0], box2[1], box2[2], box2[3]
323 | else: # transform from xywh to xyxy
324 | b1_x1, b1_x2 = box1[0] - box1[2] / 2, box1[0] + box1[2] / 2
325 | b1_y1, b1_y2 = box1[1] - box1[3] / 2, box1[1] + box1[3] / 2
326 | b2_x1, b2_x2 = box2[0] - box2[2] / 2, box2[0] + box2[2] / 2
327 | b2_y1, b2_y2 = box2[1] - box2[3] / 2, box2[1] + box2[3] / 2
328 |
329 | # Intersection area
330 | inter = (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * \
331 | (torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1)).clamp(0)
332 |
333 | # Union Area
334 | w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps
335 | w2, h2 = b2_x2 - b2_x1, b2_y2 - b2_y1 + eps
336 | union = w1 * h1 + w2 * h2 - inter + eps
337 |
338 | iou = inter / union
339 | if GIoU or DIoU or CIoU:
340 | cw = torch.max(b1_x2, b2_x2) - torch.min(b1_x1, b2_x1) # convex (smallest enclosing box) width
341 | ch = torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1) # convex height
342 | if CIoU or DIoU: # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1
343 | c2 = cw ** 2 + ch ** 2 + eps # convex diagonal squared
344 | rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 +
345 | (b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4 # center distance squared
346 | if DIoU:
347 | return iou - rho2 / c2 # DIoU
348 | elif CIoU: # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47
349 | v = (4 / math.pi ** 2) * torch.pow(torch.atan(w2 / h2) - torch.atan(w1 / h1), 2)
350 | with torch.no_grad():
351 | alpha = v / ((1 + eps) - iou + v)
352 | return iou - (rho2 / c2 + v * alpha) # CIoU
353 | else: # GIoU https://arxiv.org/pdf/1902.09630.pdf
354 | c_area = cw * ch + eps # convex area
355 | return iou - (c_area - union) / c_area # GIoU
356 | else:
357 | return iou # IoU
358 |
359 |
360 | def box_iou(box1, box2):
361 | # https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py
362 | """
363 | Return intersection-over-union (Jaccard index) of boxes.
364 | Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
365 | Arguments:
366 | box1 (Tensor[N, 4])
367 | box2 (Tensor[M, 4])
368 | Returns:
369 | iou (Tensor[N, M]): the NxM matrix containing the pairwise
370 | IoU values for every element in boxes1 and boxes2
371 | """
372 |
373 | def box_area(box):
374 | # box = 4xn
375 | return (box[2] - box[0]) * (box[3] - box[1])
376 |
377 | area1 = box_area(box1.T)
378 | area2 = box_area(box2.T)
379 |
380 | # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
381 | inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
382 | return inter / (area1[:, None] + area2 - inter) # iou = inter / (area1 + area2 - inter)
383 |
384 |
385 | def wh_iou(wh1, wh2):
386 | # Returns the nxm IoU matrix. wh1 is nx2, wh2 is mx2
387 | wh1 = wh1[:, None] # [N,1,2]
388 | wh2 = wh2[None] # [1,M,2]
389 | inter = torch.min(wh1, wh2).prod(2) # [N,M]
390 | return inter / (wh1.prod(2) + wh2.prod(2) - inter) # iou = inter / (area1 + area2 - inter)
391 |
392 |
393 | def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False,
394 | labels=()):
395 | """Runs Non-Maximum Suppression (NMS) on inference results
396 |
397 | Returns:
398 | list of detections, on (n,6) tensor per image [xyxy, conf, cls]
399 | """
400 |
401 | nc = prediction.shape[2] - 5 # number of classes
402 | xc = prediction[..., 4] > conf_thres # candidates
403 |
404 | # Settings
405 | min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height
406 | max_det = 300 # maximum number of detections per image
407 | max_nms = 30000 # maximum number of boxes into torchvision.ops.nms()
408 | time_limit = 10.0 # seconds to quit after
409 | redundant = True # require redundant detections
410 | multi_label &= nc > 1 # multiple labels per box (adds 0.5ms/img)
411 | merge = False # use merge-NMS
412 |
413 | t = time.time()
414 | output = [torch.zeros((0, 6), device=prediction.device)] * prediction.shape[0]
415 | for xi, x in enumerate(prediction): # image index, image inference
416 | # Apply constraints
417 | # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height
418 | x = x[xc[xi]] # confidence
419 |
420 | # Cat apriori labels if autolabelling
421 | if labels and len(labels[xi]):
422 | l = labels[xi]
423 | v = torch.zeros((len(l), nc + 5), device=x.device)
424 | v[:, :4] = l[:, 1:5] # box
425 | v[:, 4] = 1.0 # conf
426 | v[range(len(l)), l[:, 0].long() + 5] = 1.0 # cls
427 | x = torch.cat((x, v), 0)
428 |
429 | # If none remain process next image
430 | if not x.shape[0]:
431 | continue
432 |
433 | # Compute conf
434 | x[:, 5:] *= x[:, 4:5] # conf = obj_conf * cls_conf
435 |
436 | # Box (center x, center y, width, height) to (x1, y1, x2, y2)
437 | box = xywh2xyxy(x[:, :4])
438 |
439 | # Detections matrix nx6 (xyxy, conf, cls)
440 | if multi_label:
441 | i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T
442 | x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
443 | else: # best class only
444 | conf, j = x[:, 5:].max(1, keepdim=True)
445 | x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]
446 |
447 | # Filter by class
448 | if classes is not None:
449 | x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
450 |
451 | # Apply finite constraint
452 | # if not torch.isfinite(x).all():
453 | # x = x[torch.isfinite(x).all(1)]
454 |
455 | # Check shape
456 | n = x.shape[0] # number of boxes
457 | if not n: # no boxes
458 | continue
459 | elif n > max_nms: # excess boxes
460 | x = x[x[:, 4].argsort(descending=True)[:max_nms]] # sort by confidence
461 |
462 | # Batched NMS
463 | c = x[:, 5:6] * (0 if agnostic else max_wh) # classes
464 | boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores
465 | i = torchvision.ops.nms(boxes, scores, iou_thres) # NMS
466 | if i.shape[0] > max_det: # limit detections
467 | i = i[:max_det]
468 | if merge and (1 < n < 3E3): # Merge NMS (boxes merged using weighted mean)
469 | # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
470 | iou = box_iou(boxes[i], boxes) > iou_thres # iou matrix
471 | weights = iou * scores[None] # box weights
472 | x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True) # merged boxes
473 | if redundant:
474 | i = i[iou.sum(1) > 1] # require redundancy
475 |
476 | output[xi] = x[i]
477 | if (time.time() - t) > time_limit:
478 | print(f'WARNING: NMS time limit {time_limit}s exceeded')
479 | break # time limit exceeded
480 |
481 | return output
482 |
483 |
484 | def strip_optimizer(f='best.pt', s=''): # from utils.general import *; strip_optimizer()
485 | # Strip optimizer from 'f' to finalize training, optionally save as 's'
486 | x = torch.load(f, map_location=torch.device('cpu'))
487 | if x.get('ema'):
488 | x['model'] = x['ema'] # replace model with ema
489 | for k in 'optimizer', 'training_results', 'wandb_id', 'ema', 'updates': # keys
490 | x[k] = None
491 | x['epoch'] = -1
492 | x['model'].half() # to FP16
493 | for p in x['model'].parameters():
494 | p.requires_grad = False
495 | torch.save(x, s or f)
496 | mb = os.path.getsize(s or f) / 1E6 # filesize
497 | print(f"Optimizer stripped from {f},{(' saved as %s,' % s) if s else ''} {mb:.1f}MB")
498 |
499 |
500 | def print_mutation(hyp, results, yaml_file='hyp_evolved.yaml', bucket=''):
501 | # Print mutation results to evolve.txt (for use with train.py --evolve)
502 | a = '%10s' * len(hyp) % tuple(hyp.keys()) # hyperparam keys
503 | b = '%10.3g' * len(hyp) % tuple(hyp.values()) # hyperparam values
504 | c = '%10.4g' * len(results) % results # results (P, R, mAP@0.5, mAP@0.5:0.95, val_losses x 3)
505 | print('\n%s\n%s\nEvolved fitness: %s\n' % (a, b, c))
506 |
507 | if bucket:
508 | url = 'gs://%s/evolve.txt' % bucket
509 | if gsutil_getsize(url) > (os.path.getsize('evolve.txt') if os.path.exists('evolve.txt') else 0):
510 | os.system('gsutil cp %s .' % url) # download evolve.txt if larger than local
511 |
512 | with open('evolve.txt', 'a') as f: # append result
513 | f.write(c + b + '\n')
514 | x = np.unique(np.loadtxt('evolve.txt', ndmin=2), axis=0) # load unique rows
515 | x = x[np.argsort(-fitness(x))] # sort
516 | np.savetxt('evolve.txt', x, '%10.3g') # save sort by fitness
517 |
518 | # Save yaml
519 | for i, k in enumerate(hyp.keys()):
520 | hyp[k] = float(x[0, i + 7])
521 | with open(yaml_file, 'w') as f:
522 | results = tuple(x[0, :7])
523 | c = '%10.4g' * len(results) % results # results (P, R, mAP@0.5, mAP@0.5:0.95, val_losses x 3)
524 | f.write('# Hyperparameter Evolution Results\n# Generations: %g\n# Metrics: ' % len(x) + c + '\n\n')
525 | yaml.dump(hyp, f, sort_keys=False)
526 |
527 | if bucket:
528 | os.system('gsutil cp evolve.txt %s gs://%s' % (yaml_file, bucket)) # upload
529 |
530 |
531 | def apply_classifier(x, model, img, im0):
532 | # applies a second stage classifier to yolo outputs
533 | im0 = [im0] if isinstance(im0, np.ndarray) else im0
534 | for i, d in enumerate(x): # per image
535 | if d is not None and len(d):
536 | d = d.clone()
537 |
538 | # Reshape and pad cutouts
539 | b = xyxy2xywh(d[:, :4]) # boxes
540 | b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1) # rectangle to square
541 | b[:, 2:] = b[:, 2:] * 1.3 + 30 # pad
542 | d[:, :4] = xywh2xyxy(b).long()
543 |
544 | # Rescale boxes from img_size to im0 size
545 | scale_coords(img.shape[2:], d[:, :4], im0[i].shape)
546 |
547 | # Classes
548 | pred_cls1 = d[:, 5].long()
549 | ims = []
550 | for j, a in enumerate(d): # per item
551 | cutout = im0[i][int(a[1]):int(a[3]), int(a[0]):int(a[2])]
552 | im = cv2.resize(cutout, (224, 224)) # BGR
553 | # cv2.imwrite('test%i.jpg' % j, cutout)
554 |
555 | im = im[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
556 | im = np.ascontiguousarray(im, dtype=np.float32) # uint8 to float32
557 | im /= 255.0 # 0 - 255 to 0.0 - 1.0
558 | ims.append(im)
559 |
560 | pred_cls2 = model(torch.Tensor(ims).to(d.device)).argmax(1) # classifier prediction
561 | x[i] = x[i][pred_cls1 == pred_cls2] # retain matching class detections
562 |
563 | return x
564 |
565 |
566 | def increment_path(path, exist_ok=True, sep=''):
567 | # Increment path, i.e. runs/exp --> runs/exp{sep}0, runs/exp{sep}1 etc.
568 | path = Path(path) # os-agnostic
569 | if (path.exists() and exist_ok) or (not path.exists()):
570 | return str(path)
571 | else:
572 | dirs = glob.glob(f"{path}{sep}*") # similar paths
573 | matches = [re.search(rf"%s{sep}(\d+)" % path.stem, d) for d in dirs]
574 | i = [int(m.groups()[0]) for m in matches if m] # indices
575 | n = max(i) + 1 if i else 2 # increment number
576 | return f"{path}{sep}{n}" # update path
577 |
--------------------------------------------------------------------------------
/utils/google_utils.py:
--------------------------------------------------------------------------------
1 | # Google utils: https://cloud.google.com/storage/docs/reference/libraries
2 |
3 | import os
4 | import platform
5 | import subprocess
6 | import time
7 | from pathlib import Path
8 |
9 | import requests
10 | import torch
11 |
12 |
13 | def gsutil_getsize(url=''):
14 | # gs://bucket/file size https://cloud.google.com/storage/docs/gsutil/commands/du
15 | s = subprocess.check_output(f'gsutil du {url}', shell=True).decode('utf-8')
16 | return eval(s.split(' ')[0]) if len(s) else 0 # bytes
17 |
18 |
19 | def attempt_download(file, repo='ultralytics/yolov5'):
20 | # Attempt file download if does not exist
21 | file = Path(str(file).strip().replace("'", '').lower())
22 | if not file.exists():
23 | try:
24 | proxy={"http":None,"https":None}
25 | #response = requests.get(r'D:\实验室\程序\YOLOV5\latest_access_token=86160778ebfa6f610de04e6f8609b4ffabae3d88',proxies=proxy).json() # github api
26 | response = requests.get(f'https://api.github.com/repos/{repo}/releases/latest?access_token=86160778ebfa6f610de04e6f8609b4ffabae3d88',proxies=proxy).json() # github api
27 | assets = [x['name'] for x in response['assets']] # release assets, i.e. ['yolov5s.pt', 'yolov5m.pt', ...]
28 | tag = response['tag_name'] # i.e. 'v1.0'
29 | except: # fallback plan
30 | assets = ['yolov5.pt', 'yolov5.pt', 'yolov5l.pt', 'yolov5x.pt']
31 | tag = subprocess.check_output('git tag', shell=True).decode('utf-8').split('\n')[-2]
32 |
33 | name = file.name
34 | if name in assets:
35 | msg = f'{file} missing, try downloading from https://github.com/{repo}/releases/'
36 | redundant = False # second download option
37 | try: # GitHub
38 | url = f'https://github.com/{repo}/releases/download/{tag}/{name}'
39 | print(f'Downloading {url} to {file}...')
40 | torch.hub.download_url_to_file(url, file)
41 | assert file.exists() and file.stat().st_size > 1E6 # check
42 | except Exception as e: # GCP
43 | print(f'Download error: {e}')
44 | assert redundant, 'No secondary mirror'
45 | url = f'https://storage.googleapis.com/{repo}/ckpt/{name}'
46 | print(f'Downloading {url} to {file}...')
47 | os.system(f'curl -L {url} -o {file}') # torch.hub.download_url_to_file(url, weights)
48 | finally:
49 | if not file.exists() or file.stat().st_size < 1E6: # check
50 | file.unlink(missing_ok=True) # remove partial downloads
51 | print(f'ERROR: Download failure: {msg}')
52 | print('')
53 | return
54 |
55 |
56 | def gdrive_download(id='16TiPfZj7htmTyhntwcZyEEAejOUxuT6m', file='tmp.zip'):
57 | # Downloads a file from Google Drive. from yolov5.utils.google_utils import *; gdrive_download()
58 | t = time.time()
59 | file = Path(file)
60 | cookie = Path('cookie') # gdrive cookie
61 | print(f'Downloading https://drive.google.com/uc?export=download&id={id} as {file}... ', end='')
62 | file.unlink(missing_ok=True) # remove existing file
63 | cookie.unlink(missing_ok=True) # remove existing cookie
64 |
65 | # Attempt file download
66 | out = "NUL" if platform.system() == "Windows" else "/dev/null"
67 | os.system(f'curl -c ./cookie -s -L "drive.google.com/uc?export=download&id={id}" > {out}')
68 | if os.path.exists('cookie'): # large file
69 | s = f'curl -Lb ./cookie "drive.google.com/uc?export=download&confirm={get_token()}&id={id}" -o {file}'
70 | else: # small file
71 | s = f'curl -s -L -o {file} "drive.google.com/uc?export=download&id={id}"'
72 | r = os.system(s) # execute, capture return
73 | cookie.unlink(missing_ok=True) # remove existing cookie
74 |
75 | # Error check
76 | if r != 0:
77 | file.unlink(missing_ok=True) # remove partial
78 | print('Download error ') # raise Exception('Download error')
79 | return r
80 |
81 | # Unzip if archive
82 | if file.suffix == '.zip':
83 | print('unzipping... ', end='')
84 | os.system(f'unzip -q {file}') # unzip
85 | file.unlink() # remove zip to free space
86 |
87 | print(f'Done ({time.time() - t:.1f}s)')
88 | return r
89 |
90 |
91 | def get_token(cookie="./cookie"):
92 | with open(cookie) as f:
93 | for line in f:
94 | if "download" in line:
95 | return line.split()[-1]
96 | return ""
97 |
98 | # def upload_blob(bucket_name, source_file_name, destination_blob_name):
99 | # # Uploads a file to a bucket
100 | # # https://cloud.google.com/storage/docs/uploading-objects#storage-upload-object-python
101 | #
102 | # storage_client = storage.Client()
103 | # bucket = storage_client.get_bucket(bucket_name)
104 | # blob = bucket.blob(destination_blob_name)
105 | #
106 | # blob.upload_from_filename(source_file_name)
107 | #
108 | # print('File {} uploaded to {}.'.format(
109 | # source_file_name,
110 | # destination_blob_name))
111 | #
112 | #
113 | # def download_blob(bucket_name, source_blob_name, destination_file_name):
114 | # # Uploads a blob from a bucket
115 | # storage_client = storage.Client()
116 | # bucket = storage_client.get_bucket(bucket_name)
117 | # blob = bucket.blob(source_blob_name)
118 | #
119 | # blob.download_to_filename(destination_file_name)
120 | #
121 | # print('Blob {} downloaded to {}.'.format(
122 | # source_blob_name,
123 | # destination_file_name))
124 |
--------------------------------------------------------------------------------
/utils/loss.py:
--------------------------------------------------------------------------------
1 | # Loss functions
2 |
3 | import torch
4 | import torch.nn as nn
5 |
6 | from utils.general import bbox_iou
7 | from utils.torch_utils import is_parallel
8 |
9 | class AutomaticWeightedLoss(nn.Module):
10 | def __init__(self, num=2):
11 | super(AutomaticWeightedLoss, self).__init__()
12 | self.params = torch.nn.Parameter(torch.ones(num, requires_grad=True).cuda())
13 |
14 | def forward(self, *x):
15 | loss_sum = 0
16 | for i, loss in enumerate(x):
17 | loss_sum += 0.5 / (self.params[i] ** 2) * loss + torch.log(1 + self.params[i] ** 2)
18 | return loss_sum
19 |
20 | def smooth_BCE(eps=0.1): # https://github.com/ultralytics/yolov3/issues/238#issuecomment-598028441
21 | # return positive, negative label smoothing BCE targets
22 | return 1.0 - 0.5 * eps, 0.5 * eps
23 |
24 |
25 | class BCEBlurWithLogitsLoss(nn.Module):
26 | # BCEwithLogitLoss() with reduced missing label effects.
27 | def __init__(self, alpha=0.05):
28 | super(BCEBlurWithLogitsLoss, self).__init__()
29 | self.loss_fcn = nn.BCEWithLogitsLoss(reduction='none') # must be nn.BCEWithLogitsLoss()
30 | self.alpha = alpha
31 |
32 | def forward(self, pred, true):
33 | loss = self.loss_fcn(pred, true)
34 | pred = torch.sigmoid(pred) # prob from logits
35 | dx = pred - true # reduce only missing label effects
36 | # dx = (pred - true).abs() # reduce missing label and false label effects
37 | alpha_factor = 1 - torch.exp((dx - 1) / (self.alpha + 1e-4))
38 | loss *= alpha_factor
39 | return loss.mean()
40 |
41 |
42 | class FocalLoss(nn.Module):
43 | # Wraps focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5)
44 | def __init__(self, loss_fcn, gamma=1.5, alpha=0.25):
45 | super(FocalLoss, self).__init__()
46 | self.loss_fcn = loss_fcn # must be nn.BCEWithLogitsLoss()
47 | self.gamma = gamma
48 | self.alpha = alpha
49 | self.reduction = loss_fcn.reduction
50 | self.loss_fcn.reduction = 'none' # required to apply FL to each element
51 |
52 | def forward(self, pred, true):
53 | loss = self.loss_fcn(pred, true)
54 | # p_t = torch.exp(-loss)
55 | # loss *= self.alpha * (1.000001 - p_t) ** self.gamma # non-zero power for gradient stability
56 |
57 | # TF implementation https://github.com/tensorflow/addons/blob/v0.7.1/tensorflow_addons/losses/focal_loss.py
58 | pred_prob = torch.sigmoid(pred) # prob from logits
59 | p_t = true * pred_prob + (1 - true) * (1 - pred_prob)
60 | alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha)
61 | modulating_factor = (1.0 - p_t) ** self.gamma
62 | loss *= alpha_factor * modulating_factor
63 |
64 | if self.reduction == 'mean':
65 | return loss.mean()
66 | elif self.reduction == 'sum':
67 | return loss.sum()
68 | else: # 'none'
69 | return loss
70 |
71 |
72 | class QFocalLoss(nn.Module):
73 | # Wraps Quality focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5)
74 | def __init__(self, loss_fcn, gamma=1.5, alpha=0.25):
75 | super(QFocalLoss, self).__init__()
76 | self.loss_fcn = loss_fcn # must be nn.BCEWithLogitsLoss()
77 | self.gamma = gamma
78 | self.alpha = alpha
79 | self.reduction = loss_fcn.reduction
80 | self.loss_fcn.reduction = 'none' # required to apply FL to each element
81 |
82 | def forward(self, pred, true):
83 | loss = self.loss_fcn(pred, true)
84 |
85 | pred_prob = torch.sigmoid(pred) # prob from logits
86 | alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha)
87 | modulating_factor = torch.abs(true - pred_prob) ** self.gamma
88 | loss *= alpha_factor * modulating_factor
89 |
90 | if self.reduction == 'mean':
91 | return loss.mean()
92 | elif self.reduction == 'sum':
93 | return loss.sum()
94 | else: # 'none'
95 | return loss
96 |
97 |
98 | class ComputeLoss:
99 | # Compute losses
100 | def __init__(self, model, autobalance=False):
101 | super(ComputeLoss, self).__init__()
102 |
103 | self.multiAdapt = AutomaticWeightedLoss(4)
104 | # self.multiAdapt = AutomaticWeightedLoss(2) # ForAuto
105 |
106 | device = next(model.parameters()).device # get model device
107 | h = model.hyp # hyperparameters
108 |
109 | # Define criteria
110 | BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['cls_pw']], device=device))
111 | BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['obj_pw']], device=device))
112 |
113 | # Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3
114 | self.cp, self.cn = smooth_BCE(eps=0.0)
115 |
116 | # Focal loss
117 | g = h['fl_gamma'] # focal loss gamma
118 | if g > 0:
119 | BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)
120 |
121 | det = model.module.model[-1] if is_parallel(model) else model.model[-1] # Detect() module
122 | # self.balance = {2: [3.56, 1.0], 3: [3.67, 1.0, 0.43], 4: [3.78, 1.0, 0.39, 0.22], 5: [3.88, 1.0, 0.37, 0.17, 0.10]}[det.nl]
123 | # # self.balance = [1.0] * det.nl
124 | # self.ssi = (det.stride == 16).nonzero(as_tuple=False).item() # stride 16 index
125 | self.balance = {3: [4.0, 1.0, 0.4]}.get(det.nl, [4.0, 1.0, 0.25, 0.06, .02]) # P3-P7
126 | self.ssi = list(det.stride).index(16) if autobalance else 0 # stride 16 index
127 | self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = BCEcls, BCEobj, model.gr, h, autobalance
128 | self.dgMSE = nn.MSELoss(reduce=True, size_average=True)
129 | for k in 'na', 'nc', 'nl', 'anchors':
130 | setattr(self, k, getattr(det, k))
131 |
132 | def __call__(self, p, dgimgs, targets): # predictions, targets, model
133 | device = targets.device
134 | lcls, lbox, lobj, ldg= torch.zeros(1, device=device), torch.zeros(1, device=device), torch.zeros(1, device=device), torch.zeros(1, device=device)
135 | tcls, tbox, indices, anchors = self.build_targets(p[0], targets) # targets
136 |
137 | # Losses
138 | for i, pi in enumerate(p[0]): # layer index, layer predictions
139 | b, a, gj, gi = indices[i] # image, anchor, gridy, gridx
140 | tobj = torch.zeros_like(pi[..., 0], device=device) # target obj
141 |
142 | n = b.shape[0] # number of targets
143 | if n:
144 | ps = pi[b, a, gj, gi] # prediction subset corresponding to targets
145 |
146 | # Regression
147 | pxy = ps[:, :2].sigmoid() * 2. - 0.5
148 | pwh = (ps[:, 2:4].sigmoid() * 2) ** 2 * anchors[i]
149 | pbox = torch.cat((pxy, pwh), 1) # predicted box
150 | iou = bbox_iou(pbox.T, tbox[i], x1y1x2y2=False, CIoU=True) # iou(prediction, target)
151 | lbox += (1.0 - iou).mean() # iou loss
152 |
153 | # Objectness
154 | tobj[b, a, gj, gi] = (1.0 - self.gr) + self.gr * iou.detach().clamp(0).type(tobj.dtype) # iou ratio
155 |
156 | # Classification
157 | if self.nc > 1: # cls loss (only if multiple classes)
158 | t = torch.full_like(ps[:, 5:], self.cn, device=device) # targets
159 | t[range(n), tcls[i]] = self.cp
160 | lcls += self.BCEcls(ps[:, 5:], t) # BCE
161 |
162 | # Append targets to text file
163 | # with open('targets.txt', 'a') as file:
164 | # [file.write('%11.5g ' * 4 % tuple(x) + '\n') for x in torch.cat((txy[i], twh[i]), 1)]
165 |
166 | obji = self.BCEobj(pi[..., 4], tobj)
167 | lobj += obji * self.balance[i] # obj loss
168 | if self.autobalance:
169 | self.balance[i] = self.balance[i] * 0.9999 + 0.0001 / obji.detach().item()
170 |
171 | # "loss for patch"
172 | # if p[1]!=None:
173 | # ldgFactor = [0 if x not in targets[:, 0] else 1 for x in range(len(dgimgs))]
174 | # dgimgs = torch.stack([dgimgs[x] if ldgFactor[x] else p[1][x] for x in range(len(dgimgs))]).detach()
175 | # ldg += self.dgMSE(p[1].float(), dgimgs.float())
176 |
177 | if p[1]!=None:
178 | ldg += self.dgMSE(p[1].float(), dgimgs.float())
179 |
180 | if self.autobalance:
181 | self.balance = [x / self.balance[self.ssi] for x in self.balance]
182 | lbox *= self.hyp['box']
183 | lobj *= self.hyp['obj']
184 | lcls *= self.hyp['cls']
185 | # ldg *= 1 # ForAuto
186 | ldg *= 0.1
187 |
188 | bs = tobj.shape[0]
189 |
190 | loss = self.multiAdapt(lbox, lobj, lcls, ldg)
191 | # loss = self.multiAdapt((lbox+ lobj+ lcls), ldg) # ForAuto
192 | # loss = lbox + lobj + lcls + ldg
193 | a=torch.cat((lbox, lobj, lcls, ldg, loss))
194 | b=a.detach()
195 | return loss * bs, torch.cat((lbox, lobj, lcls, ldg, loss)).detach(), self.multiAdapt.params.detach()
196 |
197 | def build_targets(self, p, targets):
198 | # Build targets for compute_loss(), input targets(image,class,x,y,w,h)
199 | na, nt = self.na, targets.shape[0] # number of anchors, targets
200 | tcls, tbox, indices, anch = [], [], [], []
201 | gain = torch.ones(7, device=targets.device) # normalized to gridspace gain
202 | ai = torch.arange(na, device=targets.device).float().view(na, 1).repeat(1, nt) # same as .repeat_interleave(nt)
203 | targets = torch.cat((targets.repeat(na, 1, 1), ai[:, :, None]), 2) # append anchor indices
204 |
205 | g = 0.5 # bias
206 | off = torch.tensor([[0, 0],
207 | [1, 0], [0, 1], [-1, 0], [0, -1], # j,k,l,m
208 | # [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm
209 | ], device=targets.device).float() * g # offsets
210 |
211 | for i in range(self.nl):
212 | anchors = self.anchors[i]
213 | gain[2:6] = torch.tensor(p[i].shape)[[3, 2, 3, 2]] # xyxy gain
214 |
215 | # Match targets to anchors
216 | t = targets * gain
217 | if nt:
218 | # Matches
219 | r = t[:, :, 4:6] / anchors[:, None] # wh ratio
220 | j = torch.max(r, 1. / r).max(2)[0] < self.hyp['anchor_t'] # compare
221 | # j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t'] # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2))
222 | t = t[j] # filter
223 |
224 | # Offsets
225 | gxy = t[:, 2:4] # grid xy
226 | gxi = gain[[2, 3]] - gxy # inverse
227 | j, k = ((gxy % 1. < g) & (gxy > 1.)).T
228 | l, m = ((gxi % 1. < g) & (gxi > 1.)).T
229 | j = torch.stack((torch.ones_like(j), j, k, l, m))
230 | t = t.repeat((5, 1, 1))[j]
231 | offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j]
232 | else:
233 | t = targets[0]
234 | offsets = 0
235 |
236 | # Define
237 | b, c = t[:, :2].long().T # image, class
238 | gxy = t[:, 2:4] # grid xy
239 | gwh = t[:, 4:6] # grid wh
240 | gij = (gxy - offsets).long()
241 | gi, gj = gij.T # grid xy indices
242 |
243 | # Append
244 | a = t[:, 6].long() # anchor indices
245 | indices.append((b, a, gj.clamp_(0, gain[3] - 1), gi.clamp_(0, gain[2] - 1))) # image, anchor, grid indices
246 | tbox.append(torch.cat((gxy - gij, gwh), 1)) # box
247 | anch.append(anchors[a]) # anchors
248 | tcls.append(c) # class
249 |
250 | return tcls, tbox, indices, anch
251 |
--------------------------------------------------------------------------------
/utils/metrics.py:
--------------------------------------------------------------------------------
1 | # Model validation metrics
2 |
3 | from pathlib import Path
4 |
5 | import matplotlib.pyplot as plt
6 | import numpy as np
7 | import torch
8 |
9 | from . import general
10 |
11 |
12 | def fitness(x):
13 | # Model fitness as a weighted combination of metrics
14 | w = [0.0, 0.0, 0.1, 0.9] # weights for [P, R, mAP@0.5, mAP@0.5:0.95]
15 | return (x[:, :4] * w).sum(1)
16 |
17 |
18 | def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='precision-recall_curve.png', names=[]):
19 | """ Compute the average precision, given the recall and precision curves.
20 | Source: https://github.com/rafaelpadilla/Object-Detection-Metrics.
21 | # Arguments
22 | tp: True positives (nparray, nx1 or nx10).
23 | conf: Objectness value from 0-1 (nparray).
24 | pred_cls: Predicted object classes (nparray).
25 | target_cls: True object classes (nparray).
26 | plot: Plot precision-recall curve at mAP@0.5
27 | save_dir: Plot save directory
28 | # Returns
29 | The average precision as computed in py-faster-rcnn.
30 | """
31 |
32 | # Sort by objectness
33 | i = np.argsort(-conf)
34 | tp, conf, pred_cls = tp[i], conf[i], pred_cls[i]
35 |
36 | # Find unique classes
37 | unique_classes = np.unique(target_cls)
38 |
39 | # Create Precision-Recall curve and compute AP for each class
40 | px, py = np.linspace(0, 1, 1000), [] # for plotting
41 | pr_score = 0.1 # score to evaluate P and R https://github.com/ultralytics/yolov3/issues/898
42 | s = [unique_classes.shape[0], tp.shape[1]] # number class, number iou thresholds (i.e. 10 for mAP0.5...0.95)
43 | ap, p, r = np.zeros(s), np.zeros(s), np.zeros(s)
44 | for ci, c in enumerate(unique_classes):
45 | i = pred_cls == c
46 | n_l = (target_cls == c).sum() # number of labels
47 | n_p = i.sum() # number of predictions
48 |
49 | if n_p == 0 or n_l == 0:
50 | continue
51 | else:
52 | # Accumulate FPs and TPs
53 | fpc = (1 - tp[i]).cumsum(0)
54 | tpc = tp[i].cumsum(0)
55 |
56 | # Recall
57 | recall = tpc / (n_l + 1e-16) # recall curve
58 | r[ci] = np.interp(-pr_score, -conf[i], recall[:, 0]) # r at pr_score, negative x, xp because xp decreases
59 |
60 | # Precision
61 | precision = tpc / (tpc + fpc) # precision curve
62 | p[ci] = np.interp(-pr_score, -conf[i], precision[:, 0]) # p at pr_score
63 |
64 | # AP from recall-precision curve
65 | for j in range(tp.shape[1]):
66 | ap[ci, j], mpre, mrec = compute_ap(recall[:, j], precision[:, j])
67 | if plot and (j == 0):
68 | py.append(np.interp(px, mrec, mpre)) # precision at mAP@0.5
69 |
70 | # Compute F1 score (harmonic mean of precision and recall)
71 | f1 = 2 * p * r / (p + r + 1e-16)
72 |
73 | if plot:
74 | plot_pr_curve(px, py, ap, save_dir, names)
75 |
76 | return p, r, ap, f1, unique_classes.astype('int32')
77 |
78 |
79 | def compute_ap(recall, precision):
80 | """ Compute the average precision, given the recall and precision curves
81 | # Arguments
82 | recall: The recall curve (list)
83 | precision: The precision curve (list)
84 | # Returns
85 | Average precision, precision curve, recall curve
86 | """
87 |
88 | # Append sentinel values to beginning and end
89 | mrec = np.concatenate(([0.], recall, [recall[-1] + 0.01]))
90 | mpre = np.concatenate(([1.], precision, [0.]))
91 |
92 | # Compute the precision envelope
93 | mpre = np.flip(np.maximum.accumulate(np.flip(mpre)))
94 |
95 | # Integrate area under curve
96 | method = 'interp' # methods: 'continuous', 'interp'
97 | if method == 'interp':
98 | x = np.linspace(0, 1, 101) # 101-point interp (COCO)
99 | ap = np.trapz(np.interp(x, mrec, mpre), x) # integrate
100 | else: # 'continuous'
101 | i = np.where(mrec[1:] != mrec[:-1])[0] # points where x axis (recall) changes
102 | ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1]) # area under curve
103 |
104 | return ap, mpre, mrec
105 |
106 |
107 | class ConfusionMatrix:
108 | # Updated version of https://github.com/kaanakan/object_detection_confusion_matrix
109 | def __init__(self, nc, conf=0.25, iou_thres=0.45):
110 | self.matrix = np.zeros((nc + 1, nc + 1))
111 | self.nc = nc # number of classes
112 | self.conf = conf
113 | self.iou_thres = iou_thres
114 |
115 | def process_batch(self, detections, labels):
116 | """
117 | Return intersection-over-union (Jaccard index) of boxes.
118 | Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
119 | Arguments:
120 | detections (Array[N, 6]), x1, y1, x2, y2, conf, class
121 | labels (Array[M, 5]), class, x1, y1, x2, y2
122 | Returns:
123 | None, updates confusion matrix accordingly
124 | """
125 | detections = detections[detections[:, 4] > self.conf]
126 | gt_classes = labels[:, 0].int()
127 | detection_classes = detections[:, 5].int()
128 | iou = general.box_iou(labels[:, 1:], detections[:, :4])
129 |
130 | x = torch.where(iou > self.iou_thres)
131 | if x[0].shape[0]:
132 | matches = torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), 1).cpu().numpy()
133 | if x[0].shape[0] > 1:
134 | matches = matches[matches[:, 2].argsort()[::-1]]
135 | matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
136 | matches = matches[matches[:, 2].argsort()[::-1]]
137 | matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
138 | else:
139 | matches = np.zeros((0, 3))
140 |
141 | n = matches.shape[0] > 0
142 | m0, m1, _ = matches.transpose().astype(np.int16)
143 | for i, gc in enumerate(gt_classes):
144 | j = m0 == i
145 | if n and sum(j) == 1:
146 | self.matrix[gc, detection_classes[m1[j]]] += 1 # correct
147 | else:
148 | self.matrix[self.nc, gc] += 1 # background FP
149 |
150 | if n:
151 | for i, dc in enumerate(detection_classes):
152 | if not any(m1 == i):
153 | self.matrix[dc, self.nc] += 1 # background FN
154 |
155 | def matrix(self):
156 | return self.matrix
157 |
158 | def plot(self, save_dir='', names=()):
159 | try:
160 | import seaborn as sn
161 |
162 | array = self.matrix / (self.matrix.sum(0).reshape(1, self.nc + 1) + 1E-6) # normalize
163 | array[array < 0.005] = np.nan # don't annotate (would appear as 0.00)
164 |
165 | fig = plt.figure(figsize=(12, 9), tight_layout=True)
166 | sn.set(font_scale=1.0 if self.nc < 50 else 0.8) # for label size
167 | labels = (0 < len(names) < 99) and len(names) == self.nc # apply names to ticklabels
168 | sn.heatmap(array, annot=self.nc < 30, annot_kws={"size": 8}, cmap='Blues', fmt='.2f', square=True,
169 | xticklabels=names + ['background FP'] if labels else "auto",
170 | yticklabels=names + ['background FN'] if labels else "auto").set_facecolor((1, 1, 1))
171 | fig.axes[0].set_xlabel('True')
172 | fig.axes[0].set_ylabel('Predicted')
173 | fig.savefig(Path(save_dir) / 'confusion_matrix.png', dpi=250)
174 | except Exception as e:
175 | pass
176 |
177 | def print(self):
178 | for i in range(self.nc + 1):
179 | print(' '.join(map(str, self.matrix[i])))
180 |
181 |
182 | # Plots ----------------------------------------------------------------------------------------------------------------
183 |
184 | def plot_pr_curve(px, py, ap, save_dir='.', names=()):
185 | fig, ax = plt.subplots(1, 1, figsize=(9, 6), tight_layout=True)
186 | py = np.stack(py, axis=1)
187 |
188 | if 0 < len(names) < 21: # show mAP in legend if < 10 classes
189 | for i, y in enumerate(py.T):
190 | ax.plot(px, y, linewidth=1, label=f'{names[i]} %.3f' % ap[i, 0]) # plot(recall, precision)
191 | else:
192 | ax.plot(px, py, linewidth=1, color='grey') # plot(recall, precision)
193 |
194 | ax.plot(px, py.mean(1), linewidth=3, color='blue', label='all classes %.3f mAP@0.5' % ap[:, 0].mean())
195 | ax.set_xlabel('Recall')
196 | ax.set_ylabel('Precision')
197 | ax.set_xlim(0, 1)
198 | ax.set_ylim(0, 1)
199 | plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left")
200 | fig.savefig(Path(save_dir) / 'precision_recall_curve.png', dpi=250)
201 |
--------------------------------------------------------------------------------
/utils/plots.py:
--------------------------------------------------------------------------------
1 | # Plotting utils
2 |
3 | import glob
4 | import math
5 | import os
6 | import random
7 | from copy import copy
8 | from pathlib import Path
9 |
10 | import cv2
11 | import matplotlib
12 | import matplotlib.pyplot as plt
13 | import numpy as np
14 | import pandas as pd
15 | import seaborn as sns
16 | import torch
17 | import yaml
18 | from PIL import Image, ImageDraw
19 | from scipy.signal import butter, filtfilt
20 |
21 | from utils.general import xywh2xyxy, xyxy2xywh
22 | from utils.metrics import fitness
23 |
24 | # Settings
25 | matplotlib.rc('font', **{'size': 11})
26 | matplotlib.use('Agg') # for writing to files only
27 |
28 |
29 | def color_list():
30 | # Return first 10 plt colors as (r,g,b) https://stackoverflow.com/questions/51350872/python-from-color-name-to-rgb
31 | def hex2rgb(h):
32 | return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4))
33 |
34 | return [hex2rgb(h) for h in plt.rcParams['axes.prop_cycle'].by_key()['color']]
35 |
36 |
37 | def hist2d(x, y, n=100):
38 | # 2d histogram used in labels.png and evolve.png
39 | xedges, yedges = np.linspace(x.min(), x.max(), n), np.linspace(y.min(), y.max(), n)
40 | hist, xedges, yedges = np.histogram2d(x, y, (xedges, yedges))
41 | xidx = np.clip(np.digitize(x, xedges) - 1, 0, hist.shape[0] - 1)
42 | yidx = np.clip(np.digitize(y, yedges) - 1, 0, hist.shape[1] - 1)
43 | return np.log(hist[xidx, yidx])
44 |
45 |
46 | def butter_lowpass_filtfilt(data, cutoff=1500, fs=50000, order=5):
47 | # https://stackoverflow.com/questions/28536191/how-to-filter-smooth-with-scipy-numpy
48 | def butter_lowpass(cutoff, fs, order):
49 | nyq = 0.5 * fs
50 | normal_cutoff = cutoff / nyq
51 | return butter(order, normal_cutoff, btype='low', analog=False)
52 |
53 | b, a = butter_lowpass(cutoff, fs, order=order)
54 | return filtfilt(b, a, data) # forward-backward filter
55 |
56 |
57 | def plot_one_box(x, img, color=None, label=None, line_thickness=3):
58 | # Plots one bounding box on image img
59 | tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1 # line/font thickness
60 | color = color or [random.randint(0, 255) for _ in range(3)]
61 | c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
62 | cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
63 | if label:
64 | tf = max(tl - 1, 1) # font thickness
65 | t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
66 | c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
67 | cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA) # filled
68 | cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
69 |
70 |
71 | def plot_wh_methods(): # from utils.plots import *; plot_wh_methods()
72 | # Compares the two methods for width-height anchor multiplication
73 | # https://github.com/ultralytics/yolov3/issues/168
74 | x = np.arange(-4.0, 4.0, .1)
75 | ya = np.exp(x)
76 | yb = torch.sigmoid(torch.from_numpy(x)).numpy() * 2
77 |
78 | fig = plt.figure(figsize=(6, 3), tight_layout=True)
79 | plt.plot(x, ya, '.-', label='YOLOv3')
80 | plt.plot(x, yb ** 2, '.-', label='YOLOv5 ^2')
81 | plt.plot(x, yb ** 1.6, '.-', label='YOLOv5 ^1.6')
82 | plt.xlim(left=-4, right=4)
83 | plt.ylim(bottom=0, top=6)
84 | plt.xlabel('input')
85 | plt.ylabel('output')
86 | plt.grid()
87 | plt.legend()
88 | fig.savefig('comparison.png', dpi=200)
89 |
90 |
91 | def output_to_target(output):
92 | # Convert model output to target format [batch_id, class_id, x, y, w, h, conf]
93 | targets = []
94 | for i, o in enumerate(output):
95 | for *box, conf, cls in o.cpu().numpy():
96 | targets.append([i, cls, *list(*xyxy2xywh(np.array(box)[None])), conf])
97 | return np.array(targets)
98 |
99 | def plot_images_together(images, targets, out, paths=None, fname='images.jpg', names=None, max_size=640, max_subplots=16):
100 | if isinstance(images, torch.Tensor):
101 | images = images.cpu().float().numpy()
102 | if isinstance(targets, torch.Tensor):
103 | targets = targets.cpu().numpy()
104 | if isinstance(out, torch.Tensor):
105 | out = out.cpu().numpy()
106 | if images.shape[1] != 3:
107 | images = images[:, :3, :, :]
108 |
109 | # un-normalise
110 | if np.max(images[0]) <= 1:
111 | images *= 255
112 |
113 | tl = 1 # line thickness
114 | tf = max(tl - 1, 1) # font thickness
115 | bs, _, h, w = images.shape # batch size, _, height, width
116 | bs = min(bs, max_subplots) # limit plot images
117 |
118 | # Check if we should resize
119 | scale_factor = max_size / max(h, w)
120 | if scale_factor < 1:
121 | h = math.ceil(scale_factor * h)
122 | w = math.ceil(scale_factor * w)
123 |
124 | colors = color_list() # list of colors
125 | for i, img in enumerate(images):
126 | if i == max_subplots: # if last batch has fewer images than we expect
127 | break
128 | img = img.transpose(1, 2, 0)
129 | if scale_factor < 1:
130 | img = cv2.resize(img, (w, h))
131 |
132 | if len(targets) > 0:
133 | image_targets = targets[targets[:, 0] == i]
134 | boxes = xywh2xyxy(image_targets[:, 2:6]).T
135 | classes = image_targets[:, 1].astype('int')
136 | labels = image_targets.shape[1] == 6 # labels if no conf column
137 | conf = None if labels else image_targets[:, 6] # check for confidence presence (label vs pred)
138 |
139 | if boxes.shape[1]:
140 | if boxes.max() <= 1.01: # if normalized with tolerance 0.01
141 | boxes[[0, 2]] *= w # scale to pixels
142 | boxes[[1, 3]] *= h
143 | elif scale_factor < 1: # absolute coords need scale if image scales
144 | boxes *= scale_factor
145 | boxes[[0, 2]] += int(0)
146 | boxes[[1, 3]] += int(0)
147 | for j, box in enumerate(boxes.T):
148 | cls = int(classes[j])
149 | color = colors[cls % len(colors)+3]
150 | cls = names[cls] if names else cls
151 | if labels or conf[j] > 0.25: # 0.25 conf thresh
152 | label = ''
153 | img = cv2.resize(img, (w, h))
154 | plot_one_box(box, img, label=label, color=color, line_thickness=tl)
155 |
156 | if len(out) > 0:
157 | image_targets = out[out[:, 0] == i]
158 | boxes = xywh2xyxy(image_targets[:, 2:6]).T
159 | classes = image_targets[:, 1].astype('int')
160 | labels = image_targets.shape[1] == 6 # labels if no conf column
161 | conf = None if labels else image_targets[:, 6] # check for confidence presence (label vs pred)
162 |
163 | if boxes.shape[1]:
164 | if boxes.max() <= 1.01: # if normalized with tolerance 0.01
165 | boxes[[0, 2]] *= w # scale to pixels
166 | boxes[[1, 3]] *= h
167 | elif scale_factor < 1: # absolute coords need scale if image scales
168 | boxes *= scale_factor
169 | boxes[[0, 2]] += int(0)
170 | boxes[[1, 3]] += int(0)
171 | for j, box in enumerate(boxes.T):
172 | cls = int(classes[j])
173 | color = colors[cls % len(colors)]
174 | cls = names[cls] if names else cls
175 | if labels or conf[j] > 0.25: # 0.25 conf thresh
176 | label = '%s %.1f' % (cls, conf[j])
177 | img = cv2.resize(img, (w, h))
178 | plot_one_box(box, img, label=label, color=color, line_thickness=tl)
179 | cv2.imwrite(
180 | os.path.dirname(str(fname)) + os.path.sep + os.path.basename(paths[i])[:-4] + os.path.basename(str(fname)),
181 | cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
182 |
183 |
184 | def plot_images_each(images, targets, paths=None, fname='images.jpg', names=None, max_size=640, max_subplots=16):
185 | # Plot image grid with labels
186 |
187 | if isinstance(images, torch.Tensor):
188 | images = images.cpu().float().numpy()
189 | if isinstance(targets, torch.Tensor):
190 | targets = targets.cpu().numpy()
191 | if images.shape[1]!=3:
192 | images = images[:,:3,:,:]
193 |
194 | # un-normalise
195 | if np.max(images[0]) <= 1:
196 | images *= 255
197 |
198 | tl = 1 # line thickness
199 | tf = max(tl - 1, 1) # font thickness
200 | bs, _, h, w = images.shape # batch size, _, height, width
201 | bs = min(bs, max_subplots) # limit plot images
202 |
203 | # Check if we should resize
204 | scale_factor = max_size / max(h, w)
205 | if scale_factor < 1:
206 | h = math.ceil(scale_factor * h)
207 | w = math.ceil(scale_factor * w)
208 |
209 | colors = color_list() # list of colors
210 | for i, img in enumerate(images):
211 | if i == max_subplots: # if last batch has fewer images than we expect
212 | break
213 | img = img.transpose(1, 2, 0)
214 | if scale_factor < 1:
215 | img = cv2.resize(img, (w, h))
216 |
217 | if len(targets) > 0:
218 | image_targets = targets[targets[:, 0] == i]
219 | boxes = xywh2xyxy(image_targets[:, 2:6]).T
220 | classes = image_targets[:, 1].astype('int')
221 | labels = image_targets.shape[1] == 6 # labels if no conf column
222 | conf = None if labels else image_targets[:, 6] # check for confidence presence (label vs pred)
223 |
224 | if boxes.shape[1]:
225 | if boxes.max() <= 1.01: # if normalized with tolerance 0.01
226 | boxes[[0, 2]] *= w # scale to pixels
227 | boxes[[1, 3]] *= h
228 | elif scale_factor < 1: # absolute coords need scale if image scales
229 | boxes *= scale_factor
230 | boxes[[0, 2]] += int(0)
231 | boxes[[1, 3]] += int(0)
232 | for j, box in enumerate(boxes.T):
233 | cls = int(classes[j])
234 | color = colors[cls % len(colors)]
235 | cls = names[cls] if names else cls
236 | if labels or conf[j] > 0.25: # 0.25 conf thresh
237 | label = '%s' % cls if labels else '%s %.1f' % (cls, conf[j])
238 | img = cv2.resize(img, (w, h))
239 | plot_one_box(box, img, label=label, color=color, line_thickness=tl)
240 | cv2.imwrite(os.path.dirname(str(fname))+os.path.sep+ os.path.basename(paths[i])[:-4]+os.path.basename(str(fname)),cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
241 |
242 | def plot_images(images, targets, paths=None, fname='images.jpg', names=None, max_size=640, max_subplots=16):
243 | # Plot image grid with labels
244 |
245 | if isinstance(images, torch.Tensor):
246 | images = images.cpu().float().numpy()
247 | if isinstance(targets, torch.Tensor):
248 | targets = targets.cpu().numpy()
249 | if images.shape[1]!=3:
250 | images = images[:,:3,:,:]
251 |
252 | # un-normalise
253 | if np.max(images[0]) <= 1:
254 | images *= 255
255 |
256 | tl = 3 # line thickness
257 | tf = max(tl - 1, 1) # font thickness
258 | bs, _, h, w = images.shape # batch size, _, height, width
259 | bs = min(bs, max_subplots) # limit plot images
260 | ns = np.ceil(bs ** 0.5) # number of subplots (square)
261 |
262 | # Check if we should resize
263 | scale_factor = max_size / max(h, w)
264 | if scale_factor < 1:
265 | h = math.ceil(scale_factor * h)
266 | w = math.ceil(scale_factor * w)
267 |
268 | colors = color_list() # list of colors
269 | mosaic = np.full((int(ns * h), int(ns * w), 3), 255, dtype=np.uint8) # init
270 | for i, img in enumerate(images):
271 | if i == max_subplots: # if last batch has fewer images than we expect
272 | break
273 |
274 | block_x = int(w * (i // ns))
275 | block_y = int(h * (i % ns))
276 |
277 | img = img.transpose(1, 2, 0)
278 | if scale_factor < 1:
279 | img = cv2.resize(img, (w, h))
280 |
281 | mosaic[block_y:block_y + h, block_x:block_x + w, :] = img
282 | if len(targets) > 0:
283 | image_targets = targets[targets[:, 0] == i]
284 | boxes = xywh2xyxy(image_targets[:, 2:6]).T
285 | classes = image_targets[:, 1].astype('int')
286 | labels = image_targets.shape[1] == 6 # labels if no conf column
287 | conf = None if labels else image_targets[:, 6] # check for confidence presence (label vs pred)
288 |
289 | if boxes.shape[1]:
290 | if boxes.max() <= 1.01: # if normalized with tolerance 0.01
291 | boxes[[0, 2]] *= w # scale to pixels
292 | boxes[[1, 3]] *= h
293 | elif scale_factor < 1: # absolute coords need scale if image scales
294 | boxes *= scale_factor
295 | boxes[[0, 2]] += block_x
296 | boxes[[1, 3]] += block_y
297 | for j, box in enumerate(boxes.T):
298 | cls = int(classes[j])
299 | color = colors[cls % len(colors)]
300 | cls = names[cls] if names else cls
301 | if labels or conf[j] > 0.25: # 0.25 conf thresh
302 | label = '%s' % cls if labels else '%s %.1f' % (cls, conf[j])
303 | plot_one_box(box, mosaic, label=label, color=color, line_thickness=tl)
304 |
305 | # Draw image filename labels
306 | if paths:
307 | label = Path(paths[i]).name[:40] # trim to 40 char
308 | t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
309 | cv2.putText(mosaic, label, (block_x + 5, block_y + t_size[1] + 5), 0, tl / 3, [220, 220, 220], thickness=tf,
310 | lineType=cv2.LINE_AA)
311 |
312 | # Image border
313 | cv2.rectangle(mosaic, (block_x, block_y), (block_x + w, block_y + h), (255, 255, 255), thickness=3)
314 |
315 | if fname:
316 | r = min(1280. / max(h, w) / ns, 1.0) # ratio to limit image size
317 | mosaic = cv2.resize(mosaic, (int(ns * w * r), int(ns * h * r)), interpolation=cv2.INTER_AREA)
318 | # cv2.imwrite(fname, cv2.cvtColor(mosaic, cv2.COLOR_BGR2RGB)) # cv2 save
319 | Image.fromarray(mosaic).save(fname) # PIL save
320 | return mosaic
321 |
322 |
323 | def plot_lr_scheduler(optimizer, scheduler, epochs=300, save_dir=''):
324 | # Plot LR simulating training for full epochs
325 | optimizer, scheduler = copy(optimizer), copy(scheduler) # do not modify originals
326 | y = []
327 | for _ in range(epochs):
328 | scheduler.step()
329 | y.append(optimizer.param_groups[0]['lr'])
330 | plt.plot(y, '.-', label='LR')
331 | plt.xlabel('epoch')
332 | plt.ylabel('LR')
333 | plt.grid()
334 | plt.xlim(0, epochs)
335 | plt.ylim(0)
336 | plt.savefig(Path(save_dir) / 'LR.png', dpi=200)
337 | plt.close()
338 |
339 |
340 | def plot_test_txt(): # from utils.plots import *; plot_test()
341 | # Plot test.txt histograms
342 | x = np.loadtxt('test.txt', dtype=np.float32)
343 | box = xyxy2xywh(x[:, :4])
344 | cx, cy = box[:, 0], box[:, 1]
345 |
346 | fig, ax = plt.subplots(1, 1, figsize=(6, 6), tight_layout=True)
347 | ax.hist2d(cx, cy, bins=600, cmax=10, cmin=0)
348 | ax.set_aspect('equal')
349 | plt.savefig('hist2d.png', dpi=300)
350 |
351 | fig, ax = plt.subplots(1, 2, figsize=(12, 6), tight_layout=True)
352 | ax[0].hist(cx, bins=600)
353 | ax[1].hist(cy, bins=600)
354 | plt.savefig('hist1d.png', dpi=200)
355 |
356 |
357 | def plot_targets_txt(): # from utils.plots import *; plot_targets_txt()
358 | # Plot targets.txt histograms
359 | x = np.loadtxt('targets.txt', dtype=np.float32).T
360 | s = ['x targets', 'y targets', 'width targets', 'height targets']
361 | fig, ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True)
362 | ax = ax.ravel()
363 | for i in range(4):
364 | ax[i].hist(x[i], bins=100, label='%.3g +/- %.3g' % (x[i].mean(), x[i].std()))
365 | ax[i].legend()
366 | ax[i].set_title(s[i])
367 | plt.savefig('targets.jpg', dpi=200)
368 |
369 |
370 | def plot_study_txt(path='study/', x=None): # from utils.plots import *; plot_study_txt()
371 | # Plot study.txt generated by test.py
372 | fig, ax = plt.subplots(2, 4, figsize=(10, 6), tight_layout=True)
373 | ax = ax.ravel()
374 |
375 | fig2, ax2 = plt.subplots(1, 1, figsize=(8, 4), tight_layout=True)
376 | for f in [Path(path) / f'study_coco_{x}.txt' for x in ['yolov5s', 'yolov5m', 'yolov5l', 'yolov5x']]:
377 | y = np.loadtxt(f, dtype=np.float32, usecols=[0, 1, 2, 3, 7, 8, 9], ndmin=2).T
378 | x = np.arange(y.shape[1]) if x is None else np.array(x)
379 | s = ['P', 'R', 'mAP@.5', 'mAP@.5:.95', 't_inference (ms/img)', 't_NMS (ms/img)', 't_total (ms/img)']
380 | for i in range(7):
381 | ax[i].plot(x, y[i], '.-', linewidth=2, markersize=8)
382 | ax[i].set_title(s[i])
383 |
384 | j = y[3].argmax() + 1
385 | ax2.plot(y[6, :j], y[3, :j] * 1E2, '.-', linewidth=2, markersize=8,
386 | label=f.stem.replace('study_coco_', '').replace('yolo', 'YOLO'))
387 |
388 | ax2.plot(1E3 / np.array([209, 140, 97, 58, 35, 18]), [34.6, 40.5, 43.0, 47.5, 49.7, 51.5],
389 | 'k.-', linewidth=2, markersize=8, alpha=.25, label='EfficientDet')
390 |
391 | ax2.grid()
392 | ax2.set_yticks(np.arange(30, 60, 5))
393 | ax2.set_xlim(0, 30)
394 | ax2.set_ylim(29, 51)
395 | ax2.set_xlabel('GPU Speed (ms/img)')
396 | ax2.set_ylabel('COCO AP val')
397 | ax2.legend(loc='lower right')
398 | plt.savefig('test_study.png', dpi=300)
399 |
400 |
401 | def plot_labels(labels, names=(), save_dir=Path(''), loggers=None):
402 | # plot dataset labels
403 | print('Plotting labels... ')
404 | c, b = labels[:, 0], labels[:, 1:].transpose() # classes, boxes
405 | nc = int(c.max() + 1) # number of classes
406 | colors = color_list()
407 | x = pd.DataFrame(b.transpose(), columns=['x', 'y', 'width', 'height'])
408 |
409 | # seaborn correlogram
410 | sns.pairplot(x, corner=True, diag_kind='auto', kind='hist', diag_kws=dict(bins=50), plot_kws=dict(pmax=0.9))
411 | plt.savefig(save_dir / 'labels_correlogram.jpg', dpi=200)
412 | plt.close()
413 |
414 | # matplotlib labels
415 | matplotlib.use('svg') # faster
416 | ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True)[1].ravel()
417 | ax[0].hist(c, bins=np.linspace(0, nc, nc + 1) - 0.5, rwidth=0.8)
418 | ax[0].set_ylabel('instances')
419 | if 0 < len(names) < 30:
420 | ax[0].set_xticks(range(len(names)))
421 | ax[0].set_xticklabels(names, rotation=90, fontsize=10)
422 | else:
423 | ax[0].set_xlabel('classes')
424 | sns.histplot(x, x='x', y='y', ax=ax[2], bins=50, pmax=0.9)
425 | sns.histplot(x, x='width', y='height', ax=ax[3], bins=50, pmax=0.9)
426 |
427 | # rectangles
428 | labels[:, 1:3] = 0.5 # center
429 | labels[:, 1:] = xywh2xyxy(labels[:, 1:]) * 2000
430 | img = Image.fromarray(np.ones((2000, 2000, 3), dtype=np.uint8) * 255)
431 | for cls, *box in labels[:1000]:
432 | ImageDraw.Draw(img).rectangle(box, width=1, outline=colors[int(cls) % 10]) # plot
433 | ax[1].imshow(img)
434 | ax[1].axis('off')
435 |
436 | for a in [0, 1, 2, 3]:
437 | for s in ['top', 'right', 'left', 'bottom']:
438 | ax[a].spines[s].set_visible(False)
439 |
440 | plt.savefig(save_dir / 'labels.jpg', dpi=200)
441 | matplotlib.use('Agg')
442 | plt.close()
443 |
444 | # loggers
445 | for k, v in loggers.items() or {}:
446 | if k == 'wandb' and v:
447 | v.log({"Labels": [v.Image(str(x), caption=x.name) for x in save_dir.glob('*labels*.jpg')]})
448 |
449 |
450 | def plot_evolution(yaml_file='data/hyp.finetune.yaml'): # from utils.plots import *; plot_evolution()
451 | # Plot hyperparameter evolution results in evolve.txt
452 | with open(yaml_file) as f:
453 | hyp = yaml.load(f, Loader=yaml.SafeLoader)
454 | x = np.loadtxt('evolve.txt', ndmin=2)
455 | f = fitness(x)
456 | # weights = (f - f.min()) ** 2 # for weighted results
457 | plt.figure(figsize=(10, 12), tight_layout=True)
458 | matplotlib.rc('font', **{'size': 8})
459 | for i, (k, v) in enumerate(hyp.items()):
460 | y = x[:, i + 7]
461 | # mu = (y * weights).sum() / weights.sum() # best weighted result
462 | mu = y[f.argmax()] # best single result
463 | plt.subplot(6, 5, i + 1)
464 | plt.scatter(y, f, c=hist2d(y, f, 20), cmap='viridis', alpha=.8, edgecolors='none')
465 | plt.plot(mu, f.max(), 'k+', markersize=15)
466 | plt.title('%s = %.3g' % (k, mu), fontdict={'size': 9}) # limit to 40 characters
467 | if i % 5 != 0:
468 | plt.yticks([])
469 | print('%15s: %.3g' % (k, mu))
470 | plt.savefig('evolve.png', dpi=200)
471 | print('\nPlot saved as evolve.png')
472 |
473 |
474 | def profile_idetection(start=0, stop=0, labels=(), save_dir=''):
475 | # Plot iDetection '*.txt' per-image logs. from utils.plots import *; profile_idetection()
476 | ax = plt.subplots(2, 4, figsize=(12, 6), tight_layout=True)[1].ravel()
477 | s = ['Images', 'Free Storage (GB)', 'RAM Usage (GB)', 'Battery', 'dt_raw (ms)', 'dt_smooth (ms)', 'real-world FPS']
478 | files = list(Path(save_dir).glob('frames*.txt'))
479 | for fi, f in enumerate(files):
480 | try:
481 | results = np.loadtxt(f, ndmin=2).T[:, 90:-30] # clip first and last rows
482 | n = results.shape[1] # number of rows
483 | x = np.arange(start, min(stop, n) if stop else n)
484 | results = results[:, x]
485 | t = (results[0] - results[0].min()) # set t0=0s
486 | results[0] = x
487 | for i, a in enumerate(ax):
488 | if i < len(results):
489 | label = labels[fi] if len(labels) else f.stem.replace('frames_', '')
490 | a.plot(t, results[i], marker='.', label=label, linewidth=1, markersize=5)
491 | a.set_title(s[i])
492 | a.set_xlabel('time (s)')
493 | # if fi == len(files) - 1:
494 | # a.set_ylim(bottom=0)
495 | for side in ['top', 'right']:
496 | a.spines[side].set_visible(False)
497 | else:
498 | a.remove()
499 | except Exception as e:
500 | print('Warning: Plotting error for %s; %s' % (f, e))
501 |
502 | ax[1].legend()
503 | plt.savefig(Path(save_dir) / 'idetection_profile.png', dpi=200)
504 |
505 |
506 | def plot_results_overlay(start=0, stop=0): # from utils.plots import *; plot_results_overlay()
507 | # Plot training 'results*.txt', overlaying train and val losses
508 | s = ['train', 'train', 'train', 'Precision', 'mAP@0.5', 'val', 'val', 'val', 'Recall', 'mAP@0.5:0.95'] # legends
509 | t = ['Box', 'Objectness', 'Classification', 'P-R', 'mAP-F1'] # titles
510 | for f in sorted(glob.glob('results*.txt') + glob.glob('../../Downloads/results*.txt')):
511 | results = np.loadtxt(f, usecols=[2, 3, 4, 8, 9, 12, 13, 14, 10, 11], ndmin=2).T
512 | n = results.shape[1] # number of rows
513 | x = range(start, min(stop, n) if stop else n)
514 | fig, ax = plt.subplots(1, 5, figsize=(14, 3.5), tight_layout=True)
515 | ax = ax.ravel()
516 | for i in range(5):
517 | for j in [i, i + 5]:
518 | y = results[j, x]
519 | ax[i].plot(x, y, marker='.', label=s[j])
520 | # y_smooth = butter_lowpass_filtfilt(y)
521 | # ax[i].plot(x, np.gradient(y_smooth), marker='.', label=s[j])
522 |
523 | ax[i].set_title(t[i])
524 | ax[i].legend()
525 | ax[i].set_ylabel(f) if i == 0 else None # add filename
526 | fig.savefig(f.replace('.txt', '.png'), dpi=200)
527 |
528 |
529 | def plot_results(start=0, stop=0, bucket='', id=(), labels=(), save_dir=''):
530 | # Plot training 'results*.txt'. from utils.plots import *; plot_results(save_dir='runs/train/exp')
531 | fig, ax = plt.subplots(2, 5, figsize=(12, 6), tight_layout=True)
532 | ax = ax.ravel()
533 | s = ['Box', 'Objectness', 'Classification', 'Precision', 'Recall',
534 | 'val Box', 'val Objectness', 'val Classification', 'mAP@0.5', 'mAP@0.5:0.95']
535 | if bucket:
536 | # files = ['https://storage.googleapis.com/%s/results%g.txt' % (bucket, x) for x in id]
537 | files = ['results%g.txt' % x for x in id]
538 | c = ('gsutil cp ' + '%s ' * len(files) + '.') % tuple('gs://%s/results%g.txt' % (bucket, x) for x in id)
539 | os.system(c)
540 | else:
541 | files = list(Path(save_dir).glob('results*.txt'))
542 | assert len(files), 'No results.txt files found in %s, nothing to plot.' % os.path.abspath(save_dir)
543 | for fi, f in enumerate(files):
544 | try:
545 | results = np.loadtxt(f, usecols=[2, 3, 4, 8, 9, 12, 13, 14, 10, 11], ndmin=2).T
546 | n = results.shape[1] # number of rows
547 | x = range(start, min(stop, n) if stop else n)
548 | for i in range(10):
549 | y = results[i, x]
550 | if i in [0, 1, 2, 5, 6, 7]:
551 | y[y == 0] = np.nan # don't show zero loss values
552 | # y /= y[0] # normalize
553 | label = labels[fi] if len(labels) else f.stem
554 | ax[i].plot(x, y, marker='.', label=label, linewidth=2, markersize=8)
555 | ax[i].set_title(s[i])
556 | # if i in [5, 6, 7]: # share train and val loss y axes
557 | # ax[i].get_shared_y_axes().join(ax[i], ax[i - 5])
558 | except Exception as e:
559 | print('Warning: Plotting error for %s; %s' % (f, e))
560 |
561 | ax[1].legend()
562 | fig.savefig(Path(save_dir) / 'results.png', dpi=200)
563 |
564 | colors = color_list()
565 | pass
--------------------------------------------------------------------------------
/utils/torch_utils.py:
--------------------------------------------------------------------------------
1 | # PyTorch utils
2 |
3 | import logging
4 | import math
5 | import os
6 | import subprocess
7 | import time
8 | from contextlib import contextmanager
9 | from copy import deepcopy
10 | from pathlib import Path
11 |
12 | import torch
13 | import torch.backends.cudnn as cudnn
14 | import torch.nn as nn
15 | import torch.nn.functional as F
16 | import torchvision
17 |
18 | try:
19 | import thop # for FLOPS computation
20 | except ImportError:
21 | thop = None
22 | logger = logging.getLogger(__name__)
23 |
24 |
25 | @contextmanager
26 | def torch_distributed_zero_first(local_rank: int):
27 | """
28 | Decorator to make all processes in distributed training wait for each local_master to do something.
29 | """
30 | if local_rank not in [-1, 0]:
31 | torch.distributed.barrier()
32 | yield
33 | if local_rank == 0:
34 | torch.distributed.barrier()
35 |
36 |
37 | def init_torch_seeds(seed=0):
38 | # Speed-reproducibility tradeoff https://pytorch.org/docs/stable/notes/randomness.html
39 | torch.manual_seed(seed)
40 | if seed == 0: # slower, more reproducible
41 | cudnn.benchmark, cudnn.deterministic = False, True
42 | else: # faster, less reproducible
43 | cudnn.benchmark, cudnn.deterministic = True, False
44 |
45 |
46 | def git_describe():
47 | # return human-readable git description, i.e. v5.0-5-g3e25f1e https://git-scm.com/docs/git-describe
48 | if Path('.git').exists():
49 | return subprocess.check_output('git describe --tags --long --always', shell=True).decode('utf-8')[:-1]
50 | else:
51 | return ''
52 |
53 |
54 | def select_device(device='', batch_size=None):
55 | # device = 'cpu' or '0' or '0,1,2,3'
56 | s = f'YOLOv5 {git_describe()} torch {torch.__version__} ' # string
57 | cpu = device.lower() == 'cpu'
58 | if cpu:
59 | os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # force torch.cuda.is_available() = False
60 | elif device: # non-cpu device requested
61 | os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable
62 | assert torch.cuda.is_available(), f'CUDA unavailable, invalid device {device} requested' # check availability
63 |
64 | cuda = not cpu and torch.cuda.is_available()
65 | if cuda:
66 | n = torch.cuda.device_count()
67 | if n > 1 and batch_size: # check that batch_size is compatible with device_count
68 | assert batch_size % n == 0, f'batch-size {batch_size} not multiple of GPU count {n}'
69 | space = ' ' * len(s)
70 | for i, d in enumerate(device.split(',') if device else range(n)):
71 | p = torch.cuda.get_device_properties(i)
72 | s += f"{'' if i == 0 else space}CUDA:{d} ({p.name}, {p.total_memory / 1024 ** 2}MB)\n" # bytes to MB
73 | else:
74 | s += 'CPU\n'
75 |
76 | logger.info(s) # skip a line
77 | return torch.device('cuda:0' if cuda else 'cpu')
78 |
79 |
80 | def time_synchronized():
81 | # pytorch-accurate time
82 | if torch.cuda.is_available():
83 | torch.cuda.synchronize()
84 | return time.time()
85 |
86 |
87 | def profile(x, ops, n=100, device=None):
88 | # profile a pytorch module or list of modules. Example usage:
89 | # x = torch.randn(16, 3, 640, 640) # input
90 | # m1 = lambda x: x * torch.sigmoid(x)
91 | # m2 = nn.SiLU()
92 | # profile(x, [m1, m2], n=100) # profile speed over 100 iterations
93 |
94 | device = device or torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
95 | x = x.to(device)
96 | x.requires_grad = True
97 | print(torch.__version__, device.type, torch.cuda.get_device_properties(0) if device.type == 'cuda' else '')
98 | print(f"\n{'Params':>12s}{'GFLOPS':>12s}{'forward (ms)':>16s}{'backward (ms)':>16s}{'input':>24s}{'output':>24s}")
99 | for m in ops if isinstance(ops, list) else [ops]:
100 | m = m.to(device) if hasattr(m, 'to') else m # device
101 | m = m.half() if hasattr(m, 'half') and isinstance(x, torch.Tensor) and x.dtype is torch.float16 else m # type
102 | dtf, dtb, t = 0., 0., [0., 0., 0.] # dt forward, backward
103 | try:
104 | flops = thop.profile(m, inputs=(x,), verbose=False)[0] / 1E9 * 2 # GFLOPS
105 | except:
106 | flops = 0
107 |
108 | for _ in range(n):
109 | t[0] = time_synchronized()
110 | y = m(x)
111 | t[1] = time_synchronized()
112 | try:
113 | _ = y.sum().backward()
114 | t[2] = time_synchronized()
115 | except: # no backward method
116 | t[2] = float('nan')
117 | dtf += (t[1] - t[0]) * 1000 / n # ms per op forward
118 | dtb += (t[2] - t[1]) * 1000 / n # ms per op backward
119 |
120 | s_in = tuple(x.shape) if isinstance(x, torch.Tensor) else 'list'
121 | s_out = tuple(y.shape) if isinstance(y, torch.Tensor) else 'list'
122 | p = sum(list(x.numel() for x in m.parameters())) if isinstance(m, nn.Module) else 0 # parameters
123 | print(f'{p:12.4g}{flops:12.4g}{dtf:16.4g}{dtb:16.4g}{str(s_in):>24s}{str(s_out):>24s}')
124 |
125 |
126 | def is_parallel(model):
127 | return type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel)
128 |
129 |
130 | def intersect_dicts(da, db, exclude=()):
131 | # Dictionary intersection of matching keys and shapes, omitting 'exclude' keys, using da values
132 | return {k: v for k, v in da.items() if k in db and not any(x in k for x in exclude) and v.shape == db[k].shape}
133 |
134 |
135 | def initialize_weights(model):
136 | for m in model.modules():
137 | t = type(m)
138 | if t is nn.Conv2d:
139 | pass # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
140 | elif t is nn.BatchNorm2d:
141 | m.eps = 1e-3
142 | m.momentum = 0.03
143 | elif t in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6]:
144 | m.inplace = True
145 |
146 |
147 | def find_modules(model, mclass=nn.Conv2d):
148 | # Finds layer indices matching module class 'mclass'
149 | return [i for i, m in enumerate(model.module_list) if isinstance(m, mclass)]
150 |
151 |
152 | def sparsity(model):
153 | # Return global model sparsity
154 | a, b = 0., 0.
155 | for p in model.parameters():
156 | a += p.numel()
157 | b += (p == 0).sum()
158 | return b / a
159 |
160 |
161 | def prune(model, amount=0.3):
162 | # Prune model to requested global sparsity
163 | import torch.nn.utils.prune as prune
164 | print('Pruning model... ', end='')
165 | for name, m in model.named_modules():
166 | if isinstance(m, nn.Conv2d):
167 | prune.l1_unstructured(m, name='weight', amount=amount) # prune
168 | prune.remove(m, 'weight') # make permanent
169 | print(' %.3g global sparsity' % sparsity(model))
170 |
171 |
172 | def fuse_conv_and_bn(conv, bn):
173 | # Fuse convolution and batchnorm layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/
174 | fusedconv = nn.Conv2d(conv.in_channels,
175 | conv.out_channels,
176 | kernel_size=conv.kernel_size,
177 | stride=conv.stride,
178 | padding=conv.padding,
179 | groups=conv.groups,
180 | bias=True).requires_grad_(False).to(conv.weight.device)
181 |
182 | # prepare filters
183 | w_conv = conv.weight.clone().view(conv.out_channels, -1)
184 | w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var)))
185 | fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.size()))
186 |
187 | # prepare spatial bias
188 | b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device) if conv.bias is None else conv.bias
189 | b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps))
190 | fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn)
191 |
192 | return fusedconv
193 |
194 |
195 | def model_info(model, verbose=False, img_size=512):
196 | # Model information. img_size may be int or list, i.e. img_size=640 or img_size=[640, 320]
197 | n_p = sum(x.numel() for x in model.parameters()) # number parameters
198 | n_g = sum(x.numel() for x in model.parameters() if x.requires_grad) # number gradients
199 | if verbose:
200 | print('%5s %40s %9s %12s %20s %10s %10s' % ('layer', 'name', 'gradient', 'parameters', 'shape', 'mu', 'sigma'))
201 | for i, (name, p) in enumerate(model.named_parameters()):
202 | name = name.replace('module_list.', '')
203 | print('%5g %40s %9s %12g %20s %10.3g %10.3g' %
204 | (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std()))
205 |
206 | try: # FLOPS
207 | from thop import profile
208 | stride = int(model.stride.max()) if hasattr(model, 'stride') else 32
209 | stride=512
210 | img = torch.zeros((1, model.yaml.get('ch', 3), stride, stride), device=next(model.parameters()).device) # input
211 | flops = profile(deepcopy(model), inputs=(img,), verbose=False)[0] / 1E9 * 2 # stride GFLOPS
212 | img_size = img_size if isinstance(img_size, list) else [img_size, img_size] # expand if int/float
213 | fs = ', %.1f GFLOPS' % (flops * img_size[0] / stride * img_size[1] / stride) # 640x640 GFLOPS
214 | except (ImportError, Exception):
215 | fs = ''
216 |
217 | logger.info(f"Model Summary: {len(list(model.modules()))} layers, {n_p} parameters, {n_g} gradients{fs}")
218 |
219 |
220 | def load_classifier(name='resnet101', n=2):
221 | # Loads a pretrained model reshaped to n-class output
222 | model = torchvision.models.__dict__[name](pretrained=True)
223 |
224 | # ResNet model properties
225 | # input_size = [3, 224, 224]
226 | # input_space = 'RGB'
227 | # input_range = [0, 1]
228 | # mean = [0.485, 0.456, 0.406]
229 | # std = [0.229, 0.224, 0.225]
230 |
231 | # Reshape output to n classes
232 | filters = model.fc.weight.shape[1]
233 | model.fc.bias = nn.Parameter(torch.zeros(n), requires_grad=True)
234 | model.fc.weight = nn.Parameter(torch.zeros(n, filters), requires_grad=True)
235 | model.fc.out_features = n
236 | return model
237 |
238 |
239 | def scale_img(img, ratio=1.0, same_shape=False, gs=32): # img(16,3,256,416)
240 | # scales img(bs,3,y,x) by ratio constrained to gs-multiple
241 | if ratio == 1.0:
242 | return img
243 | else:
244 | h, w = img.shape[2:]
245 | s = (int(h * ratio), int(w * ratio)) # new size
246 | img = F.interpolate(img, size=s, mode='bilinear', align_corners=False) # resize
247 | if not same_shape: # pad/crop img
248 | h, w = [math.ceil(x * ratio / gs) * gs for x in (h, w)]
249 | return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447) # value = imagenet mean
250 |
251 |
252 | def copy_attr(a, b, include=(), exclude=()):
253 | # Copy attributes from b to a, options to only include [...] and to exclude [...]
254 | for k, v in b.__dict__.items():
255 | if (len(include) and k not in include) or k.startswith('_') or k in exclude:
256 | continue
257 | else:
258 | setattr(a, k, v)
259 |
260 |
261 | class ModelEMA:
262 | """ Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models
263 | Keep a moving average of everything in the model state_dict (parameters and buffers).
264 | This is intended to allow functionality like
265 | https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage
266 | A smoothed version of the weights is necessary for some training schemes to perform well.
267 | This class is sensitive where it is initialized in the sequence of model init,
268 | GPU assignment and distributed training wrappers.
269 | """
270 |
271 | def __init__(self, model, decay=0.9999, updates=0):
272 | # Create EMA
273 | self.ema = deepcopy(model.module if is_parallel(model) else model).eval() # FP32 EMA
274 | # if next(model.parameters()).device.type != 'cpu':
275 | # self.ema.half() # FP16 EMA
276 | self.updates = updates # number of EMA updates
277 | self.decay = lambda x: decay * (1 - math.exp(-x / 2000)) # decay exponential ramp (to help early epochs)
278 | for p in self.ema.parameters():
279 | p.requires_grad_(False)
280 |
281 | def update(self, model):
282 | # Update EMA parameters
283 | with torch.no_grad():
284 | self.updates += 1
285 | d = self.decay(self.updates)
286 |
287 | msd = model.module.state_dict() if is_parallel(model) else model.state_dict() # model state_dict
288 | for k, v in self.ema.state_dict().items():
289 | if v.dtype.is_floating_point:
290 | v *= d
291 | v += (1. - d) * msd[k].detach()
292 |
293 | def update_attr(self, model, include=(), exclude=('process_group', 'reducer')):
294 | # Update EMA attributes
295 | copy_attr(self.ema, model, include, exclude)
296 |
--------------------------------------------------------------------------------