YOLOv5 segmentation model support (#9052)
* [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix duplicate plots.py * Fix check_font() * # torch.use_deterministic_algorithms(True) * update doc detect->predict * Resolve precommit for segment/train and segment/val * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Resolve precommit for utils/segment * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Resolve precommit min_wh * Resolve precommit utils/segment/plots * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Resolve precommit utils/segment/general * Align NMS-seg closer to NMS * restore deterministic init_seeds code * remove easydict dependency * update * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * restore output_to_target mask * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update * cleanup * Remove unused ImageFont import * Unified NMS * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * DetectMultiBackend compatibility * segment/predict.py update * update plot colors * fix bbox shifted * sort bbox by confidence * enable overlap by default * Merge detect/segment output_to_target() function * Start segmentation CI * fix plots * Update ci-testing.yml * fix training whitespace * optimize process mask functions (can we merge both?) * Update predict/detect * Update plot_images * Update plot_images_and_masks * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix * Add train to CI * fix precommit * fix precommit CI * fix precommit pycocotools * fix val float issues * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix masks float float issues * suppress errors * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix no-predictions plotting bug * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add CSV Logger * fix val len(plot_masks) * speed up evaluation * fix process_mask * fix plots * update segment/utils build_targets * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * optimize utils/segment/general crop() * optimize utils/segment/general crop() 2 * minor updates * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * torch.where revert * downsample only if different shape * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * loss cleanup * loss cleanup 2 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * loss cleanup 3 * update project names * Rename -seg yamls from _underscore to -dash * prepare for yolov5n-seg.pt * precommit space fix * add coco128-seg.yaml * update coco128-seg comments * cleanup val.py * Major val.py cleanup * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * precommit fix * precommit fix * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * optional pycocotools * remove CI pip install pycocotools (auto-installed now) * seg yaml fix * optimize mask_iou() and masks_iou() * threaded fix * Major train.py update * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Major segments/val/process_batch() update * yolov5/val updates from segment * process_batch numpy/tensor fix * opt-in to pycocotools with --save-json * threaded pycocotools ops for 2x speed increase * Avoid permute contiguous if possible * Add max_det=300 argument to both val.py and segment/val.py * fix onnx_dynamic * speed up pycocotools ops * faster process_mask(upsample=True) for predict * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * eliminate permutations for process_mask(upsample=True) * eliminate permute-contiguous in crop(), use native dimension order * cleanup comment * Add Proto() module * fix class count * fix anchor order * broadcast mask_gti in loss for speed * Cleanup seg loss * faster indexing * faster indexing fix * faster indexing fix2 * revert faster indexing * fix validation plotting * Loss cleanup and mxyxy simplification * Loss cleanup and mxyxy simplification 2 * revert validation plotting * replace missing tanh * Eliminate last permutation * delete unneeded .float() * Remove MaskIOULoss and crop(if HWC) * Final v6.3 SegmentationModel architecture updates * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add support for TF export * remove debugger trace * add call * update * update * Merge master * Merge master * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update dataloaders.py * Restore CI * Update dataloaders.py * Fix TF/TFLite export for segmentation model * Merge master * Cleanup predict.py mask plotting * cleanup scale_masks() * rename scale_masks to scale_image * cleanup/optimize plot_masks * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add Annotator.masks() * Annotator.masks() fix * Update plots.py * Annotator mask optimization * Rename crop() to crop_mask() * Do not crop in predict.py * crop always * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Merge master * Add vid-stride from master PR * Update seg model outputs * Update seg model outputs * Add segmentation benchmarks * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add segmentation benchmarks * Add segmentation benchmarks * Add segmentation benchmarks * Fix DetectMultiBackend for OpenVINO * update Annotator.masks * fix val plot * revert val plot * clean up * revert pil * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix CI error * fix predict log * remove upsample * update interpolate * fix validation plot logging * Annotator.masks() cleanup * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Remove segmentation_model definition * Restore 0.99999 decimals Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com> Co-authored-by: Laughing-q <1185102784@qq.com> Co-authored-by: Jiacong Fang <zldrobit@126.com>
This commit is contained in:
27 files changed
+3091
-71
No files matched your search
@@ -1,161 +0,0 @@
|
||||
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
|
||||
"""
|
||||
Run YOLOv5 benchmarks on all supported export formats
|
||||
|
||||
Format | `export.py --include` | Model
|
||||
--- | --- | ---
|
||||
PyTorch | - | yolov5s.pt
|
||||
TorchScript | `torchscript` | yolov5s.torchscript
|
||||
ONNX | `onnx` | yolov5s.onnx
|
||||
OpenVINO | `openvino` | yolov5s_openvino_model/
|
||||
TensorRT | `engine` | yolov5s.engine
|
||||
CoreML | `coreml` | yolov5s.mlmodel
|
||||
TensorFlow SavedModel | `saved_model` | yolov5s_saved_model/
|
||||
TensorFlow GraphDef | `pb` | yolov5s.pb
|
||||
TensorFlow Lite | `tflite` | yolov5s.tflite
|
||||
TensorFlow Edge TPU | `edgetpu` | yolov5s_edgetpu.tflite
|
||||
TensorFlow.js | `tfjs` | yolov5s_web_model/
|
||||
|
||||
Requirements:
|
||||
$ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime openvino-dev tensorflow-cpu # CPU
|
||||
$ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime-gpu openvino-dev tensorflow # GPU
|
||||
$ pip install -U nvidia-tensorrt --index-url https://pypi.ngc.nvidia.com # TensorRT
|
||||
|
||||
Usage:
|
||||
$ python utils/benchmarks.py --weights yolov5s.pt --img 640
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import platform
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import pandas as pd
|
||||
|
||||
FILE = Path(__file__).resolve()
|
||||
ROOT = FILE.parents[1] # YOLOv5 root directory
|
||||
if str(ROOT) not in sys.path:
|
||||
sys.path.append(str(ROOT)) # add ROOT to PATH
|
||||
# ROOT = ROOT.relative_to(Path.cwd()) # relative
|
||||
|
||||
import export
|
||||
import val
|
||||
from utils import notebook_init
|
||||
from utils.general import LOGGER, check_yaml, file_size, print_args
|
||||
from utils.torch_utils import select_device
|
||||
|
||||
|
||||
def run(
|
||||
weights=ROOT / 'yolov5s.pt', # weights path
|
||||
imgsz=640, # inference size (pixels)
|
||||
batch_size=1, # batch size
|
||||
data=ROOT / 'data/coco128.yaml', # dataset.yaml path
|
||||
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
|
||||
half=False, # use FP16 half-precision inference
|
||||
test=False, # test exports only
|
||||
pt_only=False, # test PyTorch only
|
||||
hard_fail=False, # throw error on benchmark failure
|
||||
):
|
||||
y, t = [], time.time()
|
||||
device = select_device(device)
|
||||
for i, (name, f, suffix, cpu, gpu) in export.export_formats().iterrows(): # index, (name, file, suffix, CPU, GPU)
|
||||
try:
|
||||
assert i not in (9, 10, 11), 'inference not supported' # Edge TPU, TF.js and Paddle are unsupported
|
||||
assert i != 5 or platform.system() == 'Darwin', 'inference only supported on macOS>=10.13' # CoreML
|
||||
if 'cpu' in device.type:
|
||||
assert cpu, 'inference not supported on CPU'
|
||||
if 'cuda' in device.type:
|
||||
assert gpu, 'inference not supported on GPU'
|
||||
|
||||
# Export
|
||||
if f == '-':
|
||||
w = weights # PyTorch format
|
||||
else:
|
||||
w = export.run(weights=weights, imgsz=[imgsz], include=[f], device=device, half=half)[-1] # all others
|
||||
assert suffix in str(w), 'export failed'
|
||||
|
||||
# Validate
|
||||
result = val.run(data, w, batch_size, imgsz, plots=False, device=device, task='benchmark', half=half)
|
||||
metrics = result[0] # metrics (mp, mr, map50, map, *losses(box, obj, cls))
|
||||
speeds = result[2] # times (preprocess, inference, postprocess)
|
||||
y.append([name, round(file_size(w), 1), round(metrics[3], 4), round(speeds[1], 2)]) # MB, mAP, t_inference
|
||||
except Exception as e:
|
||||
if hard_fail:
|
||||
assert type(e) is AssertionError, f'Benchmark --hard-fail for {name}: {e}'
|
||||
LOGGER.warning(f'WARNING: Benchmark failure for {name}: {e}')
|
||||
y.append([name, None, None, None]) # mAP, t_inference
|
||||
if pt_only and i == 0:
|
||||
break # break after PyTorch
|
||||
|
||||
# Print results
|
||||
LOGGER.info('\n')
|
||||
parse_opt()
|
||||
notebook_init() # print system info
|
||||
c = ['Format', 'Size (MB)', 'mAP50-95', 'Inference time (ms)'] if map else ['Format', 'Export', '', '']
|
||||
py = pd.DataFrame(y, columns=c)
|
||||
LOGGER.info(f'\nBenchmarks complete ({time.time() - t:.2f}s)')
|
||||
LOGGER.info(str(py if map else py.iloc[:, :2]))
|
||||
if hard_fail and isinstance(hard_fail, str):
|
||||
metrics = py['mAP50-95'].array # values to compare to floor
|
||||
floor = eval(hard_fail) # minimum metric floor to pass, i.e. = 0.29 mAP for YOLOv5n
|
||||
assert all(x > floor for x in metrics if pd.notna(x)), f'HARD FAIL: mAP50-95 < floor {floor}'
|
||||
return py
|
||||
|
||||
|
||||
def test(
|
||||
weights=ROOT / 'yolov5s.pt', # weights path
|
||||
imgsz=640, # inference size (pixels)
|
||||
batch_size=1, # batch size
|
||||
data=ROOT / 'data/coco128.yaml', # dataset.yaml path
|
||||
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
|
||||
half=False, # use FP16 half-precision inference
|
||||
test=False, # test exports only
|
||||
pt_only=False, # test PyTorch only
|
||||
hard_fail=False, # throw error on benchmark failure
|
||||
):
|
||||
y, t = [], time.time()
|
||||
device = select_device(device)
|
||||
for i, (name, f, suffix, gpu) in export.export_formats().iterrows(): # index, (name, file, suffix, gpu-capable)
|
||||
try:
|
||||
w = weights if f == '-' else \
|
||||
export.run(weights=weights, imgsz=[imgsz], include=[f], device=device, half=half)[-1] # weights
|
||||
assert suffix in str(w), 'export failed'
|
||||
y.append([name, True])
|
||||
except Exception:
|
||||
y.append([name, False]) # mAP, t_inference
|
||||
|
||||
# Print results
|
||||
LOGGER.info('\n')
|
||||
parse_opt()
|
||||
notebook_init() # print system info
|
||||
py = pd.DataFrame(y, columns=['Format', 'Export'])
|
||||
LOGGER.info(f'\nExports complete ({time.time() - t:.2f}s)')
|
||||
LOGGER.info(str(py))
|
||||
return py
|
||||
|
||||
|
||||
def parse_opt():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='weights path')
|
||||
parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='inference size (pixels)')
|
||||
parser.add_argument('--batch-size', type=int, default=1, help='batch size')
|
||||
parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
|
||||
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
||||
parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
|
||||
parser.add_argument('--test', action='store_true', help='test exports only')
|
||||
parser.add_argument('--pt-only', action='store_true', help='test PyTorch only')
|
||||
parser.add_argument('--hard-fail', nargs='?', const=True, default=False, help='Exception on error or < min metric')
|
||||
opt = parser.parse_args()
|
||||
opt.data = check_yaml(opt.data) # check YAML
|
||||
print_args(vars(opt))
|
||||
return opt
|
||||
|
||||
|
||||
def main(opt):
|
||||
test(**vars(opt)) if opt.test else run(**vars(opt))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
opt = parse_opt()
|
||||
main(opt)
|
||||
Executable → Regular
+1
@@ -484,6 +484,7 @@ class LoadImagesAndLabels(Dataset):
|
||||
self.im_files = [self.im_files[i] for i in irect]
|
||||
self.label_files = [self.label_files[i] for i in irect]
|
||||
self.labels = [self.labels[i] for i in irect]
|
||||
self.segments = [self.segments[i] for i in irect]
|
||||
self.shapes = s[irect] # wh
|
||||
ar = ar[irect]
|
||||
|
||||
|
||||
Executable → Regular
+26
-19
@@ -798,15 +798,18 @@ def clip_coords(boxes, shape):
|
||||
boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, shape[0]) # y1, y2
|
||||
|
||||
|
||||
def non_max_suppression(prediction,
|
||||
conf_thres=0.25,
|
||||
iou_thres=0.45,
|
||||
classes=None,
|
||||
agnostic=False,
|
||||
multi_label=False,
|
||||
labels=(),
|
||||
max_det=300):
|
||||
"""Non-Maximum Suppression (NMS) on inference results to reject overlapping bounding boxes
|
||||
def non_max_suppression(
|
||||
prediction,
|
||||
conf_thres=0.25,
|
||||
iou_thres=0.45,
|
||||
classes=None,
|
||||
agnostic=False,
|
||||
multi_label=False,
|
||||
labels=(),
|
||||
max_det=300,
|
||||
nm=0, # number of masks
|
||||
):
|
||||
"""Non-Maximum Suppression (NMS) on inference results to reject overlapping detections
|
||||
|
||||
Returns:
|
||||
list of detections, on (n,6) tensor per image [xyxy, conf, cls]
|
||||
@@ -816,7 +819,7 @@ def non_max_suppression(prediction,
|
||||
prediction = prediction[0] # select only inference output
|
||||
|
||||
bs = prediction.shape[0] # batch size
|
||||
nc = prediction.shape[2] - 5 # number of classes
|
||||
nc = prediction.shape[2] - nm - 5 # number of classes
|
||||
xc = prediction[..., 4] > conf_thres # candidates
|
||||
|
||||
# Checks
|
||||
@@ -827,13 +830,14 @@ def non_max_suppression(prediction,
|
||||
# min_wh = 2 # (pixels) minimum box width and height
|
||||
max_wh = 7680 # (pixels) maximum box width and height
|
||||
max_nms = 30000 # maximum number of boxes into torchvision.ops.nms()
|
||||
time_limit = 0.3 + 0.03 * bs # seconds to quit after
|
||||
time_limit = 0.5 + 0.05 * bs # seconds to quit after
|
||||
redundant = True # require redundant detections
|
||||
multi_label &= nc > 1 # multiple labels per box (adds 0.5ms/img)
|
||||
merge = False # use merge-NMS
|
||||
|
||||
t = time.time()
|
||||
output = [torch.zeros((0, 6), device=prediction.device)] * bs
|
||||
mi = 5 + nc # mask start index
|
||||
output = [torch.zeros((0, 6 + nm), device=prediction.device)] * bs
|
||||
for xi, x in enumerate(prediction): # image index, image inference
|
||||
# Apply constraints
|
||||
# x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height
|
||||
@@ -842,7 +846,7 @@ def non_max_suppression(prediction,
|
||||
# Cat apriori labels if autolabelling
|
||||
if labels and len(labels[xi]):
|
||||
lb = labels[xi]
|
||||
v = torch.zeros((len(lb), nc + 5), device=x.device)
|
||||
v = torch.zeros((len(lb), nc + nm + 5), device=x.device)
|
||||
v[:, :4] = lb[:, 1:5] # box
|
||||
v[:, 4] = 1.0 # conf
|
||||
v[range(len(lb)), lb[:, 0].long() + 5] = 1.0 # cls
|
||||
@@ -855,16 +859,17 @@ def non_max_suppression(prediction,
|
||||
# Compute conf
|
||||
x[:, 5:] *= x[:, 4:5] # conf = obj_conf * cls_conf
|
||||
|
||||
# Box (center x, center y, width, height) to (x1, y1, x2, y2)
|
||||
box = xywh2xyxy(x[:, :4])
|
||||
# Box/Mask
|
||||
box = xywh2xyxy(x[:, :4]) # center_x, center_y, width, height) to (x1, y1, x2, y2)
|
||||
mask = x[:, mi:] # zero columns if no masks
|
||||
|
||||
# Detections matrix nx6 (xyxy, conf, cls)
|
||||
if multi_label:
|
||||
i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T
|
||||
x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
|
||||
i, j = (x[:, 5:mi] > conf_thres).nonzero(as_tuple=False).T
|
||||
x = torch.cat((box[i], x[i, 5 + j, None], j[:, None].float(), mask[i]), 1)
|
||||
else: # best class only
|
||||
conf, j = x[:, 5:].max(1, keepdim=True)
|
||||
x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]
|
||||
conf, j = x[:, 5:mi].max(1, keepdim=True)
|
||||
x = torch.cat((box, conf, j.float(), mask), 1)[conf.view(-1) > conf_thres]
|
||||
|
||||
# Filter by class
|
||||
if classes is not None:
|
||||
@@ -880,6 +885,8 @@ def non_max_suppression(prediction,
|
||||
continue
|
||||
elif n > max_nms: # excess boxes
|
||||
x = x[x[:, 4].argsort(descending=True)[:max_nms]] # sort by confidence
|
||||
else:
|
||||
x = x[x[:, 4].argsort(descending=True)] # sort by confidence
|
||||
|
||||
# Batched NMS
|
||||
c = x[:, 5:6] * (0 if agnostic else max_wh) # classes
|
||||
|
||||
+5
-5
@@ -28,7 +28,7 @@ def smooth(y, f=0.05):
|
||||
return np.convolve(yp, np.ones(nf) / nf, mode='valid') # y-smoothed
|
||||
|
||||
|
||||
def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='.', names=(), eps=1e-16):
|
||||
def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='.', names=(), eps=1e-16, prefix=""):
|
||||
""" Compute the average precision, given the recall and precision curves.
|
||||
Source: https://github.com/rafaelpadilla/Object-Detection-Metrics.
|
||||
# Arguments
|
||||
@@ -83,10 +83,10 @@ def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='.', names
|
||||
names = [v for k, v in names.items() if k in unique_classes] # list: only classes that have data
|
||||
names = dict(enumerate(names)) # to dict
|
||||
if plot:
|
||||
plot_pr_curve(px, py, ap, Path(save_dir) / 'PR_curve.png', names)
|
||||
plot_mc_curve(px, f1, Path(save_dir) / 'F1_curve.png', names, ylabel='F1')
|
||||
plot_mc_curve(px, p, Path(save_dir) / 'P_curve.png', names, ylabel='Precision')
|
||||
plot_mc_curve(px, r, Path(save_dir) / 'R_curve.png', names, ylabel='Recall')
|
||||
plot_pr_curve(px, py, ap, Path(save_dir) / f'{prefix}PR_curve.png', names)
|
||||
plot_mc_curve(px, f1, Path(save_dir) / f'{prefix}F1_curve.png', names, ylabel='F1')
|
||||
plot_mc_curve(px, p, Path(save_dir) / f'{prefix}P_curve.png', names, ylabel='Precision')
|
||||
plot_mc_curve(px, r, Path(save_dir) / f'{prefix}R_curve.png', names, ylabel='Recall')
|
||||
|
||||
i = smooth(f1.mean(0), 0.1).argmax() # max F1 index
|
||||
p, r, f1 = p[:, i], r[:, i], f1[:, i]
|
||||
|
||||
+64
-7
@@ -23,6 +23,7 @@ from utils import TryExcept, threaded
|
||||
from utils.general import (CONFIG_DIR, FONT, LOGGER, check_font, check_requirements, clip_coords, increment_path,
|
||||
is_ascii, xywh2xyxy, xyxy2xywh)
|
||||
from utils.metrics import fitness
|
||||
from utils.segment.general import scale_image
|
||||
|
||||
# Settings
|
||||
RANK = int(os.getenv('RANK', -1))
|
||||
@@ -113,6 +114,52 @@ class Annotator:
|
||||
thickness=tf,
|
||||
lineType=cv2.LINE_AA)
|
||||
|
||||
def masks(self, masks, colors, im_gpu=None, alpha=0.5):
|
||||
"""Plot masks at once.
|
||||
Args:
|
||||
masks (tensor): predicted masks on cuda, shape: [n, h, w]
|
||||
colors (List[List[Int]]): colors for predicted masks, [[r, g, b] * n]
|
||||
im_gpu (tensor): img is in cuda, shape: [3, h, w], range: [0, 1]
|
||||
alpha (float): mask transparency: 0.0 fully transparent, 1.0 opaque
|
||||
"""
|
||||
if self.pil:
|
||||
# convert to numpy first
|
||||
self.im = np.asarray(self.im).copy()
|
||||
if im_gpu is None:
|
||||
# Add multiple masks of shape(h,w,n) with colors list([r,g,b], [r,g,b], ...)
|
||||
if len(masks) == 0:
|
||||
return
|
||||
if isinstance(masks, torch.Tensor):
|
||||
masks = torch.as_tensor(masks, dtype=torch.uint8)
|
||||
masks = masks.permute(1, 2, 0).contiguous()
|
||||
masks = masks.cpu().numpy()
|
||||
# masks = np.ascontiguousarray(masks.transpose(1, 2, 0))
|
||||
masks = scale_image(masks.shape[:2], masks, self.im.shape)
|
||||
masks = np.asarray(masks, dtype=np.float32)
|
||||
colors = np.asarray(colors, dtype=np.float32) # shape(n,3)
|
||||
s = masks.sum(2, keepdims=True).clip(0, 1) # add all masks together
|
||||
masks = (masks @ colors).clip(0, 255) # (h,w,n) @ (n,3) = (h,w,3)
|
||||
self.im[:] = masks * alpha + self.im * (1 - s * alpha)
|
||||
else:
|
||||
if len(masks) == 0:
|
||||
self.im[:] = im_gpu.permute(1, 2, 0).contiguous().cpu().numpy() * 255
|
||||
colors = torch.tensor(colors, device=im_gpu.device, dtype=torch.float32) / 255.0
|
||||
colors = colors[:, None, None] # shape(n,1,1,3)
|
||||
masks = masks.unsqueeze(3) # shape(n,h,w,1)
|
||||
masks_color = masks * (colors * alpha) # shape(n,h,w,3)
|
||||
|
||||
inv_alph_masks = (1 - masks * alpha).cumprod(0) # shape(n,h,w,1)
|
||||
mcs = (masks_color * inv_alph_masks).sum(0) * 2 # mask color summand shape(n,h,w,3)
|
||||
|
||||
im_gpu = im_gpu.flip(dims=[0]) # flip channel
|
||||
im_gpu = im_gpu.permute(1, 2, 0).contiguous() # shape(h,w,3)
|
||||
im_gpu = im_gpu * inv_alph_masks[-1] + mcs
|
||||
im_mask = (im_gpu * 255).byte().cpu().numpy()
|
||||
self.im[:] = scale_image(im_gpu.shape, im_mask, self.im.shape)
|
||||
if self.pil:
|
||||
# convert im back to PIL and update draw
|
||||
self.fromarray(self.im)
|
||||
|
||||
def rectangle(self, xy, fill=None, outline=None, width=1):
|
||||
# Add rectangle to image (PIL-only)
|
||||
self.draw.rectangle(xy, fill, outline, width)
|
||||
@@ -124,6 +171,11 @@ class Annotator:
|
||||
xy[1] += 1 - h
|
||||
self.draw.text(xy, text, fill=txt_color, font=self.font)
|
||||
|
||||
def fromarray(self, im):
|
||||
# Update self.im from a numpy array
|
||||
self.im = im if isinstance(im, Image.Image) else Image.fromarray(im)
|
||||
self.draw = ImageDraw.Draw(self.im)
|
||||
|
||||
def result(self):
|
||||
# Return annotated image as array
|
||||
return np.asarray(self.im)
|
||||
@@ -180,26 +232,31 @@ def butter_lowpass_filtfilt(data, cutoff=1500, fs=50000, order=5):
|
||||
return filtfilt(b, a, data) # forward-backward filter
|
||||
|
||||
|
||||
def output_to_target(output):
|
||||
# Convert model output to target format [batch_id, class_id, x, y, w, h, conf]
|
||||
def output_to_target(output, max_det=300):
|
||||
# Convert model output to target format [batch_id, class_id, x, y, w, h, conf] for plotting
|
||||
targets = []
|
||||
for i, o in enumerate(output):
|
||||
targets.extend([i, cls, *list(*xyxy2xywh(np.array(box)[None])), conf] for *box, conf, cls in o.cpu().numpy())
|
||||
return np.array(targets)
|
||||
box, conf, cls = o[:max_det, :6].cpu().split((4, 1, 1), 1)
|
||||
j = torch.full((conf.shape[0], 1), i)
|
||||
targets.append(torch.cat((j, cls, xyxy2xywh(box), conf), 1))
|
||||
return torch.cat(targets, 0).numpy()
|
||||
|
||||
|
||||
@threaded
|
||||
def plot_images(images, targets, paths=None, fname='images.jpg', names=None, max_size=1920, max_subplots=16):
|
||||
def plot_images(images, targets, paths=None, fname='images.jpg', names=None):
|
||||
# Plot image grid with labels
|
||||
if isinstance(images, torch.Tensor):
|
||||
images = images.cpu().float().numpy()
|
||||
if isinstance(targets, torch.Tensor):
|
||||
targets = targets.cpu().numpy()
|
||||
if np.max(images[0]) <= 1:
|
||||
images *= 255 # de-normalise (optional)
|
||||
|
||||
max_size = 1920 # max image size
|
||||
max_subplots = 16 # max image subplots, i.e. 4x4
|
||||
bs, _, h, w = images.shape # batch size, _, height, width
|
||||
bs = min(bs, max_subplots) # limit plot images
|
||||
ns = np.ceil(bs ** 0.5) # number of subplots (square)
|
||||
if np.max(images[0]) <= 1:
|
||||
images *= 255 # de-normalise (optional)
|
||||
|
||||
# Build Image
|
||||
mosaic = np.full((int(ns * h), int(ns * w), 3), 255, dtype=np.uint8) # init
|
||||
|
||||
Whitespace-only changes.
@@ -0,0 +1,104 @@
|
||||
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
|
||||
"""
|
||||
Image augmentation functions
|
||||
"""
|
||||
|
||||
import math
|
||||
import random
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from ..augmentations import box_candidates
|
||||
from ..general import resample_segments, segment2box
|
||||
|
||||
|
||||
def mixup(im, labels, segments, im2, labels2, segments2):
|
||||
# Applies MixUp augmentation https://arxiv.org/pdf/1710.09412.pdf
|
||||
r = np.random.beta(32.0, 32.0) # mixup ratio, alpha=beta=32.0
|
||||
im = (im * r + im2 * (1 - r)).astype(np.uint8)
|
||||
labels = np.concatenate((labels, labels2), 0)
|
||||
segments = np.concatenate((segments, segments2), 0)
|
||||
return im, labels, segments
|
||||
|
||||
|
||||
def random_perspective(im,
|
||||
targets=(),
|
||||
segments=(),
|
||||
degrees=10,
|
||||
translate=.1,
|
||||
scale=.1,
|
||||
shear=10,
|
||||
perspective=0.0,
|
||||
border=(0, 0)):
|
||||
# torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(.1, .1), scale=(.9, 1.1), shear=(-10, 10))
|
||||
# targets = [cls, xyxy]
|
||||
|
||||
height = im.shape[0] + border[0] * 2 # shape(h,w,c)
|
||||
width = im.shape[1] + border[1] * 2
|
||||
|
||||
# Center
|
||||
C = np.eye(3)
|
||||
C[0, 2] = -im.shape[1] / 2 # x translation (pixels)
|
||||
C[1, 2] = -im.shape[0] / 2 # y translation (pixels)
|
||||
|
||||
# Perspective
|
||||
P = np.eye(3)
|
||||
P[2, 0] = random.uniform(-perspective, perspective) # x perspective (about y)
|
||||
P[2, 1] = random.uniform(-perspective, perspective) # y perspective (about x)
|
||||
|
||||
# Rotation and Scale
|
||||
R = np.eye(3)
|
||||
a = random.uniform(-degrees, degrees)
|
||||
# a += random.choice([-180, -90, 0, 90]) # add 90deg rotations to small rotations
|
||||
s = random.uniform(1 - scale, 1 + scale)
|
||||
# s = 2 ** random.uniform(-scale, scale)
|
||||
R[:2] = cv2.getRotationMatrix2D(angle=a, center=(0, 0), scale=s)
|
||||
|
||||
# Shear
|
||||
S = np.eye(3)
|
||||
S[0, 1] = math.tan(random.uniform(-shear, shear) * math.pi / 180) # x shear (deg)
|
||||
S[1, 0] = math.tan(random.uniform(-shear, shear) * math.pi / 180) # y shear (deg)
|
||||
|
||||
# Translation
|
||||
T = np.eye(3)
|
||||
T[0, 2] = (random.uniform(0.5 - translate, 0.5 + translate) * width) # x translation (pixels)
|
||||
T[1, 2] = (random.uniform(0.5 - translate, 0.5 + translate) * height) # y translation (pixels)
|
||||
|
||||
# Combined rotation matrix
|
||||
M = T @ S @ R @ P @ C # order of operations (right to left) is IMPORTANT
|
||||
if (border[0] != 0) or (border[1] != 0) or (M != np.eye(3)).any(): # image changed
|
||||
if perspective:
|
||||
im = cv2.warpPerspective(im, M, dsize=(width, height), borderValue=(114, 114, 114))
|
||||
else: # affine
|
||||
im = cv2.warpAffine(im, M[:2], dsize=(width, height), borderValue=(114, 114, 114))
|
||||
|
||||
# Visualize
|
||||
# import matplotlib.pyplot as plt
|
||||
# ax = plt.subplots(1, 2, figsize=(12, 6))[1].ravel()
|
||||
# ax[0].imshow(im[:, :, ::-1]) # base
|
||||
# ax[1].imshow(im2[:, :, ::-1]) # warped
|
||||
|
||||
# Transform label coordinates
|
||||
n = len(targets)
|
||||
new_segments = []
|
||||
if n:
|
||||
new = np.zeros((n, 4))
|
||||
segments = resample_segments(segments) # upsample
|
||||
for i, segment in enumerate(segments):
|
||||
xy = np.ones((len(segment), 3))
|
||||
xy[:, :2] = segment
|
||||
xy = xy @ M.T # transform
|
||||
xy = (xy[:, :2] / xy[:, 2:3] if perspective else xy[:, :2]) # perspective rescale or affine
|
||||
|
||||
# clip
|
||||
new[i] = segment2box(xy, width, height)
|
||||
new_segments.append(xy)
|
||||
|
||||
# filter candidates
|
||||
i = box_candidates(box1=targets[:, 1:5].T * s, box2=new.T, area_thr=0.01)
|
||||
targets = targets[i]
|
||||
targets[:, 1:5] = new[i]
|
||||
new_segments = np.array(new_segments)[i]
|
||||
|
||||
return im, targets, new_segments
|
||||
@@ -0,0 +1,330 @@
|
||||
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
|
||||
"""
|
||||
Dataloaders
|
||||
"""
|
||||
|
||||
import os
|
||||
import random
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.utils.data import DataLoader, distributed
|
||||
|
||||
from ..augmentations import augment_hsv, copy_paste, letterbox
|
||||
from ..dataloaders import InfiniteDataLoader, LoadImagesAndLabels, seed_worker
|
||||
from ..general import LOGGER, xyn2xy, xywhn2xyxy, xyxy2xywhn
|
||||
from ..torch_utils import torch_distributed_zero_first
|
||||
from .augmentations import mixup, random_perspective
|
||||
|
||||
|
||||
def create_dataloader(path,
|
||||
imgsz,
|
||||
batch_size,
|
||||
stride,
|
||||
single_cls=False,
|
||||
hyp=None,
|
||||
augment=False,
|
||||
cache=False,
|
||||
pad=0.0,
|
||||
rect=False,
|
||||
rank=-1,
|
||||
workers=8,
|
||||
image_weights=False,
|
||||
quad=False,
|
||||
prefix='',
|
||||
shuffle=False,
|
||||
mask_downsample_ratio=1,
|
||||
overlap_mask=False):
|
||||
if rect and shuffle:
|
||||
LOGGER.warning('WARNING: --rect is incompatible with DataLoader shuffle, setting shuffle=False')
|
||||
shuffle = False
|
||||
with torch_distributed_zero_first(rank): # init dataset *.cache only once if DDP
|
||||
dataset = LoadImagesAndLabelsAndMasks(
|
||||
path,
|
||||
imgsz,
|
||||
batch_size,
|
||||
augment=augment, # augmentation
|
||||
hyp=hyp, # hyperparameters
|
||||
rect=rect, # rectangular batches
|
||||
cache_images=cache,
|
||||
single_cls=single_cls,
|
||||
stride=int(stride),
|
||||
pad=pad,
|
||||
image_weights=image_weights,
|
||||
prefix=prefix,
|
||||
downsample_ratio=mask_downsample_ratio,
|
||||
overlap=overlap_mask)
|
||||
|
||||
batch_size = min(batch_size, len(dataset))
|
||||
nd = torch.cuda.device_count() # number of CUDA devices
|
||||
nw = min([os.cpu_count() // max(nd, 1), batch_size if batch_size > 1 else 0, workers]) # number of workers
|
||||
sampler = None if rank == -1 else distributed.DistributedSampler(dataset, shuffle=shuffle)
|
||||
loader = DataLoader if image_weights else InfiniteDataLoader # only DataLoader allows for attribute updates
|
||||
# generator = torch.Generator()
|
||||
# generator.manual_seed(0)
|
||||
return loader(
|
||||
dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=shuffle and sampler is None,
|
||||
num_workers=nw,
|
||||
sampler=sampler,
|
||||
pin_memory=True,
|
||||
collate_fn=LoadImagesAndLabelsAndMasks.collate_fn4 if quad else LoadImagesAndLabelsAndMasks.collate_fn,
|
||||
worker_init_fn=seed_worker,
|
||||
# generator=generator,
|
||||
), dataset
|
||||
|
||||
|
||||
class LoadImagesAndLabelsAndMasks(LoadImagesAndLabels): # for training/testing
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
path,
|
||||
img_size=640,
|
||||
batch_size=16,
|
||||
augment=False,
|
||||
hyp=None,
|
||||
rect=False,
|
||||
image_weights=False,
|
||||
cache_images=False,
|
||||
single_cls=False,
|
||||
stride=32,
|
||||
pad=0,
|
||||
prefix="",
|
||||
downsample_ratio=1,
|
||||
overlap=False,
|
||||
):
|
||||
super().__init__(path, img_size, batch_size, augment, hyp, rect, image_weights, cache_images, single_cls,
|
||||
stride, pad, prefix)
|
||||
self.downsample_ratio = downsample_ratio
|
||||
self.overlap = overlap
|
||||
|
||||
def __getitem__(self, index):
|
||||
index = self.indices[index] # linear, shuffled, or image_weights
|
||||
|
||||
hyp = self.hyp
|
||||
mosaic = self.mosaic and random.random() < hyp['mosaic']
|
||||
masks = []
|
||||
if mosaic:
|
||||
# Load mosaic
|
||||
img, labels, segments = self.load_mosaic(index)
|
||||
shapes = None
|
||||
|
||||
# MixUp augmentation
|
||||
if random.random() < hyp["mixup"]:
|
||||
img, labels, segments = mixup(img, labels, segments, *self.load_mosaic(random.randint(0, self.n - 1)))
|
||||
|
||||
else:
|
||||
# Load image
|
||||
img, (h0, w0), (h, w) = self.load_image(index)
|
||||
|
||||
# Letterbox
|
||||
shape = self.batch_shapes[self.batch[index]] if self.rect else self.img_size # final letterboxed shape
|
||||
img, ratio, pad = letterbox(img, shape, auto=False, scaleup=self.augment)
|
||||
shapes = (h0, w0), ((h / h0, w / w0), pad) # for COCO mAP rescaling
|
||||
|
||||
labels = self.labels[index].copy()
|
||||
# [array, array, ....], array.shape=(num_points, 2), xyxyxyxy
|
||||
segments = self.segments[index].copy()
|
||||
if len(segments):
|
||||
for i_s in range(len(segments)):
|
||||
segments[i_s] = xyn2xy(
|
||||
segments[i_s],
|
||||
ratio[0] * w,
|
||||
ratio[1] * h,
|
||||
padw=pad[0],
|
||||
padh=pad[1],
|
||||
)
|
||||
if labels.size: # normalized xywh to pixel xyxy format
|
||||
labels[:, 1:] = xywhn2xyxy(labels[:, 1:], ratio[0] * w, ratio[1] * h, padw=pad[0], padh=pad[1])
|
||||
|
||||
if self.augment:
|
||||
img, labels, segments = random_perspective(
|
||||
img,
|
||||
labels,
|
||||
segments=segments,
|
||||
degrees=hyp["degrees"],
|
||||
translate=hyp["translate"],
|
||||
scale=hyp["scale"],
|
||||
shear=hyp["shear"],
|
||||
perspective=hyp["perspective"],
|
||||
return_seg=True,
|
||||
)
|
||||
|
||||
nl = len(labels) # number of labels
|
||||
if nl:
|
||||
labels[:, 1:5] = xyxy2xywhn(labels[:, 1:5], w=img.shape[1], h=img.shape[0], clip=True, eps=1e-3)
|
||||
if self.overlap:
|
||||
masks, sorted_idx = polygons2masks_overlap(img.shape[:2],
|
||||
segments,
|
||||
downsample_ratio=self.downsample_ratio)
|
||||
masks = masks[None] # (640, 640) -> (1, 640, 640)
|
||||
labels = labels[sorted_idx]
|
||||
else:
|
||||
masks = polygons2masks(img.shape[:2], segments, color=1, downsample_ratio=self.downsample_ratio)
|
||||
|
||||
masks = (torch.from_numpy(masks) if len(masks) else torch.zeros(1 if self.overlap else nl, img.shape[0] //
|
||||
self.downsample_ratio, img.shape[1] //
|
||||
self.downsample_ratio))
|
||||
# TODO: albumentations support
|
||||
if self.augment:
|
||||
# Albumentations
|
||||
# there are some augmentation that won't change boxes and masks,
|
||||
# so just be it for now.
|
||||
img, labels = self.albumentations(img, labels)
|
||||
nl = len(labels) # update after albumentations
|
||||
|
||||
# HSV color-space
|
||||
augment_hsv(img, hgain=hyp["hsv_h"], sgain=hyp["hsv_s"], vgain=hyp["hsv_v"])
|
||||
|
||||
# Flip up-down
|
||||
if random.random() < hyp["flipud"]:
|
||||
img = np.flipud(img)
|
||||
if nl:
|
||||
labels[:, 2] = 1 - labels[:, 2]
|
||||
masks = torch.flip(masks, dims=[1])
|
||||
|
||||
# Flip left-right
|
||||
if random.random() < hyp["fliplr"]:
|
||||
img = np.fliplr(img)
|
||||
if nl:
|
||||
labels[:, 1] = 1 - labels[:, 1]
|
||||
masks = torch.flip(masks, dims=[2])
|
||||
|
||||
# Cutouts # labels = cutout(img, labels, p=0.5)
|
||||
|
||||
labels_out = torch.zeros((nl, 6))
|
||||
if nl:
|
||||
labels_out[:, 1:] = torch.from_numpy(labels)
|
||||
|
||||
# Convert
|
||||
img = img.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB
|
||||
img = np.ascontiguousarray(img)
|
||||
|
||||
return (torch.from_numpy(img), labels_out, self.im_files[index], shapes, masks)
|
||||
|
||||
def load_mosaic(self, index):
|
||||
# YOLOv5 4-mosaic loader. Loads 1 image + 3 random images into a 4-image mosaic
|
||||
labels4, segments4 = [], []
|
||||
s = self.img_size
|
||||
yc, xc = (int(random.uniform(-x, 2 * s + x)) for x in self.mosaic_border) # mosaic center x, y
|
||||
|
||||
# 3 additional image indices
|
||||
indices = [index] + random.choices(self.indices, k=3) # 3 additional image indices
|
||||
for i, index in enumerate(indices):
|
||||
# Load image
|
||||
img, _, (h, w) = self.load_image(index)
|
||||
|
||||
# place img in img4
|
||||
if i == 0: # top left
|
||||
img4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8) # base image with 4 tiles
|
||||
x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc # xmin, ymin, xmax, ymax (large image)
|
||||
x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h # xmin, ymin, xmax, ymax (small image)
|
||||
elif i == 1: # top right
|
||||
x1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), yc
|
||||
x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), h
|
||||
elif i == 2: # bottom left
|
||||
x1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(s * 2, yc + h)
|
||||
x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, w, min(y2a - y1a, h)
|
||||
elif i == 3: # bottom right
|
||||
x1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(s * 2, yc + h)
|
||||
x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h)
|
||||
|
||||
img4[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b] # img4[ymin:ymax, xmin:xmax]
|
||||
padw = x1a - x1b
|
||||
padh = y1a - y1b
|
||||
|
||||
labels, segments = self.labels[index].copy(), self.segments[index].copy()
|
||||
|
||||
if labels.size:
|
||||
labels[:, 1:] = xywhn2xyxy(labels[:, 1:], w, h, padw, padh) # normalized xywh to pixel xyxy format
|
||||
segments = [xyn2xy(x, w, h, padw, padh) for x in segments]
|
||||
labels4.append(labels)
|
||||
segments4.extend(segments)
|
||||
|
||||
# Concat/clip labels
|
||||
labels4 = np.concatenate(labels4, 0)
|
||||
for x in (labels4[:, 1:], *segments4):
|
||||
np.clip(x, 0, 2 * s, out=x) # clip when using random_perspective()
|
||||
# img4, labels4 = replicate(img4, labels4) # replicate
|
||||
|
||||
# Augment
|
||||
img4, labels4, segments4 = copy_paste(img4, labels4, segments4, p=self.hyp["copy_paste"])
|
||||
img4, labels4, segments4 = random_perspective(img4,
|
||||
labels4,
|
||||
segments4,
|
||||
degrees=self.hyp["degrees"],
|
||||
translate=self.hyp["translate"],
|
||||
scale=self.hyp["scale"],
|
||||
shear=self.hyp["shear"],
|
||||
perspective=self.hyp["perspective"],
|
||||
border=self.mosaic_border) # border to remove
|
||||
return img4, labels4, segments4
|
||||
|
||||
@staticmethod
|
||||
def collate_fn(batch):
|
||||
img, label, path, shapes, masks = zip(*batch) # transposed
|
||||
batched_masks = torch.cat(masks, 0)
|
||||
for i, l in enumerate(label):
|
||||
l[:, 0] = i # add target image index for build_targets()
|
||||
return torch.stack(img, 0), torch.cat(label, 0), path, shapes, batched_masks
|
||||
|
||||
|
||||
def polygon2mask(img_size, polygons, color=1, downsample_ratio=1):
|
||||
"""
|
||||
Args:
|
||||
img_size (tuple): The image size.
|
||||
polygons (np.ndarray): [N, M], N is the number of polygons,
|
||||
M is the number of points(Be divided by 2).
|
||||
"""
|
||||
mask = np.zeros(img_size, dtype=np.uint8)
|
||||
polygons = np.asarray(polygons)
|
||||
polygons = polygons.astype(np.int32)
|
||||
shape = polygons.shape
|
||||
polygons = polygons.reshape(shape[0], -1, 2)
|
||||
cv2.fillPoly(mask, polygons, color=color)
|
||||
nh, nw = (img_size[0] // downsample_ratio, img_size[1] // downsample_ratio)
|
||||
# NOTE: fillPoly firstly then resize is trying the keep the same way
|
||||
# of loss calculation when mask-ratio=1.
|
||||
mask = cv2.resize(mask, (nw, nh))
|
||||
return mask
|
||||
|
||||
|
||||
def polygons2masks(img_size, polygons, color, downsample_ratio=1):
|
||||
"""
|
||||
Args:
|
||||
img_size (tuple): The image size.
|
||||
polygons (list[np.ndarray]): each polygon is [N, M],
|
||||
N is the number of polygons,
|
||||
M is the number of points(Be divided by 2).
|
||||
"""
|
||||
masks = []
|
||||
for si in range(len(polygons)):
|
||||
mask = polygon2mask(img_size, [polygons[si].reshape(-1)], color, downsample_ratio)
|
||||
masks.append(mask)
|
||||
return np.array(masks)
|
||||
|
||||
|
||||
def polygons2masks_overlap(img_size, segments, downsample_ratio=1):
|
||||
"""Return a (640, 640) overlap mask."""
|
||||
masks = np.zeros((img_size[0] // downsample_ratio, img_size[1] // downsample_ratio), dtype=np.uint8)
|
||||
areas = []
|
||||
ms = []
|
||||
for si in range(len(segments)):
|
||||
mask = polygon2mask(
|
||||
img_size,
|
||||
[segments[si].reshape(-1)],
|
||||
downsample_ratio=downsample_ratio,
|
||||
color=1,
|
||||
)
|
||||
ms.append(mask)
|
||||
areas.append(mask.sum())
|
||||
areas = np.asarray(areas)
|
||||
index = np.argsort(-areas)
|
||||
ms = np.array(ms)[index]
|
||||
for i in range(len(segments)):
|
||||
mask = ms[i] * (i + 1)
|
||||
masks = masks + mask
|
||||
masks = np.clip(masks, a_min=0, a_max=i + 1)
|
||||
return masks, index
|
||||
@@ -0,0 +1,120 @@
|
||||
import cv2
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def crop_mask(masks, boxes):
|
||||
"""
|
||||
"Crop" predicted masks by zeroing out everything not in the predicted bbox.
|
||||
Vectorized by Chong (thanks Chong).
|
||||
|
||||
Args:
|
||||
- masks should be a size [h, w, n] tensor of masks
|
||||
- boxes should be a size [n, 4] tensor of bbox coords in relative point form
|
||||
"""
|
||||
|
||||
n, h, w = masks.shape
|
||||
x1, y1, x2, y2 = torch.chunk(boxes[:, :, None], 4, 1) # x1 shape(1,1,n)
|
||||
r = torch.arange(w, device=masks.device, dtype=x1.dtype)[None, None, :] # rows shape(1,w,1)
|
||||
c = torch.arange(h, device=masks.device, dtype=x1.dtype)[None, :, None] # cols shape(h,1,1)
|
||||
|
||||
return masks * ((r >= x1) * (r < x2) * (c >= y1) * (c < y2))
|
||||
|
||||
|
||||
def process_mask_upsample(protos, masks_in, bboxes, shape):
|
||||
"""
|
||||
Crop after upsample.
|
||||
proto_out: [mask_dim, mask_h, mask_w]
|
||||
out_masks: [n, mask_dim], n is number of masks after nms
|
||||
bboxes: [n, 4], n is number of masks after nms
|
||||
shape:input_image_size, (h, w)
|
||||
|
||||
return: h, w, n
|
||||
"""
|
||||
|
||||
c, mh, mw = protos.shape # CHW
|
||||
masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw)
|
||||
masks = F.interpolate(masks[None], shape, mode='bilinear', align_corners=False)[0] # CHW
|
||||
masks = crop_mask(masks, bboxes) # CHW
|
||||
return masks.gt_(0.5)
|
||||
|
||||
|
||||
def process_mask(protos, masks_in, bboxes, shape, upsample=False):
|
||||
"""
|
||||
Crop before upsample.
|
||||
proto_out: [mask_dim, mask_h, mask_w]
|
||||
out_masks: [n, mask_dim], n is number of masks after nms
|
||||
bboxes: [n, 4], n is number of masks after nms
|
||||
shape:input_image_size, (h, w)
|
||||
|
||||
return: h, w, n
|
||||
"""
|
||||
|
||||
c, mh, mw = protos.shape # CHW
|
||||
ih, iw = shape
|
||||
masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw) # CHW
|
||||
|
||||
downsampled_bboxes = bboxes.clone()
|
||||
downsampled_bboxes[:, 0] *= mw / iw
|
||||
downsampled_bboxes[:, 2] *= mw / iw
|
||||
downsampled_bboxes[:, 3] *= mh / ih
|
||||
downsampled_bboxes[:, 1] *= mh / ih
|
||||
|
||||
masks = crop_mask(masks, downsampled_bboxes) # CHW
|
||||
if upsample:
|
||||
masks = F.interpolate(masks[None], shape, mode='bilinear', align_corners=False)[0] # CHW
|
||||
return masks.gt_(0.5)
|
||||
|
||||
|
||||
def scale_image(im1_shape, masks, im0_shape, ratio_pad=None):
|
||||
"""
|
||||
img1_shape: model input shape, [h, w]
|
||||
img0_shape: origin pic shape, [h, w, 3]
|
||||
masks: [h, w, num]
|
||||
"""
|
||||
# Rescale coordinates (xyxy) from im1_shape to im0_shape
|
||||
if ratio_pad is None: # calculate from im0_shape
|
||||
gain = min(im1_shape[0] / im0_shape[0], im1_shape[1] / im0_shape[1]) # gain = old / new
|
||||
pad = (im1_shape[1] - im0_shape[1] * gain) / 2, (im1_shape[0] - im0_shape[0] * gain) / 2 # wh padding
|
||||
else:
|
||||
pad = ratio_pad[1]
|
||||
top, left = int(pad[1]), int(pad[0]) # y, x
|
||||
bottom, right = int(im1_shape[0] - pad[1]), int(im1_shape[1] - pad[0])
|
||||
|
||||
if len(masks.shape) < 2:
|
||||
raise ValueError(f'"len of masks shape" should be 2 or 3, but got {len(masks.shape)}')
|
||||
masks = masks[top:bottom, left:right]
|
||||
# masks = masks.permute(2, 0, 1).contiguous()
|
||||
# masks = F.interpolate(masks[None], im0_shape[:2], mode='bilinear', align_corners=False)[0]
|
||||
# masks = masks.permute(1, 2, 0).contiguous()
|
||||
masks = cv2.resize(masks, (im0_shape[1], im0_shape[0]))
|
||||
|
||||
if len(masks.shape) == 2:
|
||||
masks = masks[:, :, None]
|
||||
return masks
|
||||
|
||||
|
||||
def mask_iou(mask1, mask2, eps=1e-7):
|
||||
"""
|
||||
mask1: [N, n] m1 means number of predicted objects
|
||||
mask2: [M, n] m2 means number of gt objects
|
||||
Note: n means image_w x image_h
|
||||
|
||||
return: masks iou, [N, M]
|
||||
"""
|
||||
intersection = torch.matmul(mask1, mask2.t()).clamp(0)
|
||||
union = (mask1.sum(1)[:, None] + mask2.sum(1)[None]) - intersection # (area1 + area2) - intersection
|
||||
return intersection / (union + eps)
|
||||
|
||||
|
||||
def masks_iou(mask1, mask2, eps=1e-7):
|
||||
"""
|
||||
mask1: [N, n] m1 means number of predicted objects
|
||||
mask2: [N, n] m2 means number of gt objects
|
||||
Note: n means image_w x image_h
|
||||
|
||||
return: masks iou, (N, )
|
||||
"""
|
||||
intersection = (mask1 * mask2).sum(1).clamp(0) # (N, )
|
||||
union = (mask1.sum(1) + mask2.sum(1))[None] - intersection # (area1 + area2) - intersection
|
||||
return intersection / (union + eps)
|
||||
@@ -0,0 +1,186 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from ..general import xywh2xyxy
|
||||
from ..loss import FocalLoss, smooth_BCE
|
||||
from ..metrics import bbox_iou
|
||||
from ..torch_utils import de_parallel
|
||||
from .general import crop_mask
|
||||
|
||||
|
||||
class ComputeLoss:
|
||||
# Compute losses
|
||||
def __init__(self, model, autobalance=False, overlap=False):
|
||||
self.sort_obj_iou = False
|
||||
self.overlap = overlap
|
||||
device = next(model.parameters()).device # get model device
|
||||
h = model.hyp # hyperparameters
|
||||
self.device = device
|
||||
|
||||
# Define criteria
|
||||
BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['cls_pw']], device=device))
|
||||
BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['obj_pw']], device=device))
|
||||
|
||||
# Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3
|
||||
self.cp, self.cn = smooth_BCE(eps=h.get('label_smoothing', 0.0)) # positive, negative BCE targets
|
||||
|
||||
# Focal loss
|
||||
g = h['fl_gamma'] # focal loss gamma
|
||||
if g > 0:
|
||||
BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)
|
||||
|
||||
m = de_parallel(model).model[-1] # Detect() module
|
||||
self.balance = {3: [4.0, 1.0, 0.4]}.get(m.nl, [4.0, 1.0, 0.25, 0.06, 0.02]) # P3-P7
|
||||
self.ssi = list(m.stride).index(16) if autobalance else 0 # stride 16 index
|
||||
self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = BCEcls, BCEobj, 1.0, h, autobalance
|
||||
self.na = m.na # number of anchors
|
||||
self.nc = m.nc # number of classes
|
||||
self.nl = m.nl # number of layers
|
||||
self.nm = m.nm # number of masks
|
||||
self.anchors = m.anchors
|
||||
self.device = device
|
||||
|
||||
def __call__(self, preds, targets, masks): # predictions, targets, model
|
||||
p, proto = preds
|
||||
bs, nm, mask_h, mask_w = proto.shape # batch size, number of masks, mask height, mask width
|
||||
lcls = torch.zeros(1, device=self.device)
|
||||
lbox = torch.zeros(1, device=self.device)
|
||||
lobj = torch.zeros(1, device=self.device)
|
||||
lseg = torch.zeros(1, device=self.device)
|
||||
tcls, tbox, indices, anchors, tidxs, xywhn = self.build_targets(p, targets) # targets
|
||||
|
||||
# Losses
|
||||
for i, pi in enumerate(p): # layer index, layer predictions
|
||||
b, a, gj, gi = indices[i] # image, anchor, gridy, gridx
|
||||
tobj = torch.zeros(pi.shape[:4], dtype=pi.dtype, device=self.device) # target obj
|
||||
|
||||
n = b.shape[0] # number of targets
|
||||
if n:
|
||||
pxy, pwh, _, pcls, pmask = pi[b, a, gj, gi].split((2, 2, 1, self.nc, nm), 1) # subset of predictions
|
||||
|
||||
# Box regression
|
||||
pxy = pxy.sigmoid() * 2 - 0.5
|
||||
pwh = (pwh.sigmoid() * 2) ** 2 * anchors[i]
|
||||
pbox = torch.cat((pxy, pwh), 1) # predicted box
|
||||
iou = bbox_iou(pbox, tbox[i], CIoU=True).squeeze() # iou(prediction, target)
|
||||
lbox += (1.0 - iou).mean() # iou loss
|
||||
|
||||
# Objectness
|
||||
iou = iou.detach().clamp(0).type(tobj.dtype)
|
||||
if self.sort_obj_iou:
|
||||
j = iou.argsort()
|
||||
b, a, gj, gi, iou = b[j], a[j], gj[j], gi[j], iou[j]
|
||||
if self.gr < 1:
|
||||
iou = (1.0 - self.gr) + self.gr * iou
|
||||
tobj[b, a, gj, gi] = iou # iou ratio
|
||||
|
||||
# Classification
|
||||
if self.nc > 1: # cls loss (only if multiple classes)
|
||||
t = torch.full_like(pcls, self.cn, device=self.device) # targets
|
||||
t[range(n), tcls[i]] = self.cp
|
||||
lcls += self.BCEcls(pcls, t) # BCE
|
||||
|
||||
# Mask regression
|
||||
if tuple(masks.shape[-2:]) != (mask_h, mask_w): # downsample
|
||||
masks = F.interpolate(masks[None], (mask_h, mask_w), mode="nearest")[0]
|
||||
marea = xywhn[i][:, 2:].prod(1) # mask width, height normalized
|
||||
mxyxy = xywh2xyxy(xywhn[i] * torch.tensor([mask_w, mask_h, mask_w, mask_h], device=self.device))
|
||||
for bi in b.unique():
|
||||
j = b == bi # matching index
|
||||
if self.overlap:
|
||||
mask_gti = torch.where(masks[bi][None] == tidxs[i][j].view(-1, 1, 1), 1.0, 0.0)
|
||||
else:
|
||||
mask_gti = masks[tidxs[i]][j]
|
||||
lseg += self.single_mask_loss(mask_gti, pmask[j], proto[bi], mxyxy[j], marea[j])
|
||||
|
||||
obji = self.BCEobj(pi[..., 4], tobj)
|
||||
lobj += obji * self.balance[i] # obj loss
|
||||
if self.autobalance:
|
||||
self.balance[i] = self.balance[i] * 0.9999 + 0.0001 / obji.detach().item()
|
||||
|
||||
if self.autobalance:
|
||||
self.balance = [x / self.balance[self.ssi] for x in self.balance]
|
||||
lbox *= self.hyp["box"]
|
||||
lobj *= self.hyp["obj"]
|
||||
lcls *= self.hyp["cls"]
|
||||
lseg *= self.hyp["box"] / bs
|
||||
|
||||
loss = lbox + lobj + lcls + lseg
|
||||
return loss * bs, torch.cat((lbox, lseg, lobj, lcls)).detach()
|
||||
|
||||
def single_mask_loss(self, gt_mask, pred, proto, xyxy, area):
|
||||
# Mask loss for one image
|
||||
pred_mask = (pred @ proto.view(self.nm, -1)).view(-1, *proto.shape[1:]) # (n,32) @ (32,80,80) -> (n,80,80)
|
||||
loss = F.binary_cross_entropy_with_logits(pred_mask, gt_mask, reduction="none")
|
||||
return (crop_mask(loss, xyxy).mean(dim=(1, 2)) / area).mean()
|
||||
|
||||
def build_targets(self, p, targets):
|
||||
# Build targets for compute_loss(), input targets(image,class,x,y,w,h)
|
||||
na, nt = self.na, targets.shape[0] # number of anchors, targets
|
||||
tcls, tbox, indices, anch, tidxs, xywhn = [], [], [], [], [], []
|
||||
gain = torch.ones(8, device=self.device) # normalized to gridspace gain
|
||||
ai = torch.arange(na, device=self.device).float().view(na, 1).repeat(1, nt) # same as .repeat_interleave(nt)
|
||||
if self.overlap:
|
||||
batch = p[0].shape[0]
|
||||
ti = []
|
||||
for i in range(batch):
|
||||
num = (targets[:, 0] == i).sum() # find number of targets of each image
|
||||
ti.append(torch.arange(num, device=self.device).float().view(1, num).repeat(na, 1) + 1) # (na, num)
|
||||
ti = torch.cat(ti, 1) # (na, nt)
|
||||
else:
|
||||
ti = torch.arange(nt, device=self.device).float().view(1, nt).repeat(na, 1)
|
||||
targets = torch.cat((targets.repeat(na, 1, 1), ai[..., None], ti[..., None]), 2) # append anchor indices
|
||||
|
||||
g = 0.5 # bias
|
||||
off = torch.tensor(
|
||||
[
|
||||
[0, 0],
|
||||
[1, 0],
|
||||
[0, 1],
|
||||
[-1, 0],
|
||||
[0, -1], # j,k,l,m
|
||||
# [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm
|
||||
],
|
||||
device=self.device).float() * g # offsets
|
||||
|
||||
for i in range(self.nl):
|
||||
anchors, shape = self.anchors[i], p[i].shape
|
||||
gain[2:6] = torch.tensor(shape)[[3, 2, 3, 2]] # xyxy gain
|
||||
|
||||
# Match targets to anchors
|
||||
t = targets * gain # shape(3,n,7)
|
||||
if nt:
|
||||
# Matches
|
||||
r = t[..., 4:6] / anchors[:, None] # wh ratio
|
||||
j = torch.max(r, 1 / r).max(2)[0] < self.hyp['anchor_t'] # compare
|
||||
# j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t'] # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2))
|
||||
t = t[j] # filter
|
||||
|
||||
# Offsets
|
||||
gxy = t[:, 2:4] # grid xy
|
||||
gxi = gain[[2, 3]] - gxy # inverse
|
||||
j, k = ((gxy % 1 < g) & (gxy > 1)).T
|
||||
l, m = ((gxi % 1 < g) & (gxi > 1)).T
|
||||
j = torch.stack((torch.ones_like(j), j, k, l, m))
|
||||
t = t.repeat((5, 1, 1))[j]
|
||||
offsets = (torch.zeros_like(gxy)[None] + off[:, None])[j]
|
||||
else:
|
||||
t = targets[0]
|
||||
offsets = 0
|
||||
|
||||
# Define
|
||||
bc, gxy, gwh, at = t.chunk(4, 1) # (image, class), grid xy, grid wh, anchors
|
||||
(a, tidx), (b, c) = at.long().T, bc.long().T # anchors, image, class
|
||||
gij = (gxy - offsets).long()
|
||||
gi, gj = gij.T # grid indices
|
||||
|
||||
# Append
|
||||
indices.append((b, a, gj.clamp_(0, shape[2] - 1), gi.clamp_(0, shape[3] - 1))) # image, anchor, grid
|
||||
tbox.append(torch.cat((gxy - gij, gwh), 1)) # box
|
||||
anch.append(anchors[a]) # anchors
|
||||
tcls.append(c) # class
|
||||
tidxs.append(tidx)
|
||||
xywhn.append(torch.cat((gxy, gwh), 1) / gain[2:6]) # xywh normalized
|
||||
|
||||
return tcls, tbox, indices, anch, tidxs, xywhn
|
||||
@@ -0,0 +1,210 @@
|
||||
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
|
||||
"""
|
||||
Model validation metrics
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
from ..metrics import ap_per_class
|
||||
|
||||
|
||||
def fitness(x):
|
||||
# Model fitness as a weighted combination of metrics
|
||||
w = [0.0, 0.0, 0.1, 0.9, 0.0, 0.0, 0.1, 0.9]
|
||||
return (x[:, :8] * w).sum(1)
|
||||
|
||||
|
||||
def ap_per_class_box_and_mask(
|
||||
tp_m,
|
||||
tp_b,
|
||||
conf,
|
||||
pred_cls,
|
||||
target_cls,
|
||||
plot=False,
|
||||
save_dir=".",
|
||||
names=(),
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
tp_b: tp of boxes.
|
||||
tp_m: tp of masks.
|
||||
other arguments see `func: ap_per_class`.
|
||||
"""
|
||||
results_boxes = ap_per_class(tp_b,
|
||||
conf,
|
||||
pred_cls,
|
||||
target_cls,
|
||||
plot=plot,
|
||||
save_dir=save_dir,
|
||||
names=names,
|
||||
prefix="Box")[2:]
|
||||
results_masks = ap_per_class(tp_m,
|
||||
conf,
|
||||
pred_cls,
|
||||
target_cls,
|
||||
plot=plot,
|
||||
save_dir=save_dir,
|
||||
names=names,
|
||||
prefix="Mask")[2:]
|
||||
|
||||
results = {
|
||||
"boxes": {
|
||||
"p": results_boxes[0],
|
||||
"r": results_boxes[1],
|
||||
"ap": results_boxes[3],
|
||||
"f1": results_boxes[2],
|
||||
"ap_class": results_boxes[4]},
|
||||
"masks": {
|
||||
"p": results_masks[0],
|
||||
"r": results_masks[1],
|
||||
"ap": results_masks[3],
|
||||
"f1": results_masks[2],
|
||||
"ap_class": results_masks[4]}}
|
||||
return results
|
||||
|
||||
|
||||
class Metric:
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.p = [] # (nc, )
|
||||
self.r = [] # (nc, )
|
||||
self.f1 = [] # (nc, )
|
||||
self.all_ap = [] # (nc, 10)
|
||||
self.ap_class_index = [] # (nc, )
|
||||
|
||||
@property
|
||||
def ap50(self):
|
||||
"""AP@0.5 of all classes.
|
||||
Return:
|
||||
(nc, ) or [].
|
||||
"""
|
||||
return self.all_ap[:, 0] if len(self.all_ap) else []
|
||||
|
||||
@property
|
||||
def ap(self):
|
||||
"""AP@0.5:0.95
|
||||
Return:
|
||||
(nc, ) or [].
|
||||
"""
|
||||
return self.all_ap.mean(1) if len(self.all_ap) else []
|
||||
|
||||
@property
|
||||
def mp(self):
|
||||
"""mean precision of all classes.
|
||||
Return:
|
||||
float.
|
||||
"""
|
||||
return self.p.mean() if len(self.p) else 0.0
|
||||
|
||||
@property
|
||||
def mr(self):
|
||||
"""mean recall of all classes.
|
||||
Return:
|
||||
float.
|
||||
"""
|
||||
return self.r.mean() if len(self.r) else 0.0
|
||||
|
||||
@property
|
||||
def map50(self):
|
||||
"""Mean AP@0.5 of all classes.
|
||||
Return:
|
||||
float.
|
||||
"""
|
||||
return self.all_ap[:, 0].mean() if len(self.all_ap) else 0.0
|
||||
|
||||
@property
|
||||
def map(self):
|
||||
"""Mean AP@0.5:0.95 of all classes.
|
||||
Return:
|
||||
float.
|
||||
"""
|
||||
return self.all_ap.mean() if len(self.all_ap) else 0.0
|
||||
|
||||
def mean_results(self):
|
||||
"""Mean of results, return mp, mr, map50, map"""
|
||||
return (self.mp, self.mr, self.map50, self.map)
|
||||
|
||||
def class_result(self, i):
|
||||
"""class-aware result, return p[i], r[i], ap50[i], ap[i]"""
|
||||
return (self.p[i], self.r[i], self.ap50[i], self.ap[i])
|
||||
|
||||
def get_maps(self, nc):
|
||||
maps = np.zeros(nc) + self.map
|
||||
for i, c in enumerate(self.ap_class_index):
|
||||
maps[c] = self.ap[i]
|
||||
return maps
|
||||
|
||||
def update(self, results):
|
||||
"""
|
||||
Args:
|
||||
results: tuple(p, r, ap, f1, ap_class)
|
||||
"""
|
||||
p, r, all_ap, f1, ap_class_index = results
|
||||
self.p = p
|
||||
self.r = r
|
||||
self.all_ap = all_ap
|
||||
self.f1 = f1
|
||||
self.ap_class_index = ap_class_index
|
||||
|
||||
|
||||
class Metrics:
|
||||
"""Metric for boxes and masks."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.metric_box = Metric()
|
||||
self.metric_mask = Metric()
|
||||
|
||||
def update(self, results):
|
||||
"""
|
||||
Args:
|
||||
results: Dict{'boxes': Dict{}, 'masks': Dict{}}
|
||||
"""
|
||||
self.metric_box.update(list(results["boxes"].values()))
|
||||
self.metric_mask.update(list(results["masks"].values()))
|
||||
|
||||
def mean_results(self):
|
||||
return self.metric_box.mean_results() + self.metric_mask.mean_results()
|
||||
|
||||
def class_result(self, i):
|
||||
return self.metric_box.class_result(i) + self.metric_mask.class_result(i)
|
||||
|
||||
def get_maps(self, nc):
|
||||
return self.metric_box.get_maps(nc) + self.metric_mask.get_maps(nc)
|
||||
|
||||
@property
|
||||
def ap_class_index(self):
|
||||
# boxes and masks have the same ap_class_index
|
||||
return self.metric_box.ap_class_index
|
||||
|
||||
|
||||
KEYS = [
|
||||
"train/box_loss",
|
||||
"train/seg_loss", # train loss
|
||||
"train/obj_loss",
|
||||
"train/cls_loss",
|
||||
"metrics/precision(B)",
|
||||
"metrics/recall(B)",
|
||||
"metrics/mAP_0.5(B)",
|
||||
"metrics/mAP_0.5:0.95(B)", # metrics
|
||||
"metrics/precision(M)",
|
||||
"metrics/recall(M)",
|
||||
"metrics/mAP_0.5(M)",
|
||||
"metrics/mAP_0.5:0.95(M)", # metrics
|
||||
"val/box_loss",
|
||||
"val/seg_loss", # val loss
|
||||
"val/obj_loss",
|
||||
"val/cls_loss",
|
||||
"x/lr0",
|
||||
"x/lr1",
|
||||
"x/lr2",]
|
||||
|
||||
BEST_KEYS = [
|
||||
"best/epoch",
|
||||
"best/precision(B)",
|
||||
"best/recall(B)",
|
||||
"best/mAP_0.5(B)",
|
||||
"best/mAP_0.5:0.95(B)",
|
||||
"best/precision(M)",
|
||||
"best/recall(M)",
|
||||
"best/mAP_0.5(M)",
|
||||
"best/mAP_0.5:0.95(M)",]
|
||||
@@ -0,0 +1,143 @@
|
||||
import contextlib
|
||||
import math
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import torch
|
||||
|
||||
from .. import threaded
|
||||
from ..general import xywh2xyxy
|
||||
from ..plots import Annotator, colors
|
||||
|
||||
|
||||
@threaded
|
||||
def plot_images_and_masks(images, targets, masks, paths=None, fname='images.jpg', names=None):
|
||||
# Plot image grid with labels
|
||||
if isinstance(images, torch.Tensor):
|
||||
images = images.cpu().float().numpy()
|
||||
if isinstance(targets, torch.Tensor):
|
||||
targets = targets.cpu().numpy()
|
||||
if isinstance(masks, torch.Tensor):
|
||||
masks = masks.cpu().numpy().astype(int)
|
||||
|
||||
max_size = 1920 # max image size
|
||||
max_subplots = 16 # max image subplots, i.e. 4x4
|
||||
bs, _, h, w = images.shape # batch size, _, height, width
|
||||
bs = min(bs, max_subplots) # limit plot images
|
||||
ns = np.ceil(bs ** 0.5) # number of subplots (square)
|
||||
if np.max(images[0]) <= 1:
|
||||
images *= 255 # de-normalise (optional)
|
||||
|
||||
# Build Image
|
||||
mosaic = np.full((int(ns * h), int(ns * w), 3), 255, dtype=np.uint8) # init
|
||||
for i, im in enumerate(images):
|
||||
if i == max_subplots: # if last batch has fewer images than we expect
|
||||
break
|
||||
x, y = int(w * (i // ns)), int(h * (i % ns)) # block origin
|
||||
im = im.transpose(1, 2, 0)
|
||||
mosaic[y:y + h, x:x + w, :] = im
|
||||
|
||||
# Resize (optional)
|
||||
scale = max_size / ns / max(h, w)
|
||||
if scale < 1:
|
||||
h = math.ceil(scale * h)
|
||||
w = math.ceil(scale * w)
|
||||
mosaic = cv2.resize(mosaic, tuple(int(x * ns) for x in (w, h)))
|
||||
|
||||
# Annotate
|
||||
fs = int((h + w) * ns * 0.01) # font size
|
||||
annotator = Annotator(mosaic, line_width=round(fs / 10), font_size=fs, pil=True, example=names)
|
||||
for i in range(i + 1):
|
||||
x, y = int(w * (i // ns)), int(h * (i % ns)) # block origin
|
||||
annotator.rectangle([x, y, x + w, y + h], None, (255, 255, 255), width=2) # borders
|
||||
if paths:
|
||||
annotator.text((x + 5, y + 5 + h), text=Path(paths[i]).name[:40], txt_color=(220, 220, 220)) # filenames
|
||||
if len(targets) > 0:
|
||||
idx = targets[:, 0] == i
|
||||
ti = targets[idx] # image targets
|
||||
|
||||
boxes = xywh2xyxy(ti[:, 2:6]).T
|
||||
classes = ti[:, 1].astype('int')
|
||||
labels = ti.shape[1] == 6 # labels if no conf column
|
||||
conf = None if labels else ti[:, 6] # check for confidence presence (label vs pred)
|
||||
|
||||
if boxes.shape[1]:
|
||||
if boxes.max() <= 1.01: # if normalized with tolerance 0.01
|
||||
boxes[[0, 2]] *= w # scale to pixels
|
||||
boxes[[1, 3]] *= h
|
||||
elif scale < 1: # absolute coords need scale if image scales
|
||||
boxes *= scale
|
||||
boxes[[0, 2]] += x
|
||||
boxes[[1, 3]] += y
|
||||
for j, box in enumerate(boxes.T.tolist()):
|
||||
cls = classes[j]
|
||||
color = colors(cls)
|
||||
cls = names[cls] if names else cls
|
||||
if labels or conf[j] > 0.25: # 0.25 conf thresh
|
||||
label = f'{cls}' if labels else f'{cls} {conf[j]:.1f}'
|
||||
annotator.box_label(box, label, color=color)
|
||||
|
||||
# Plot masks
|
||||
if len(masks):
|
||||
if masks.max() > 1.0: # mean that masks are overlap
|
||||
image_masks = masks[[i]] # (1, 640, 640)
|
||||
nl = len(ti)
|
||||
index = np.arange(nl).reshape(nl, 1, 1) + 1
|
||||
image_masks = np.repeat(image_masks, nl, axis=0)
|
||||
image_masks = np.where(image_masks == index, 1.0, 0.0)
|
||||
else:
|
||||
image_masks = masks[idx]
|
||||
|
||||
im = np.asarray(annotator.im).copy()
|
||||
for j, box in enumerate(boxes.T.tolist()):
|
||||
if labels or conf[j] > 0.25: # 0.25 conf thresh
|
||||
color = colors(classes[j])
|
||||
mh, mw = image_masks[j].shape
|
||||
if mh != h or mw != w:
|
||||
mask = image_masks[j].astype(np.uint8)
|
||||
mask = cv2.resize(mask, (w, h))
|
||||
mask = mask.astype(np.bool)
|
||||
else:
|
||||
mask = image_masks[j].astype(np.bool)
|
||||
with contextlib.suppress(Exception):
|
||||
im[y:y + h, x:x + w, :][mask] = im[y:y + h, x:x + w, :][mask] * 0.4 + np.array(color) * 0.6
|
||||
annotator.fromarray(im)
|
||||
annotator.im.save(fname) # save
|
||||
|
||||
|
||||
def plot_results_with_masks(file="path/to/results.csv", dir="", best=True):
|
||||
# Plot training results.csv. Usage: from utils.plots import *; plot_results('path/to/results.csv')
|
||||
save_dir = Path(file).parent if file else Path(dir)
|
||||
fig, ax = plt.subplots(2, 8, figsize=(18, 6), tight_layout=True)
|
||||
ax = ax.ravel()
|
||||
files = list(save_dir.glob("results*.csv"))
|
||||
assert len(files), f"No results.csv files found in {save_dir.resolve()}, nothing to plot."
|
||||
for f in files:
|
||||
try:
|
||||
data = pd.read_csv(f)
|
||||
index = np.argmax(0.9 * data.values[:, 8] + 0.1 * data.values[:, 7] + 0.9 * data.values[:, 12] +
|
||||
0.1 * data.values[:, 11])
|
||||
s = [x.strip() for x in data.columns]
|
||||
x = data.values[:, 0]
|
||||
for i, j in enumerate([1, 2, 3, 4, 5, 6, 9, 10, 13, 14, 15, 16, 7, 8, 11, 12]):
|
||||
y = data.values[:, j]
|
||||
# y[y == 0] = np.nan # don't show zero values
|
||||
ax[i].plot(x, y, marker=".", label=f.stem, linewidth=2, markersize=2)
|
||||
if best:
|
||||
# best
|
||||
ax[i].scatter(index, y[index], color="r", label=f"best:{index}", marker="*", linewidth=3)
|
||||
ax[i].set_title(s[j] + f"\n{round(y[index], 5)}")
|
||||
else:
|
||||
# last
|
||||
ax[i].scatter(x[-1], y[-1], color="r", label="last", marker="*", linewidth=3)
|
||||
ax[i].set_title(s[j] + f"\n{round(y[-1], 5)}")
|
||||
# if j in [8, 9, 10]: # share train and val loss y axes
|
||||
# ax[i].get_shared_y_axes().join(ax[i], ax[i - 5])
|
||||
except Exception as e:
|
||||
print(f"Warning: Plotting error for {f}: {e}")
|
||||
ax[1].legend()
|
||||
fig.savefig(save_dir / "results.png", dpi=200)
|
||||
plt.close()
|
||||
Reference in new issue
Block a user