| Car Semantic Segmentation | 63 |Car Semantic Segmentation | 64 |Person Segmentation + Tracking | 65 |
|
68 | |
69 | |
70 |
| Car Semantic Segmentation | 160 |Car Semantic Segmentation | 161 |Person Segmentation + Tracking | 162 |
|
165 |  |
166 | |
167 |
4 |
5 | ## About ClearML
6 |
7 | [ClearML](https://cutt.ly/yolov5-tutorial-clearml) is an [open-source](https://github.com/allegroai/clearml) toolbox designed to save you time ⏱️.
8 |
9 | 🔨 Track every YOLOv5 training run in the experiment manager
10 |
11 | 🔧 Version and easily access your custom training data with the integrated ClearML Data Versioning Tool
12 |
13 | 🔦 Remotely train and monitor your YOLOv5 training runs using ClearML Agent
14 |
15 | 🔬 Get the very best mAP using ClearML Hyperparameter Optimization
16 |
17 | 🔭 Turn your newly trained YOLOv5 model into an API with just a few commands using ClearML Serving
18 |
19 | 
36 | 
37 |
38 | 
$ python train.py --upload_data val
78 |
79 | 
80 |
81 | {dataset}_wandb.yaml file which can be used to train from dataset artifact.
85 | $ python utils/logger/wandb/log_dataset.py --project ... --name ... --data ..
88 |
89 | 
90 |
91 | $ python train.py --data {data}_wandb.yaml
99 |
100 | 
101 |
102 | $ python train.py --save_period 1
111 |
112 | 
113 |
114 | --resume argument starts with wandb-artifact:// prefix followed by the run path, i.e, wandb-artifact://username/project/runid . This doesn't require the model checkpoint to be present on the local system.
120 |
121 | $ python train.py --resume wandb-artifact://{run_path}
124 |
125 | 
126 |
127 | --upload_dataset or
132 | train from _wandb.yaml file and set --save_period
133 |
134 | $ python train.py --resume wandb-artifact://{run_path}
137 |
138 | 
139 |
140 | 
148 |
149 | ## Environments
150 |
151 | YOLOv5 may be run in any of the following up-to-date verified environments (with all dependencies including [CUDA](https://developer.nvidia.com/cuda)/[CUDNN](https://developer.nvidia.com/cudnn), [Python](https://www.python.org/) and [PyTorch](https://pytorch.org/) preinstalled):
152 |
153 | - **Google Colab and Kaggle** notebooks with free GPU:
154 | - **Google Cloud** Deep Learning VM. See [GCP Quickstart Guide](https://github.com/ultralytics/yolov5/wiki/GCP-Quickstart)
155 | - **Amazon** Deep Learning AMI. See [AWS Quickstart Guide](https://github.com/ultralytics/yolov5/wiki/AWS-Quickstart)
156 | - **Docker Image**. See [Docker Quickstart Guide](https://github.com/ultralytics/yolov5/wiki/Docker-Quickstart) 
157 |
158 | ## Status
159 |
160 | 
161 |
162 | If this badge is green, all [YOLOv5 GitHub Actions](https://github.com/ultralytics/yolov5/actions) Continuous Integration (CI) tests are currently passing. CI tests verify correct operation of YOLOv5 training ([train.py](https://github.com/ultralytics/yolov5/blob/master/train.py)), validation ([val.py](https://github.com/ultralytics/yolov5/blob/master/val.py)), inference ([detect.py](https://github.com/ultralytics/yolov5/blob/master/detect.py)) and export ([export.py](https://github.com/ultralytics/yolov5/blob/master/export.py)) on macOS, Windows, and Ubuntu every 24 hours and on every commit.
163 |
--------------------------------------------------------------------------------
/utils/loggers/wandb/__init__.py:
--------------------------------------------------------------------------------
1 | # init
--------------------------------------------------------------------------------
/utils/loggers/wandb/log_dataset.py:
--------------------------------------------------------------------------------
1 | import argparse
2 |
3 | from wandb_utils import WandbLogger
4 |
5 | from utils.general import LOGGER
6 |
7 | WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
8 |
9 |
10 | def create_dataset_artifact(opt):
11 | logger = WandbLogger(opt, None, job_type='Dataset Creation') # TODO: return value unused
12 | if not logger.wandb:
13 | LOGGER.info("install wandb using `pip install wandb` to log the dataset")
14 |
15 |
16 | if __name__ == '__main__':
17 | parser = argparse.ArgumentParser()
18 | parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path')
19 | parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
20 | parser.add_argument('--project', type=str, default='YOLOv5', help='name of W&B Project')
21 | parser.add_argument('--entity', default=None, help='W&B entity')
22 | parser.add_argument('--name', type=str, default='log dataset', help='name of W&B run')
23 |
24 | opt = parser.parse_args()
25 | opt.resume = False # Explicitly disallow resume check for dataset upload job
26 |
27 | create_dataset_artifact(opt)
28 |
--------------------------------------------------------------------------------
/utils/loggers/wandb/sweep.py:
--------------------------------------------------------------------------------
1 | import sys
2 | from pathlib import Path
3 |
4 | import wandb
5 |
6 | FILE = Path(__file__).resolve()
7 | ROOT = FILE.parents[3] # YOLOv5 root directory
8 | if str(ROOT) not in sys.path:
9 | sys.path.append(str(ROOT)) # add ROOT to PATH
10 |
11 | from train import parse_opt, train
12 | from utils.callbacks import Callbacks
13 | from utils.general import increment_path
14 | from utils.torch_utils import select_device
15 |
16 |
17 | def sweep():
18 | wandb.init()
19 | # Get hyp dict from sweep agent. Copy because train() modifies parameters which confused wandb.
20 | hyp_dict = vars(wandb.config).get("_items").copy()
21 |
22 | # Workaround: get necessary opt args
23 | opt = parse_opt(known=True)
24 | opt.batch_size = hyp_dict.get("batch_size")
25 | opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok or opt.evolve))
26 | opt.epochs = hyp_dict.get("epochs")
27 | opt.nosave = True
28 | opt.data = hyp_dict.get("data")
29 | opt.weights = str(opt.weights)
30 | opt.cfg = str(opt.cfg)
31 | opt.data = str(opt.data)
32 | opt.hyp = str(opt.hyp)
33 | opt.project = str(opt.project)
34 | device = select_device(opt.device, batch_size=opt.batch_size)
35 |
36 | # train
37 | train(hyp_dict, opt, device, callbacks=Callbacks())
38 |
39 |
40 | if __name__ == "__main__":
41 | sweep()
42 |
--------------------------------------------------------------------------------
/utils/loggers/wandb/sweep.yaml:
--------------------------------------------------------------------------------
1 | # Hyperparameters for training
2 | # To set range-
3 | # Provide min and max values as:
4 | # parameter:
5 | #
6 | # min: scalar
7 | # max: scalar
8 | # OR
9 | #
10 | # Set a specific list of search space-
11 | # parameter:
12 | # values: [scalar1, scalar2, scalar3...]
13 | #
14 | # You can use grid, bayesian and hyperopt search strategy
15 | # For more info on configuring sweeps visit - https://docs.wandb.ai/guides/sweeps/configuration
16 |
17 | program: utils/loggers/wandb/sweep.py
18 | method: random
19 | metric:
20 | name: metrics/mAP_0.5
21 | goal: maximize
22 |
23 | parameters:
24 | # hyperparameters: set either min, max range or values list
25 | data:
26 | value: "data/coco128.yaml"
27 | batch_size:
28 | values: [64]
29 | epochs:
30 | values: [10]
31 |
32 | lr0:
33 | distribution: uniform
34 | min: 1e-5
35 | max: 1e-1
36 | lrf:
37 | distribution: uniform
38 | min: 0.01
39 | max: 1.0
40 | momentum:
41 | distribution: uniform
42 | min: 0.6
43 | max: 0.98
44 | weight_decay:
45 | distribution: uniform
46 | min: 0.0
47 | max: 0.001
48 | warmup_epochs:
49 | distribution: uniform
50 | min: 0.0
51 | max: 5.0
52 | warmup_momentum:
53 | distribution: uniform
54 | min: 0.0
55 | max: 0.95
56 | warmup_bias_lr:
57 | distribution: uniform
58 | min: 0.0
59 | max: 0.2
60 | box:
61 | distribution: uniform
62 | min: 0.02
63 | max: 0.2
64 | cls:
65 | distribution: uniform
66 | min: 0.2
67 | max: 4.0
68 | cls_pw:
69 | distribution: uniform
70 | min: 0.5
71 | max: 2.0
72 | obj:
73 | distribution: uniform
74 | min: 0.2
75 | max: 4.0
76 | obj_pw:
77 | distribution: uniform
78 | min: 0.5
79 | max: 2.0
80 | iou_t:
81 | distribution: uniform
82 | min: 0.1
83 | max: 0.7
84 | anchor_t:
85 | distribution: uniform
86 | min: 2.0
87 | max: 8.0
88 | fl_gamma:
89 | distribution: uniform
90 | min: 0.0
91 | max: 4.0
92 | hsv_h:
93 | distribution: uniform
94 | min: 0.0
95 | max: 0.1
96 | hsv_s:
97 | distribution: uniform
98 | min: 0.0
99 | max: 0.9
100 | hsv_v:
101 | distribution: uniform
102 | min: 0.0
103 | max: 0.9
104 | degrees:
105 | distribution: uniform
106 | min: 0.0
107 | max: 45.0
108 | translate:
109 | distribution: uniform
110 | min: 0.0
111 | max: 0.9
112 | scale:
113 | distribution: uniform
114 | min: 0.0
115 | max: 0.9
116 | shear:
117 | distribution: uniform
118 | min: 0.0
119 | max: 10.0
120 | perspective:
121 | distribution: uniform
122 | min: 0.0
123 | max: 0.001
124 | flipud:
125 | distribution: uniform
126 | min: 0.0
127 | max: 1.0
128 | fliplr:
129 | distribution: uniform
130 | min: 0.0
131 | max: 1.0
132 | mosaic:
133 | distribution: uniform
134 | min: 0.0
135 | max: 1.0
136 | mixup:
137 | distribution: uniform
138 | min: 0.0
139 | max: 1.0
140 | copy_paste:
141 | distribution: uniform
142 | min: 0.0
143 | max: 1.0
144 |
--------------------------------------------------------------------------------
/utils/loss.py:
--------------------------------------------------------------------------------
1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
2 | """
3 | Loss functions
4 | """
5 |
6 | import torch
7 | import torch.nn as nn
8 |
9 | from utils.metrics import bbox_iou
10 | from utils.torch_utils import de_parallel
11 |
12 |
13 | def smooth_BCE(eps=0.1): # https://github.com/ultralytics/yolov3/issues/238#issuecomment-598028441
14 | # return positive, negative label smoothing BCE targets
15 | return 1.0 - 0.5 * eps, 0.5 * eps
16 |
17 |
18 | class BCEBlurWithLogitsLoss(nn.Module):
19 | # BCEwithLogitLoss() with reduced missing label effects.
20 | def __init__(self, alpha=0.05):
21 | super().__init__()
22 | self.loss_fcn = nn.BCEWithLogitsLoss(reduction='none') # must be nn.BCEWithLogitsLoss()
23 | self.alpha = alpha
24 |
25 | def forward(self, pred, true):
26 | loss = self.loss_fcn(pred, true)
27 | pred = torch.sigmoid(pred) # prob from logits
28 | dx = pred - true # reduce only missing label effects
29 | # dx = (pred - true).abs() # reduce missing label and false label effects
30 | alpha_factor = 1 - torch.exp((dx - 1) / (self.alpha + 1e-4))
31 | loss *= alpha_factor
32 | return loss.mean()
33 |
34 |
35 | class FocalLoss(nn.Module):
36 | # Wraps focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5)
37 | def __init__(self, loss_fcn, gamma=1.5, alpha=0.25):
38 | super().__init__()
39 | self.loss_fcn = loss_fcn # must be nn.BCEWithLogitsLoss()
40 | self.gamma = gamma
41 | self.alpha = alpha
42 | self.reduction = loss_fcn.reduction
43 | self.loss_fcn.reduction = 'none' # required to apply FL to each element
44 |
45 | def forward(self, pred, true):
46 | loss = self.loss_fcn(pred, true)
47 | # p_t = torch.exp(-loss)
48 | # loss *= self.alpha * (1.000001 - p_t) ** self.gamma # non-zero power for gradient stability
49 |
50 | # TF implementation https://github.com/tensorflow/addons/blob/v0.7.1/tensorflow_addons/losses/focal_loss.py
51 | pred_prob = torch.sigmoid(pred) # prob from logits
52 | p_t = true * pred_prob + (1 - true) * (1 - pred_prob)
53 | alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha)
54 | modulating_factor = (1.0 - p_t) ** self.gamma
55 | loss *= alpha_factor * modulating_factor
56 |
57 | if self.reduction == 'mean':
58 | return loss.mean()
59 | elif self.reduction == 'sum':
60 | return loss.sum()
61 | else: # 'none'
62 | return loss
63 |
64 |
65 | class QFocalLoss(nn.Module):
66 | # Wraps Quality focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5)
67 | def __init__(self, loss_fcn, gamma=1.5, alpha=0.25):
68 | super().__init__()
69 | self.loss_fcn = loss_fcn # must be nn.BCEWithLogitsLoss()
70 | self.gamma = gamma
71 | self.alpha = alpha
72 | self.reduction = loss_fcn.reduction
73 | self.loss_fcn.reduction = 'none' # required to apply FL to each element
74 |
75 | def forward(self, pred, true):
76 | loss = self.loss_fcn(pred, true)
77 |
78 | pred_prob = torch.sigmoid(pred) # prob from logits
79 | alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha)
80 | modulating_factor = torch.abs(true - pred_prob) ** self.gamma
81 | loss *= alpha_factor * modulating_factor
82 |
83 | if self.reduction == 'mean':
84 | return loss.mean()
85 | elif self.reduction == 'sum':
86 | return loss.sum()
87 | else: # 'none'
88 | return loss
89 |
90 |
91 | class ComputeLoss:
92 | sort_obj_iou = False
93 |
94 | # Compute losses
95 | def __init__(self, model, autobalance=False):
96 | device = next(model.parameters()).device # get model device
97 | h = model.hyp # hyperparameters
98 |
99 | # Define criteria
100 | BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['cls_pw']], device=device))
101 | BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['obj_pw']], device=device))
102 |
103 | # Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3
104 | self.cp, self.cn = smooth_BCE(eps=h.get('label_smoothing', 0.0)) # positive, negative BCE targets
105 |
106 | # Focal loss
107 | g = h['fl_gamma'] # focal loss gamma
108 | if g > 0:
109 | BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)
110 |
111 | m = de_parallel(model).model[-1] # Detect() module
112 | self.balance = {3: [4.0, 1.0, 0.4]}.get(m.nl, [4.0, 1.0, 0.25, 0.06, 0.02]) # P3-P7
113 | self.ssi = list(m.stride).index(16) if autobalance else 0 # stride 16 index
114 | self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = BCEcls, BCEobj, 1.0, h, autobalance
115 | self.na = m.na # number of anchors
116 | self.nc = m.nc # number of classes
117 | self.nl = m.nl # number of layers
118 | self.anchors = m.anchors
119 | self.device = device
120 |
121 | def __call__(self, p, targets): # predictions, targets
122 | lcls = torch.zeros(1, device=self.device) # class loss
123 | lbox = torch.zeros(1, device=self.device) # box loss
124 | lobj = torch.zeros(1, device=self.device) # object loss
125 | tcls, tbox, indices, anchors = self.build_targets(p, targets) # targets
126 |
127 | # Losses
128 | for i, pi in enumerate(p): # layer index, layer predictions
129 | b, a, gj, gi = indices[i] # image, anchor, gridy, gridx
130 | tobj = torch.zeros(pi.shape[:4], dtype=pi.dtype, device=self.device) # target obj
131 |
132 | n = b.shape[0] # number of targets
133 | if n:
134 | # pxy, pwh, _, pcls = pi[b, a, gj, gi].tensor_split((2, 4, 5), dim=1) # faster, requires torch 1.8.0
135 | pxy, pwh, _, pcls = pi[b, a, gj, gi].split((2, 2, 1, self.nc), 1) # target-subset of predictions
136 |
137 | # Regression
138 | pxy = pxy.sigmoid() * 2 - 0.5
139 | pwh = (pwh.sigmoid() * 2) ** 2 * anchors[i]
140 | pbox = torch.cat((pxy, pwh), 1) # predicted box
141 | iou = bbox_iou(pbox, tbox[i], CIoU=True).squeeze() # iou(prediction, target)
142 | lbox += (1.0 - iou).mean() # iou loss
143 |
144 | # Objectness
145 | iou = iou.detach().clamp(0).type(tobj.dtype)
146 | if self.sort_obj_iou:
147 | j = iou.argsort()
148 | b, a, gj, gi, iou = b[j], a[j], gj[j], gi[j], iou[j]
149 | if self.gr < 1:
150 | iou = (1.0 - self.gr) + self.gr * iou
151 | tobj[b, a, gj, gi] = iou # iou ratio
152 |
153 | # Classification
154 | if self.nc > 1: # cls loss (only if multiple classes)
155 | t = torch.full_like(pcls, self.cn, device=self.device) # targets
156 | t[range(n), tcls[i]] = self.cp
157 | lcls += self.BCEcls(pcls, t) # BCE
158 |
159 | # Append targets to text file
160 | # with open('targets.txt', 'a') as file:
161 | # [file.write('%11.5g ' * 4 % tuple(x) + '\n') for x in torch.cat((txy[i], twh[i]), 1)]
162 |
163 | obji = self.BCEobj(pi[..., 4], tobj)
164 | lobj += obji * self.balance[i] # obj loss
165 | if self.autobalance:
166 | self.balance[i] = self.balance[i] * 0.9999 + 0.0001 / obji.detach().item()
167 |
168 | if self.autobalance:
169 | self.balance = [x / self.balance[self.ssi] for x in self.balance]
170 | lbox *= self.hyp['box']
171 | lobj *= self.hyp['obj']
172 | lcls *= self.hyp['cls']
173 | bs = tobj.shape[0] # batch size
174 |
175 | return (lbox + lobj + lcls) * bs, torch.cat((lbox, lobj, lcls)).detach()
176 |
177 | def build_targets(self, p, targets):
178 | # Build targets for compute_loss(), input targets(image,class,x,y,w,h)
179 | na, nt = self.na, targets.shape[0] # number of anchors, targets
180 | tcls, tbox, indices, anch = [], [], [], []
181 | gain = torch.ones(7, device=self.device) # normalized to gridspace gain
182 | ai = torch.arange(na, device=self.device).float().view(na, 1).repeat(1, nt) # same as .repeat_interleave(nt)
183 | targets = torch.cat((targets.repeat(na, 1, 1), ai[..., None]), 2) # append anchor indices
184 |
185 | g = 0.5 # bias
186 | off = torch.tensor(
187 | [
188 | [0, 0],
189 | [1, 0],
190 | [0, 1],
191 | [-1, 0],
192 | [0, -1], # j,k,l,m
193 | # [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm
194 | ],
195 | device=self.device).float() * g # offsets
196 |
197 | for i in range(self.nl):
198 | anchors, shape = self.anchors[i], p[i].shape
199 | gain[2:6] = torch.tensor(shape)[[3, 2, 3, 2]] # xyxy gain
200 |
201 | # Match targets to anchors
202 | t = targets * gain # shape(3,n,7)
203 | if nt:
204 | # Matches
205 | r = t[..., 4:6] / anchors[:, None] # wh ratio
206 | j = torch.max(r, 1 / r).max(2)[0] < self.hyp['anchor_t'] # compare
207 | # j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t'] # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2))
208 | t = t[j] # filter
209 |
210 | # Offsets
211 | gxy = t[:, 2:4] # grid xy
212 | gxi = gain[[2, 3]] - gxy # inverse
213 | j, k = ((gxy % 1 < g) & (gxy > 1)).T
214 | l, m = ((gxi % 1 < g) & (gxi > 1)).T
215 | j = torch.stack((torch.ones_like(j), j, k, l, m))
216 | t = t.repeat((5, 1, 1))[j]
217 | offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j]
218 | else:
219 | t = targets[0]
220 | offsets = 0
221 |
222 | # Define
223 | bc, gxy, gwh, a = t.chunk(4, 1) # (image, class), grid xy, grid wh, anchors
224 | a, (b, c) = a.long().view(-1), bc.long().T # anchors, image, class
225 | gij = (gxy - offsets).long()
226 | gi, gj = gij.T # grid indices
227 |
228 | # Append
229 | indices.append((b, a, gj.clamp_(0, shape[2] - 1), gi.clamp_(0, shape[3] - 1))) # image, anchor, grid
230 | tbox.append(torch.cat((gxy - gij, gwh), 1)) # box
231 | anch.append(anchors[a]) # anchors
232 | tcls.append(c) # class
233 |
234 | return tcls, tbox, indices, anch
235 |
--------------------------------------------------------------------------------
/utils/segment/__init__.py:
--------------------------------------------------------------------------------
1 | # init
--------------------------------------------------------------------------------
/utils/segment/augmentations.py:
--------------------------------------------------------------------------------
1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
2 | """
3 | Image augmentation functions
4 | """
5 |
6 | import math
7 | import random
8 |
9 | import cv2
10 | import numpy as np
11 |
12 | from ..augmentations import box_candidates
13 | from ..general import resample_segments, segment2box
14 |
15 |
16 | def mixup(im, labels, segments, im2, labels2, segments2):
17 | # Applies MixUp augmentation https://arxiv.org/pdf/1710.09412.pdf
18 | r = np.random.beta(32.0, 32.0) # mixup ratio, alpha=beta=32.0
19 | im = (im * r + im2 * (1 - r)).astype(np.uint8)
20 | labels = np.concatenate((labels, labels2), 0)
21 | segments = np.concatenate((segments, segments2), 0)
22 | return im, labels, segments
23 |
24 |
25 | def random_perspective(im,
26 | targets=(),
27 | segments=(),
28 | degrees=10,
29 | translate=.1,
30 | scale=.1,
31 | shear=10,
32 | perspective=0.0,
33 | border=(0, 0)):
34 | # torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(.1, .1), scale=(.9, 1.1), shear=(-10, 10))
35 | # targets = [cls, xyxy]
36 |
37 | height = im.shape[0] + border[0] * 2 # shape(h,w,c)
38 | width = im.shape[1] + border[1] * 2
39 |
40 | # Center
41 | C = np.eye(3)
42 | C[0, 2] = -im.shape[1] / 2 # x translation (pixels)
43 | C[1, 2] = -im.shape[0] / 2 # y translation (pixels)
44 |
45 | # Perspective
46 | P = np.eye(3)
47 | P[2, 0] = random.uniform(-perspective, perspective) # x perspective (about y)
48 | P[2, 1] = random.uniform(-perspective, perspective) # y perspective (about x)
49 |
50 | # Rotation and Scale
51 | R = np.eye(3)
52 | a = random.uniform(-degrees, degrees)
53 | # a += random.choice([-180, -90, 0, 90]) # add 90deg rotations to small rotations
54 | s = random.uniform(1 - scale, 1 + scale)
55 | # s = 2 ** random.uniform(-scale, scale)
56 | R[:2] = cv2.getRotationMatrix2D(angle=a, center=(0, 0), scale=s)
57 |
58 | # Shear
59 | S = np.eye(3)
60 | S[0, 1] = math.tan(random.uniform(-shear, shear) * math.pi / 180) # x shear (deg)
61 | S[1, 0] = math.tan(random.uniform(-shear, shear) * math.pi / 180) # y shear (deg)
62 |
63 | # Translation
64 | T = np.eye(3)
65 | T[0, 2] = (random.uniform(0.5 - translate, 0.5 + translate) * width) # x translation (pixels)
66 | T[1, 2] = (random.uniform(0.5 - translate, 0.5 + translate) * height) # y translation (pixels)
67 |
68 | # Combined rotation matrix
69 | M = T @ S @ R @ P @ C # order of operations (right to left) is IMPORTANT
70 | if (border[0] != 0) or (border[1] != 0) or (M != np.eye(3)).any(): # image changed
71 | if perspective:
72 | im = cv2.warpPerspective(im, M, dsize=(width, height), borderValue=(114, 114, 114))
73 | else: # affine
74 | im = cv2.warpAffine(im, M[:2], dsize=(width, height), borderValue=(114, 114, 114))
75 |
76 | # Visualize
77 | # import matplotlib.pyplot as plt
78 | # ax = plt.subplots(1, 2, figsize=(12, 6))[1].ravel()
79 | # ax[0].imshow(im[:, :, ::-1]) # base
80 | # ax[1].imshow(im2[:, :, ::-1]) # warped
81 |
82 | # Transform label coordinates
83 | n = len(targets)
84 | new_segments = []
85 | if n:
86 | new = np.zeros((n, 4))
87 | segments = resample_segments(segments) # upsample
88 | for i, segment in enumerate(segments):
89 | xy = np.ones((len(segment), 3))
90 | xy[:, :2] = segment
91 | xy = xy @ M.T # transform
92 | xy = (xy[:, :2] / xy[:, 2:3] if perspective else xy[:, :2]) # perspective rescale or affine
93 |
94 | # clip
95 | new[i] = segment2box(xy, width, height)
96 | new_segments.append(xy)
97 |
98 | # filter candidates
99 | i = box_candidates(box1=targets[:, 1:5].T * s, box2=new.T, area_thr=0.01)
100 | targets = targets[i]
101 | targets[:, 1:5] = new[i]
102 | new_segments = np.array(new_segments)[i]
103 |
104 | return im, targets, new_segments
105 |
--------------------------------------------------------------------------------
/utils/segment/dataloaders.py:
--------------------------------------------------------------------------------
1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
2 | """
3 | Dataloaders
4 | """
5 |
6 | import os
7 | import random
8 |
9 | import cv2
10 | import numpy as np
11 | import torch
12 | from torch.utils.data import DataLoader, distributed
13 |
14 | from ..augmentations import augment_hsv, copy_paste, letterbox
15 | from ..dataloaders import InfiniteDataLoader, LoadImagesAndLabels, seed_worker
16 | from ..general import LOGGER, xyn2xy, xywhn2xyxy, xyxy2xywhn
17 | from ..torch_utils import torch_distributed_zero_first
18 | from .augmentations import mixup, random_perspective
19 |
20 |
21 | def create_dataloader(path,
22 | imgsz,
23 | batch_size,
24 | stride,
25 | single_cls=False,
26 | hyp=None,
27 | augment=False,
28 | cache=False,
29 | pad=0.0,
30 | rect=False,
31 | rank=-1,
32 | workers=8,
33 | image_weights=False,
34 | quad=False,
35 | prefix='',
36 | shuffle=False,
37 | mask_downsample_ratio=1,
38 | overlap_mask=False):
39 | if rect and shuffle:
40 | LOGGER.warning('WARNING: --rect is incompatible with DataLoader shuffle, setting shuffle=False')
41 | shuffle = False
42 | with torch_distributed_zero_first(rank): # init dataset *.cache only once if DDP
43 | dataset = LoadImagesAndLabelsAndMasks(
44 | path,
45 | imgsz,
46 | batch_size,
47 | augment=augment, # augmentation
48 | hyp=hyp, # hyperparameters
49 | rect=rect, # rectangular batches
50 | cache_images=cache,
51 | single_cls=single_cls,
52 | stride=int(stride),
53 | pad=pad,
54 | image_weights=image_weights,
55 | prefix=prefix,
56 | downsample_ratio=mask_downsample_ratio,
57 | overlap=overlap_mask)
58 |
59 | batch_size = min(batch_size, len(dataset))
60 | nd = torch.cuda.device_count() # number of CUDA devices
61 | nw = min([os.cpu_count() // max(nd, 1), batch_size if batch_size > 1 else 0, workers]) # number of workers
62 | sampler = None if rank == -1 else distributed.DistributedSampler(dataset, shuffle=shuffle)
63 | loader = DataLoader if image_weights else InfiniteDataLoader # only DataLoader allows for attribute updates
64 | # generator = torch.Generator()
65 | # generator.manual_seed(0)
66 | return loader(
67 | dataset,
68 | batch_size=batch_size,
69 | shuffle=shuffle and sampler is None,
70 | num_workers=nw,
71 | sampler=sampler,
72 | pin_memory=True,
73 | collate_fn=LoadImagesAndLabelsAndMasks.collate_fn4 if quad else LoadImagesAndLabelsAndMasks.collate_fn,
74 | worker_init_fn=seed_worker,
75 | # generator=generator,
76 | ), dataset
77 |
78 |
79 | class LoadImagesAndLabelsAndMasks(LoadImagesAndLabels): # for training/testing
80 |
81 | def __init__(
82 | self,
83 | path,
84 | img_size=640,
85 | batch_size=16,
86 | augment=False,
87 | hyp=None,
88 | rect=False,
89 | image_weights=False,
90 | cache_images=False,
91 | single_cls=False,
92 | stride=32,
93 | pad=0,
94 | prefix="",
95 | downsample_ratio=1,
96 | overlap=False,
97 | ):
98 | super().__init__(path, img_size, batch_size, augment, hyp, rect, image_weights, cache_images, single_cls,
99 | stride, pad, prefix)
100 | self.downsample_ratio = downsample_ratio
101 | self.overlap = overlap
102 |
103 | def __getitem__(self, index):
104 | index = self.indices[index] # linear, shuffled, or image_weights
105 |
106 | hyp = self.hyp
107 | mosaic = self.mosaic and random.random() < hyp['mosaic']
108 | masks = []
109 | if mosaic:
110 | # Load mosaic
111 | img, labels, segments = self.load_mosaic(index)
112 | shapes = None
113 |
114 | # MixUp augmentation
115 | if random.random() < hyp["mixup"]:
116 | img, labels, segments = mixup(img, labels, segments, *self.load_mosaic(random.randint(0, self.n - 1)))
117 |
118 | else:
119 | # Load image
120 | img, (h0, w0), (h, w) = self.load_image(index)
121 |
122 | # Letterbox
123 | shape = self.batch_shapes[self.batch[index]] if self.rect else self.img_size # final letterboxed shape
124 | img, ratio, pad = letterbox(img, shape, auto=False, scaleup=self.augment)
125 | shapes = (h0, w0), ((h / h0, w / w0), pad) # for COCO mAP rescaling
126 |
127 | labels = self.labels[index].copy()
128 | # [array, array, ....], array.shape=(num_points, 2), xyxyxyxy
129 | segments = self.segments[index].copy()
130 | if len(segments):
131 | for i_s in range(len(segments)):
132 | segments[i_s] = xyn2xy(
133 | segments[i_s],
134 | ratio[0] * w,
135 | ratio[1] * h,
136 | padw=pad[0],
137 | padh=pad[1],
138 | )
139 | if labels.size: # normalized xywh to pixel xyxy format
140 | labels[:, 1:] = xywhn2xyxy(labels[:, 1:], ratio[0] * w, ratio[1] * h, padw=pad[0], padh=pad[1])
141 |
142 | if self.augment:
143 | img, labels, segments = random_perspective(
144 | img,
145 | labels,
146 | segments=segments,
147 | degrees=hyp["degrees"],
148 | translate=hyp["translate"],
149 | scale=hyp["scale"],
150 | shear=hyp["shear"],
151 | perspective=hyp["perspective"],
152 | return_seg=True,
153 | )
154 |
155 | nl = len(labels) # number of labels
156 | if nl:
157 | labels[:, 1:5] = xyxy2xywhn(labels[:, 1:5], w=img.shape[1], h=img.shape[0], clip=True, eps=1e-3)
158 | if self.overlap:
159 | masks, sorted_idx = polygons2masks_overlap(img.shape[:2],
160 | segments,
161 | downsample_ratio=self.downsample_ratio)
162 | masks = masks[None] # (640, 640) -> (1, 640, 640)
163 | labels = labels[sorted_idx]
164 | else:
165 | masks = polygons2masks(img.shape[:2], segments, color=1, downsample_ratio=self.downsample_ratio)
166 |
167 | masks = (torch.from_numpy(masks) if len(masks) else torch.zeros(1 if self.overlap else nl, img.shape[0] //
168 | self.downsample_ratio, img.shape[1] //
169 | self.downsample_ratio))
170 | # TODO: albumentations support
171 | if self.augment:
172 | # Albumentations
173 | # there are some augmentation that won't change boxes and masks,
174 | # so just be it for now.
175 | img, labels = self.albumentations(img, labels)
176 | nl = len(labels) # update after albumentations
177 |
178 | # HSV color-space
179 | augment_hsv(img, hgain=hyp["hsv_h"], sgain=hyp["hsv_s"], vgain=hyp["hsv_v"])
180 |
181 | # Flip up-down
182 | if random.random() < hyp["flipud"]:
183 | img = np.flipud(img)
184 | if nl:
185 | labels[:, 2] = 1 - labels[:, 2]
186 | masks = torch.flip(masks, dims=[1])
187 |
188 | # Flip left-right
189 | if random.random() < hyp["fliplr"]:
190 | img = np.fliplr(img)
191 | if nl:
192 | labels[:, 1] = 1 - labels[:, 1]
193 | masks = torch.flip(masks, dims=[2])
194 |
195 | # Cutouts # labels = cutout(img, labels, p=0.5)
196 |
197 | labels_out = torch.zeros((nl, 6))
198 | if nl:
199 | labels_out[:, 1:] = torch.from_numpy(labels)
200 |
201 | # Convert
202 | img = img.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB
203 | img = np.ascontiguousarray(img)
204 |
205 | return (torch.from_numpy(img), labels_out, self.im_files[index], shapes, masks)
206 |
207 | def load_mosaic(self, index):
208 | # YOLOv5 4-mosaic loader. Loads 1 image + 3 random images into a 4-image mosaic
209 | labels4, segments4 = [], []
210 | s = self.img_size
211 | yc, xc = (int(random.uniform(-x, 2 * s + x)) for x in self.mosaic_border) # mosaic center x, y
212 |
213 | # 3 additional image indices
214 | indices = [index] + random.choices(self.indices, k=3) # 3 additional image indices
215 | for i, index in enumerate(indices):
216 | # Load image
217 | img, _, (h, w) = self.load_image(index)
218 |
219 | # place img in img4
220 | if i == 0: # top left
221 | img4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8) # base image with 4 tiles
222 | x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc # xmin, ymin, xmax, ymax (large image)
223 | x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h # xmin, ymin, xmax, ymax (small image)
224 | elif i == 1: # top right
225 | x1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), yc
226 | x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), h
227 | elif i == 2: # bottom left
228 | x1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(s * 2, yc + h)
229 | x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, w, min(y2a - y1a, h)
230 | elif i == 3: # bottom right
231 | x1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(s * 2, yc + h)
232 | x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h)
233 |
234 | img4[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b] # img4[ymin:ymax, xmin:xmax]
235 | padw = x1a - x1b
236 | padh = y1a - y1b
237 |
238 | labels, segments = self.labels[index].copy(), self.segments[index].copy()
239 |
240 | if labels.size:
241 | labels[:, 1:] = xywhn2xyxy(labels[:, 1:], w, h, padw, padh) # normalized xywh to pixel xyxy format
242 | segments = [xyn2xy(x, w, h, padw, padh) for x in segments]
243 | labels4.append(labels)
244 | segments4.extend(segments)
245 |
246 | # Concat/clip labels
247 | labels4 = np.concatenate(labels4, 0)
248 | for x in (labels4[:, 1:], *segments4):
249 | np.clip(x, 0, 2 * s, out=x) # clip when using random_perspective()
250 | # img4, labels4 = replicate(img4, labels4) # replicate
251 |
252 | # Augment
253 | img4, labels4, segments4 = copy_paste(img4, labels4, segments4, p=self.hyp["copy_paste"])
254 | img4, labels4, segments4 = random_perspective(img4,
255 | labels4,
256 | segments4,
257 | degrees=self.hyp["degrees"],
258 | translate=self.hyp["translate"],
259 | scale=self.hyp["scale"],
260 | shear=self.hyp["shear"],
261 | perspective=self.hyp["perspective"],
262 | border=self.mosaic_border) # border to remove
263 | return img4, labels4, segments4
264 |
265 | @staticmethod
266 | def collate_fn(batch):
267 | img, label, path, shapes, masks = zip(*batch) # transposed
268 | batched_masks = torch.cat(masks, 0)
269 | for i, l in enumerate(label):
270 | l[:, 0] = i # add target image index for build_targets()
271 | return torch.stack(img, 0), torch.cat(label, 0), path, shapes, batched_masks
272 |
273 |
274 | def polygon2mask(img_size, polygons, color=1, downsample_ratio=1):
275 | """
276 | Args:
277 | img_size (tuple): The image size.
278 | polygons (np.ndarray): [N, M], N is the number of polygons,
279 | M is the number of points(Be divided by 2).
280 | """
281 | mask = np.zeros(img_size, dtype=np.uint8)
282 | polygons = np.asarray(polygons)
283 | polygons = polygons.astype(np.int32)
284 | shape = polygons.shape
285 | polygons = polygons.reshape(shape[0], -1, 2)
286 | cv2.fillPoly(mask, polygons, color=color)
287 | nh, nw = (img_size[0] // downsample_ratio, img_size[1] // downsample_ratio)
288 | # NOTE: fillPoly firstly then resize is trying the keep the same way
289 | # of loss calculation when mask-ratio=1.
290 | mask = cv2.resize(mask, (nw, nh))
291 | return mask
292 |
293 |
294 | def polygons2masks(img_size, polygons, color, downsample_ratio=1):
295 | """
296 | Args:
297 | img_size (tuple): The image size.
298 | polygons (list[np.ndarray]): each polygon is [N, M],
299 | N is the number of polygons,
300 | M is the number of points(Be divided by 2).
301 | """
302 | masks = []
303 | for si in range(len(polygons)):
304 | mask = polygon2mask(img_size, [polygons[si].reshape(-1)], color, downsample_ratio)
305 | masks.append(mask)
306 | return np.array(masks)
307 |
308 |
309 | def polygons2masks_overlap(img_size, segments, downsample_ratio=1):
310 | """Return a (640, 640) overlap mask."""
311 | masks = np.zeros((img_size[0] // downsample_ratio, img_size[1] // downsample_ratio), dtype=np.uint8)
312 | areas = []
313 | ms = []
314 | for si in range(len(segments)):
315 | mask = polygon2mask(
316 | img_size,
317 | [segments[si].reshape(-1)],
318 | downsample_ratio=downsample_ratio,
319 | color=1,
320 | )
321 | ms.append(mask)
322 | areas.append(mask.sum())
323 | areas = np.asarray(areas)
324 | index = np.argsort(-areas)
325 | ms = np.array(ms)[index]
326 | for i in range(len(segments)):
327 | mask = ms[i] * (i + 1)
328 | masks = masks + mask
329 | masks = np.clip(masks, a_min=0, a_max=i + 1)
330 | return masks, index
331 |
--------------------------------------------------------------------------------
/utils/segment/general.py:
--------------------------------------------------------------------------------
1 | import cv2
2 | import torch
3 | import torch.nn.functional as F
4 | import numpy as np
5 |
6 | def crop(masks, boxes):
7 | """
8 | "Crop" predicted masks by zeroing out everything not in the predicted bbox.
9 | Vectorized by Chong (thanks Chong).
10 |
11 | Args:
12 | - masks should be a size [h, w, n] tensor of masks
13 | - boxes should be a size [n, 4] tensor of bbox coords in relative point form
14 | """
15 |
16 | n, h, w = masks.shape
17 | x1, y1, x2, y2 = torch.chunk(boxes[:, :, None], 4, 1) # x1 shape(1,1,n)
18 | r = torch.arange(w, device=masks.device, dtype=x1.dtype)[None, None, :] # rows shape(1,w,1)
19 | c = torch.arange(h, device=masks.device, dtype=x1.dtype)[None, :, None] # cols shape(h,1,1)
20 |
21 | return masks * ((r >= x1) * (r < x2) * (c >= y1) * (c < y2))
22 |
23 |
24 | def process_mask_upsample(protos, masks_in, bboxes, shape):
25 | """
26 | Crop after upsample.
27 | proto_out: [mask_dim, mask_h, mask_w]
28 | out_masks: [n, mask_dim], n is number of masks after nms
29 | bboxes: [n, 4], n is number of masks after nms
30 | shape:input_image_size, (h, w)
31 |
32 | return: h, w, n
33 | """
34 |
35 | c, mh, mw = protos.shape # CHW
36 | masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw)
37 | masks = F.interpolate(masks[None], shape, mode='bilinear', align_corners=False)[0] # CHW
38 | masks = crop(masks, bboxes) # CHW
39 | return masks.gt_(0.5)
40 |
41 |
42 | def process_mask(protos, masks_in, bboxes, shape, upsample=False):
43 | """
44 | Crop before upsample.
45 | proto_out: [mask_dim, mask_h, mask_w]
46 | out_masks: [n, mask_dim], n is number of masks after nms
47 | bboxes: [n, 4], n is number of masks after nms
48 | shape:input_image_size, (h, w)
49 |
50 | return: h, w, n
51 | """
52 |
53 | c, mh, mw = protos.shape # CHW
54 | ih, iw = shape
55 | masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw) # CHW
56 |
57 | downsampled_bboxes = bboxes.clone()
58 | downsampled_bboxes[:, 0] *= mw / iw
59 | downsampled_bboxes[:, 2] *= mw / iw
60 | downsampled_bboxes[:, 3] *= mh / ih
61 | downsampled_bboxes[:, 1] *= mh / ih
62 |
63 | masks = crop(masks, downsampled_bboxes) # CHW
64 | if upsample:
65 | masks = F.interpolate(masks[None], shape, mode='bilinear', align_corners=False)[0] # CHW
66 | return masks.gt_(0.5)
67 |
68 |
69 | def scale_masks(img1_shape, masks, img0_shape, ratio_pad=None):
70 | """
71 | img1_shape: model input shape, [h, w]
72 | img0_shape: origin pic shape, [h, w, 3]
73 | masks: [h, w, num]
74 | resize for the most time
75 | """
76 | # Rescale coords (xyxy) from img1_shape to img0_shape
77 | if ratio_pad is None: # calculate from img0_shape
78 | gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain = old / new
79 | pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding
80 | else:
81 | gain = ratio_pad[0][0]
82 | pad = ratio_pad[1]
83 | tl_pad = int(pad[1]), int(pad[0]) # y, x
84 | br_pad = int(img1_shape[0] - pad[1]), int(img1_shape[1] - pad[0])
85 |
86 | if len(masks.shape) < 2:
87 | raise ValueError(f'"len of masks shape" should be 2 or 3, but got {len(masks.shape)}')
88 | # masks_h, masks_w, n
89 | masks = masks[tl_pad[0]:br_pad[0], tl_pad[1]:br_pad[1]]
90 | # 1, n, masks_h, masks_w
91 | # masks = masks.permute(2, 0, 1).contiguous()[None, :]
92 | # # shape = [1, n, masks_h, masks_w] after F.interpolate, so take first element
93 | # masks = F.interpolate(masks, img0_shape[:2], mode='bilinear', align_corners=False)[0]
94 | # masks = masks.permute(1, 2, 0).contiguous()
95 | # masks_h, masks_w, n
96 | masks = cv2.resize(masks, (img0_shape[1], img0_shape[0]))
97 |
98 | # keepdim
99 | if len(masks.shape) == 2:
100 | masks = masks[:, :, None]
101 |
102 | return masks
103 |
104 |
105 | def mask_iou(mask1, mask2, eps=1e-7):
106 | """
107 | mask1: [N, n] m1 means number of predicted objects
108 | mask2: [M, n] m2 means number of gt objects
109 | Note: n means image_w x image_h
110 |
111 | return: masks iou, [N, M]
112 | """
113 | intersection = torch.matmul(mask1, mask2.t()).clamp(0)
114 | union = (mask1.sum(1)[:, None] + mask2.sum(1)[None]) - intersection # (area1 + area2) - intersection
115 | return intersection / (union + eps)
116 |
117 |
118 | def masks_iou(mask1, mask2, eps=1e-7):
119 | """
120 | mask1: [N, n] m1 means number of predicted objects
121 | mask2: [N, n] m2 means number of gt objects
122 | Note: n means image_w x image_h
123 |
124 | return: masks iou, (N, )
125 | """
126 | intersection = (mask1 * mask2).sum(1).clamp(0) # (N, )
127 | union = (mask1.sum(1) + mask2.sum(1))[None] - intersection # (area1 + area2) - intersection
128 | return intersection / (union + eps)
129 |
130 |
131 | def masks2segments(masks, strategy='largest'):
132 | # Convert masks(n,160,160) into segments(n,xy)
133 | segments = []
134 | for x in masks.int().cpu().numpy().astype('uint8'):
135 | c = cv2.findContours(x, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
136 | if c:
137 | if strategy == 'concat': # concatenate all segments
138 | c = np.concatenate([x.reshape(-1, 2) for x in c])
139 | elif strategy == 'largest': # select largest segment
140 | c = np.array(c[np.array([len(x) for x in c]).argmax()]).reshape(-1, 2)
141 | else:
142 | c = np.zeros((0, 2)) # no segments found
143 | segments.append(c.astype('float32'))
144 | return segments
--------------------------------------------------------------------------------
/utils/segment/loss.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import torch.nn as nn
3 | import torch.nn.functional as F
4 |
5 | from ..general import xywh2xyxy
6 | from ..loss import FocalLoss, smooth_BCE
7 | from ..metrics import bbox_iou
8 | from ..torch_utils import de_parallel
9 | from .general import crop
10 |
11 |
12 | class ComputeLoss:
13 | # Compute losses
14 | def __init__(self, model, autobalance=False, overlap=False):
15 | self.sort_obj_iou = False
16 | self.overlap = overlap
17 | device = next(model.parameters()).device # get model device
18 | h = model.hyp # hyperparameters
19 | self.device = device
20 |
21 | # Define criteria
22 | BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['cls_pw']], device=device))
23 | BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['obj_pw']], device=device))
24 |
25 | # Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3
26 | self.cp, self.cn = smooth_BCE(eps=h.get('label_smoothing', 0.0)) # positive, negative BCE targets
27 |
28 | # Focal loss
29 | g = h['fl_gamma'] # focal loss gamma
30 | if g > 0:
31 | BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)
32 |
33 | m = de_parallel(model).model[-1] # Detect() module
34 | self.balance = {3: [4.0, 1.0, 0.4]}.get(m.nl, [4.0, 1.0, 0.25, 0.06, 0.02]) # P3-P7
35 | self.ssi = list(m.stride).index(16) if autobalance else 0 # stride 16 index
36 | self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = BCEcls, BCEobj, 1.0, h, autobalance
37 | self.na = m.na # number of anchors
38 | self.nc = m.nc # number of classes
39 | self.nl = m.nl # number of layers
40 | self.nm = m.nm # number of masks
41 | self.anchors = m.anchors
42 | self.device = device
43 |
44 | def __call__(self, preds, targets, masks): # predictions, targets, model
45 | p, proto = preds
46 | bs, nm, mask_h, mask_w = proto.shape # batch size, number of masks, mask height, mask width
47 | lcls = torch.zeros(1, device=self.device)
48 | lbox = torch.zeros(1, device=self.device)
49 | lobj = torch.zeros(1, device=self.device)
50 | lseg = torch.zeros(1, device=self.device)
51 | tcls, tbox, indices, anchors, tidxs, xywhn = self.build_targets(p, targets) # targets
52 |
53 | # Losses
54 | for i, pi in enumerate(p): # layer index, layer predictions
55 | b, a, gj, gi = indices[i] # image, anchor, gridy, gridx
56 | tobj = torch.zeros(pi.shape[:4], dtype=pi.dtype, device=self.device) # target obj
57 |
58 | n = b.shape[0] # number of targets
59 | if n:
60 | pxy, pwh, _, pcls, pmask = pi[b, a, gj, gi].split((2, 2, 1, self.nc, nm), 1) # subset of predictions
61 |
62 | # Box regression
63 | pxy = pxy.sigmoid() * 2 - 0.5
64 | pwh = (pwh.sigmoid() * 2) ** 2 * anchors[i]
65 | pbox = torch.cat((pxy, pwh), 1) # predicted box
66 | iou = bbox_iou(pbox, tbox[i], CIoU=True).squeeze() # iou(prediction, target)
67 | lbox += (1.0 - iou).mean() # iou loss
68 |
69 | # Objectness
70 | iou = iou.detach().clamp(0).type(tobj.dtype)
71 | if self.sort_obj_iou:
72 | j = iou.argsort()
73 | b, a, gj, gi, iou = b[j], a[j], gj[j], gi[j], iou[j]
74 | if self.gr < 1:
75 | iou = (1.0 - self.gr) + self.gr * iou
76 | tobj[b, a, gj, gi] = iou # iou ratio
77 |
78 | # Classification
79 | if self.nc > 1: # cls loss (only if multiple classes)
80 | t = torch.full_like(pcls, self.cn, device=self.device) # targets
81 | t[range(n), tcls[i]] = self.cp
82 | lcls += self.BCEcls(pcls, t) # BCE
83 |
84 | # Mask regression
85 | if tuple(masks.shape[-2:]) != (mask_h, mask_w): # downsample
86 | masks = F.interpolate(masks[None], (mask_h, mask_w), mode="bilinear", align_corners=False)[0]
87 | marea = xywhn[i][:, 2:].prod(1) # mask width, height normalized
88 | mxyxy = xywh2xyxy(xywhn[i] * torch.tensor([mask_w, mask_h, mask_w, mask_h], device=self.device))
89 | for bi in b.unique():
90 | j = b == bi # matching index
91 | if self.overlap:
92 | mask_gti = torch.where(masks[bi][None] == tidxs[i][j].view(-1, 1, 1), 1.0, 0.0)
93 | else:
94 | mask_gti = masks[tidxs[i]][j]
95 | lseg += self.single_mask_loss(mask_gti, pmask[j], proto[bi], mxyxy[j], marea[j])
96 |
97 | obji = self.BCEobj(pi[..., 4], tobj)
98 | lobj += obji * self.balance[i] # obj loss
99 | if self.autobalance:
100 | self.balance[i] = self.balance[i] * 0.9999 + 0.0001 / obji.detach().item()
101 |
102 | if self.autobalance:
103 | self.balance = [x / self.balance[self.ssi] for x in self.balance]
104 | lbox *= self.hyp["box"]
105 | lobj *= self.hyp["obj"]
106 | lcls *= self.hyp["cls"]
107 | lseg *= self.hyp["box"] / bs
108 |
109 | loss = lbox + lobj + lcls + lseg
110 | return loss * bs, torch.cat((lbox, lseg, lobj, lcls)).detach()
111 |
112 | def single_mask_loss(self, gt_mask, pred, proto, xyxy, area):
113 | # Mask loss for one image
114 | pred_mask = (pred @ proto.view(self.nm, -1)).view(-1, *proto.shape[1:]) # (n,32) @ (32,80,80) -> (n,80,80)
115 | loss = F.binary_cross_entropy_with_logits(pred_mask, gt_mask, reduction="none")
116 | return (crop(loss, xyxy).mean(dim=(1, 2)) / area).mean()
117 |
118 | def build_targets(self, p, targets):
119 | # Build targets for compute_loss(), input targets(image,class,x,y,w,h)
120 | na, nt = self.na, targets.shape[0] # number of anchors, targets
121 | tcls, tbox, indices, anch, tidxs, xywhn = [], [], [], [], [], []
122 | gain = torch.ones(8, device=self.device) # normalized to gridspace gain
123 | ai = torch.arange(na, device=self.device).float().view(na, 1).repeat(1, nt) # same as .repeat_interleave(nt)
124 | if self.overlap:
125 | batch = p[0].shape[0]
126 | ti = []
127 | for i in range(batch):
128 | num = (targets[:, 0] == i).sum() # find number of targets of each image
129 | ti.append(torch.arange(num, device=self.device).float().view(1, num).repeat(na, 1) + 1) # (na, num)
130 | ti = torch.cat(ti, 1) # (na, nt)
131 | else:
132 | ti = torch.arange(nt, device=self.device).float().view(1, nt).repeat(na, 1)
133 | targets = torch.cat((targets.repeat(na, 1, 1), ai[..., None], ti[..., None]), 2) # append anchor indices
134 |
135 | g = 0.5 # bias
136 | off = torch.tensor(
137 | [
138 | [0, 0],
139 | [1, 0],
140 | [0, 1],
141 | [-1, 0],
142 | [0, -1], # j,k,l,m
143 | # [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm
144 | ],
145 | device=self.device).float() * g # offsets
146 |
147 | for i in range(self.nl):
148 | anchors, shape = self.anchors[i], p[i].shape
149 | gain[2:6] = torch.tensor(shape)[[3, 2, 3, 2]] # xyxy gain
150 |
151 | # Match targets to anchors
152 | t = targets * gain # shape(3,n,7)
153 | if nt:
154 | # Matches
155 | r = t[..., 4:6] / anchors[:, None] # wh ratio
156 | j = torch.max(r, 1 / r).max(2)[0] < self.hyp['anchor_t'] # compare
157 | # j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t'] # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2))
158 | t = t[j] # filter
159 |
160 | # Offsets
161 | gxy = t[:, 2:4] # grid xy
162 | gxi = gain[[2, 3]] - gxy # inverse
163 | j, k = ((gxy % 1 < g) & (gxy > 1)).T
164 | l, m = ((gxi % 1 < g) & (gxi > 1)).T
165 | j = torch.stack((torch.ones_like(j), j, k, l, m))
166 | t = t.repeat((5, 1, 1))[j]
167 | offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j]
168 | else:
169 | t = targets[0]
170 | offsets = 0
171 |
172 | # Define
173 | bc, gxy, gwh, at = t.chunk(4, 1) # (image, class), grid xy, grid wh, anchors
174 | (a, tidx), (b, c) = at.long().T, bc.long().T # anchors, image, class
175 | gij = (gxy - offsets).long()
176 | gi, gj = gij.T # grid indices
177 |
178 | # Append
179 | indices.append((b, a, gj.clamp_(0, shape[2] - 1), gi.clamp_(0, shape[3] - 1))) # image, anchor, grid
180 | tbox.append(torch.cat((gxy - gij, gwh), 1)) # box
181 | anch.append(anchors[a]) # anchors
182 | tcls.append(c) # class
183 | tidxs.append(tidx)
184 | xywhn.append(torch.cat((gxy, gwh), 1) / gain[2:6]) # xywh normalized
185 |
186 | return tcls, tbox, indices, anch, tidxs, xywhn
187 |
--------------------------------------------------------------------------------
/utils/segment/metrics.py:
--------------------------------------------------------------------------------
1 | # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
2 | """
3 | Model validation metrics
4 | """
5 |
6 | import numpy as np
7 |
8 | from ..metrics import ap_per_class
9 |
10 |
11 | def fitness(x):
12 | # Model fitness as a weighted combination of metrics
13 | w = [0.0, 0.0, 0.1, 0.9, 0.0, 0.0, 0.1, 0.9]
14 | return (x[:, :8] * w).sum(1)
15 |
16 |
17 | def ap_per_class_box_and_mask(
18 | tp_m,
19 | tp_b,
20 | conf,
21 | pred_cls,
22 | target_cls,
23 | plot=False,
24 | save_dir=".",
25 | names=(),
26 | ):
27 | """
28 | Args:
29 | tp_b: tp of boxes.
30 | tp_m: tp of masks.
31 | other arguments see `func: ap_per_class`.
32 | """
33 | results_boxes = ap_per_class(tp_b,
34 | conf,
35 | pred_cls,
36 | target_cls,
37 | plot=plot,
38 | save_dir=save_dir,
39 | names=names,
40 | prefix="Box")[2:]
41 | results_masks = ap_per_class(tp_m,
42 | conf,
43 | pred_cls,
44 | target_cls,
45 | plot=plot,
46 | save_dir=save_dir,
47 | names=names,
48 | prefix="Mask")[2:]
49 |
50 | results = {
51 | "boxes": {
52 | "p": results_boxes[0],
53 | "r": results_boxes[1],
54 | "ap": results_boxes[3],
55 | "f1": results_boxes[2],
56 | "ap_class": results_boxes[4]},
57 | "masks": {
58 | "p": results_masks[0],
59 | "r": results_masks[1],
60 | "ap": results_masks[3],
61 | "f1": results_masks[2],
62 | "ap_class": results_masks[4]}}
63 | return results
64 |
65 |
66 | class Metric:
67 |
68 | def __init__(self) -> None:
69 | self.p = [] # (nc, )
70 | self.r = [] # (nc, )
71 | self.f1 = [] # (nc, )
72 | self.all_ap = [] # (nc, 10)
73 | self.ap_class_index = [] # (nc, )
74 |
75 | @property
76 | def ap50(self):
77 | """AP@0.5 of all classes.
78 | Return:
79 | (nc, ) or [].
80 | """
81 | return self.all_ap[:, 0] if len(self.all_ap) else []
82 |
83 | @property
84 | def ap(self):
85 | """AP@0.5:0.95
86 | Return:
87 | (nc, ) or [].
88 | """
89 | return self.all_ap.mean(1) if len(self.all_ap) else []
90 |
91 | @property
92 | def mp(self):
93 | """mean precision of all classes.
94 | Return:
95 | float.
96 | """
97 | return self.p.mean() if len(self.p) else 0.0
98 |
99 | @property
100 | def mr(self):
101 | """mean recall of all classes.
102 | Return:
103 | float.
104 | """
105 | return self.r.mean() if len(self.r) else 0.0
106 |
107 | @property
108 | def map50(self):
109 | """Mean AP@0.5 of all classes.
110 | Return:
111 | float.
112 | """
113 | return self.all_ap[:, 0].mean() if len(self.all_ap) else 0.0
114 |
115 | @property
116 | def map(self):
117 | """Mean AP@0.5:0.95 of all classes.
118 | Return:
119 | float.
120 | """
121 | return self.all_ap.mean() if len(self.all_ap) else 0.0
122 |
123 | def mean_results(self):
124 | """Mean of results, return mp, mr, map50, map"""
125 | return (self.mp, self.mr, self.map50, self.map)
126 |
127 | def class_result(self, i):
128 | """class-aware result, return p[i], r[i], ap50[i], ap[i]"""
129 | return (self.p[i], self.r[i], self.ap50[i], self.ap[i])
130 |
131 | def get_maps(self, nc):
132 | maps = np.zeros(nc) + self.map
133 | for i, c in enumerate(self.ap_class_index):
134 | maps[c] = self.ap[i]
135 | return maps
136 |
137 | def update(self, results):
138 | """
139 | Args:
140 | results: tuple(p, r, ap, f1, ap_class)
141 | """
142 | p, r, all_ap, f1, ap_class_index = results
143 | self.p = p
144 | self.r = r
145 | self.all_ap = all_ap
146 | self.f1 = f1
147 | self.ap_class_index = ap_class_index
148 |
149 |
150 | class Metrics:
151 | """Metric for boxes and masks."""
152 |
153 | def __init__(self) -> None:
154 | self.metric_box = Metric()
155 | self.metric_mask = Metric()
156 |
157 | def update(self, results):
158 | """
159 | Args:
160 | results: Dict{'boxes': Dict{}, 'masks': Dict{}}
161 | """
162 | self.metric_box.update(list(results["boxes"].values()))
163 | self.metric_mask.update(list(results["masks"].values()))
164 |
165 | def mean_results(self):
166 | return self.metric_box.mean_results() + self.metric_mask.mean_results()
167 |
168 | def class_result(self, i):
169 | return self.metric_box.class_result(i) + self.metric_mask.class_result(i)
170 |
171 | def get_maps(self, nc):
172 | return self.metric_box.get_maps(nc) + self.metric_mask.get_maps(nc)
173 |
174 | @property
175 | def ap_class_index(self):
176 | # boxes and masks have the same ap_class_index
177 | return self.metric_box.ap_class_index
178 |
179 |
180 | KEYS = [
181 | "train/box_loss",
182 | "train/seg_loss", # train loss
183 | "train/obj_loss",
184 | "train/cls_loss",
185 | "metrics/precision(B)",
186 | "metrics/recall(B)",
187 | "metrics/mAP_0.5(B)",
188 | "metrics/mAP_0.5:0.95(B)", # metrics
189 | "metrics/precision(M)",
190 | "metrics/recall(M)",
191 | "metrics/mAP_0.5(M)",
192 | "metrics/mAP_0.5:0.95(M)", # metrics
193 | "val/box_loss",
194 | "val/seg_loss", # val loss
195 | "val/obj_loss",
196 | "val/cls_loss",
197 | "x/lr0",
198 | "x/lr1",
199 | "x/lr2",]
200 |
201 | BEST_KEYS = [
202 | "best/epoch",
203 | "best/precision(B)",
204 | "best/recall(B)",
205 | "best/mAP_0.5(B)",
206 | "best/mAP_0.5:0.95(B)",
207 | "best/precision(M)",
208 | "best/recall(M)",
209 | "best/mAP_0.5(M)",
210 | "best/mAP_0.5:0.95(M)",]
211 |
--------------------------------------------------------------------------------
/utils/segment/plots.py:
--------------------------------------------------------------------------------
1 | import contextlib
2 | import math
3 | from pathlib import Path
4 |
5 | import cv2
6 | import matplotlib.pyplot as plt
7 | import numpy as np
8 | import pandas as pd
9 | import torch
10 |
11 | from .. import threaded
12 | from ..general import xywh2xyxy
13 | from ..plots import Annotator, colors
14 |
15 |
16 | def plot_masks(img, masks, colors, alpha=0.5):
17 | """
18 | Args:
19 | img (tensor): img is in cuda, shape: [3, h, w], range: [0, 1]
20 | masks (tensor): predicted masks on cuda, shape: [n, h, w]
21 | colors (List[List[Int]]): colors for predicted masks, [[r, g, b] * n]
22 | Return:
23 | ndarray: img after draw masks, shape: [h, w, 3]
24 |
25 | transform colors and send img_gpu to cpu for the most time.
26 | """
27 | img_gpu = img.clone()
28 | num_masks = len(masks)
29 | if num_masks == 0:
30 | return img.permute(1, 2, 0).contiguous().cpu().numpy() * 255
31 |
32 | # [n, 1, 1, 3]
33 | # faster this way to transform colors
34 | colors = torch.tensor(colors, device=img.device).float() / 255.0
35 | colors = colors[:, None, None, :]
36 | # [n, h, w, 1]
37 | masks = masks[:, :, :, None]
38 | masks_color = masks.repeat(1, 1, 1, 3) * colors * alpha
39 | inv_alph_masks = masks * (-alpha) + 1
40 | masks_color_summand = masks_color[0]
41 | if num_masks > 1:
42 | inv_alph_cumul = inv_alph_masks[:(num_masks - 1)].cumprod(dim=0)
43 | masks_color_cumul = masks_color[1:] * inv_alph_cumul
44 | masks_color_summand += masks_color_cumul.sum(dim=0)
45 |
46 | # print(inv_alph_masks.prod(dim=0).shape) # [h, w, 1]
47 | img_gpu = img_gpu.flip(dims=[0]) # filp channel for opencv
48 | img_gpu = img_gpu.permute(1, 2, 0).contiguous()
49 | # [h, w, 3]
50 | img_gpu = img_gpu * inv_alph_masks.prod(dim=0) + masks_color_summand
51 | return (img_gpu * 255).byte().cpu().numpy()
52 |
53 |
54 | @threaded
55 | def plot_images_and_masks(images, targets, masks, paths=None, fname='images.jpg', names=None):
56 | # Plot image grid with labels
57 | if isinstance(images, torch.Tensor):
58 | images = images.cpu().float().numpy()
59 | if isinstance(targets, torch.Tensor):
60 | targets = targets.cpu().numpy()
61 | if isinstance(masks, torch.Tensor):
62 | masks = masks.cpu().numpy().astype(int)
63 |
64 | max_size = 1920 # max image size
65 | max_subplots = 16 # max image subplots, i.e. 4x4
66 | bs, _, h, w = images.shape # batch size, _, height, width
67 | bs = min(bs, max_subplots) # limit plot images
68 | ns = np.ceil(bs ** 0.5) # number of subplots (square)
69 | if np.max(images[0]) <= 1:
70 | images *= 255 # de-normalise (optional)
71 |
72 | # Build Image
73 | mosaic = np.full((int(ns * h), int(ns * w), 3), 255, dtype=np.uint8) # init
74 | for i, im in enumerate(images):
75 | if i == max_subplots: # if last batch has fewer images than we expect
76 | break
77 | x, y = int(w * (i // ns)), int(h * (i % ns)) # block origin
78 | im = im.transpose(1, 2, 0)
79 | mosaic[y:y + h, x:x + w, :] = im
80 |
81 | # Resize (optional)
82 | scale = max_size / ns / max(h, w)
83 | if scale < 1:
84 | h = math.ceil(scale * h)
85 | w = math.ceil(scale * w)
86 | mosaic = cv2.resize(mosaic, tuple(int(x * ns) for x in (w, h)))
87 |
88 | # Annotate
89 | fs = int((h + w) * ns * 0.01) # font size
90 | annotator = Annotator(mosaic, line_width=round(fs / 10), font_size=fs, pil=True, example=names)
91 | for i in range(i + 1):
92 | x, y = int(w * (i // ns)), int(h * (i % ns)) # block origin
93 | annotator.rectangle([x, y, x + w, y + h], None, (255, 255, 255), width=2) # borders
94 | if paths:
95 | annotator.text((x + 5, y + 5 + h), text=Path(paths[i]).name[:40], txt_color=(220, 220, 220)) # filenames
96 | if len(targets) > 0:
97 | idx = targets[:, 0] == i
98 | ti = targets[idx] # image targets
99 |
100 | boxes = xywh2xyxy(ti[:, 2:6]).T
101 | classes = ti[:, 1].astype('int')
102 | labels = ti.shape[1] == 6 # labels if no conf column
103 | conf = None if labels else ti[:, 6] # check for confidence presence (label vs pred)
104 |
105 | if boxes.shape[1]:
106 | if boxes.max() <= 1.01: # if normalized with tolerance 0.01
107 | boxes[[0, 2]] *= w # scale to pixels
108 | boxes[[1, 3]] *= h
109 | elif scale < 1: # absolute coords need scale if image scales
110 | boxes *= scale
111 | boxes[[0, 2]] += x
112 | boxes[[1, 3]] += y
113 | for j, box in enumerate(boxes.T.tolist()):
114 | cls = classes[j]
115 | color = colors(cls)
116 | cls = names[cls] if names else cls
117 | if labels or conf[j] > 0.25: # 0.25 conf thresh
118 | label = f'{cls}' if labels else f'{cls} {conf[j]:.1f}'
119 | annotator.box_label(box, label, color=color)
120 |
121 | # Plot masks
122 | if len(masks):
123 | if masks.max() > 1.0: # mean that masks are overlap
124 | image_masks = masks[[i]] # (1, 640, 640)
125 | nl = len(ti)
126 | index = np.arange(nl).reshape(nl, 1, 1) + 1
127 | image_masks = np.repeat(image_masks, nl, axis=0)
128 | image_masks = np.where(image_masks == index, 1.0, 0.0)
129 | else:
130 | image_masks = masks[idx]
131 |
132 | im = np.asarray(annotator.im).copy()
133 | for j, box in enumerate(boxes.T.tolist()):
134 | if labels or conf[j] > 0.25: # 0.25 conf thresh
135 | color = colors(classes[j])
136 | mh, mw = image_masks[j].shape
137 | if mh != h or mw != w:
138 | mask = image_masks[j].astype(np.uint8)
139 | mask = cv2.resize(mask, (w, h))
140 | mask = mask.astype(np.bool)
141 | else:
142 | mask = image_masks[j].astype(np.bool)
143 | with contextlib.suppress(Exception):
144 | im[y:y + h, x:x + w, :][mask] = im[y:y + h, x:x + w, :][mask] * 0.4 + np.array(color) * 0.6
145 | annotator.fromarray(im)
146 | annotator.im.save(fname) # save
147 |
148 |
149 | def plot_results_with_masks(file="path/to/results.csv", dir="", best=True):
150 | # Plot training results.csv. Usage: from utils.plots import *; plot_results('path/to/results.csv')
151 | save_dir = Path(file).parent if file else Path(dir)
152 | fig, ax = plt.subplots(2, 8, figsize=(18, 6), tight_layout=True)
153 | ax = ax.ravel()
154 | files = list(save_dir.glob("results*.csv"))
155 | assert len(files), f"No results.csv files found in {save_dir.resolve()}, nothing to plot."
156 | for f in files:
157 | try:
158 | data = pd.read_csv(f)
159 | index = np.argmax(
160 | 0.9 * data.values[:, 8] + 0.1 * data.values[:, 7] + 0.9 * data.values[:, 12] +
161 | 0.1 * data.values[:, 11],)
162 | s = [x.strip() for x in data.columns]
163 | x = data.values[:, 0]
164 | for i, j in enumerate([1, 2, 3, 4, 5, 6, 9, 10, 13, 14, 15, 16, 7, 8, 11, 12]):
165 | y = data.values[:, j]
166 | # y[y == 0] = np.nan # don't show zero values
167 | ax[i].plot(x, y, marker=".", label=f.stem, linewidth=2, markersize=2)
168 | if best:
169 | # best
170 | ax[i].scatter(index, y[index], color="r", label=f"best:{index}", marker="*", linewidth=3)
171 | ax[i].set_title(s[j] + f"\n{round(y[index], 5)}")
172 | else:
173 | # last
174 | ax[i].scatter(x[-1], y[-1], color="r", label="last", marker="*", linewidth=3)
175 | ax[i].set_title(s[j] + f"\n{round(y[-1], 5)}")
176 | # if j in [8, 9, 10]: # share train and val loss y axes
177 | # ax[i].get_shared_y_axes().join(ax[i], ax[i - 5])
178 | except Exception as e:
179 | print(f"Warning: Plotting error for {f}: {e}")
180 | ax[1].legend()
181 | fig.savefig(save_dir / "results.png", dpi=200)
182 | plt.close()
183 |
--------------------------------------------------------------------------------
/utils/wandb_logging/__init__.py:
--------------------------------------------------------------------------------
1 | # init
--------------------------------------------------------------------------------
/utils/wandb_logging/log_dataset.py:
--------------------------------------------------------------------------------
1 | import argparse
2 |
3 | import yaml
4 |
5 | from wandb_utils import WandbLogger
6 |
7 | WANDB_ARTIFACT_PREFIX = 'wandb-artifact://'
8 |
9 |
10 | def create_dataset_artifact(opt):
11 | with open(opt.data) as f:
12 | data = yaml.load(f, Loader=yaml.SafeLoader) # data dict
13 | logger = WandbLogger(opt, '', None, data, job_type='Dataset Creation')
14 |
15 |
16 | if __name__ == '__main__':
17 | parser = argparse.ArgumentParser()
18 | parser.add_argument('--data', type=str, default='data/coco.yaml', help='data.yaml path')
19 | parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
20 | parser.add_argument('--project', type=str, default='YOLOR', help='name of W&B Project')
21 | opt = parser.parse_args()
22 | opt.resume = False # Explicitly disallow resume check for dataset upload job
23 |
24 | create_dataset_artifact(opt)
25 |
--------------------------------------------------------------------------------