def infer(loader, net, device):
net.eval()
metric = ConfusionMatrix(len(loader.dataset.classes) - 1)
for i, (x, gt) in enumerate(loader):
x = x.to(device, non_blocking=True)
gt = gt.squeeze(1).to(device, dtype=torch.long, non_blocking=True)
pred = net(x)["fine"].argmax(1)
metric.update(pred, gt)
mIoU = metric.mIoU()
logging.info(f"[Infer] mIOU : {mIoU}")
return mIoU
def __init__(self,
criterion,
optimizer,
n_class,
size0,
size_g,
size_p,
n,
sub_batch_size=6,
mode=1,
lamb_fmreg=0.15):
self.criterion = criterion
self.optimizer = optimizer
self.metrics_global = ConfusionMatrix(n_class)
self.metrics_local = ConfusionMatrix(n_class)
self.metrics = ConfusionMatrix(n_class)
self.n_class = n_class
self.size0 = size0
self.size_g = size_g
self.size_p = size_p
self.n = n
self.sub_batch_size = sub_batch_size
self.mode = mode
self.lamb_fmreg = lamb_fmreg
self.ratio = float(size_p[0]) / size0[0]
self.step = (size0[0] - size_p[0]) // (n - 1)
self.template, self.coordinates = template_patch2global(
size0, size_p, n, self.step)
def __init__(self,
n_class,
size_g,
size_p,
sub_batch_size=6,
mode=1,
test=False):
self.metrics_global = ConfusionMatrix(n_class)
self.metrics_local = ConfusionMatrix(n_class)
self.metrics = ConfusionMatrix(n_class)
self.n_class = n_class
# self.size0 = size0
self.size_g = size_g
self.size_p = size_p
# self.n = sub_batch_size
self.sub_batch_size = sub_batch_size
self.mode = mode
self.test = test
# self.ratio = float(size_p[0]) / size0[0]
# self.step = (size0[0] - size_p[0]) // (self.n - 1)
# self.template, self.coordinates = template_patch2global(size0, size_p, self.n, self.step)
if test:
self.flip_range = [False, True]
self.rotate_range = [0, 1, 2, 3]
else:
self.flip_range = [False]
self.rotate_range = [0]
def __init__(self, n_class, sub_batchsize, mode=1, test=False):
self.metrics_global = ConfusionMatrix(n_class)
self.metrics_local = ConfusionMatrix(n_class)
self.metrics = ConfusionMatrix(n_class)
self.n_class = n_class
self.sub_batchsize = sub_batchsize
self.mode = mode
self.test = test
def __init__(self, cf,
test=None,
reference=None,
labels=None,
metrics=None,
advanced_metrics=None,
nan_for_nonexisting=True):
self.cf = cf
self.test = None
self.reference = None
self.confusion_matrix = ConfusionMatrix()
self.labels = None
self.nan_for_nonexisting = nan_for_nonexisting
self.result = None
self.metrics = []
if metrics is None:
for m in self.default_metrics:
self.metrics.append(m)
else:
for m in metrics:
self.metrics.append(m)
self.advanced_metrics = []
if advanced_metrics is None:
for m in self.default_advanced_metrics:
self.advanced_metrics.append(m)
else:
for m in advanced_metrics:
self.advanced_metrics.append(m)
self.set_reference(reference)
self.set_test(test)
if labels is not None:
self.set_labels(labels)
else:
if test is not None and reference is not None:
self.gen_labels()
def __init__(self, criterion, optimizer, n_class, sub_batchsize, mode=1, fmreg=0.15):
self.criterion = criterion
self.optimizer = optimizer
self.metrics_global = ConfusionMatrix(n_class)
self.metrics_local = ConfusionMatrix(n_class)
self.metrics = ConfusionMatrix(n_class)
self.n_class = n_class
self.sub_batchsize = sub_batchsize
self.mode = mode
self.fmreg = fmreg # regulization item
def __init__(self, criterion, optimizer, n_class, size_g, size_p, sub_batch_size=6, mode=1, lamb_fmreg=0.15):
self.criterion = criterion
self.optimizer = optimizer
self.metrics_global = ConfusionMatrix(n_class)
self.metrics_local = ConfusionMatrix(n_class)
self.metrics = ConfusionMatrix(n_class)
self.n_class = n_class
self.size_g = size_g
self.size_p = size_p
self.sub_batch_size = sub_batch_size
self.mode = mode
self.lamb_fmreg = lamb_fmreg
def __init__(self, n_class, size_g, size_p, sub_batch_size=6, mode=1, test=False):
self.metrics_global = ConfusionMatrix(n_class)
self.metrics_local = ConfusionMatrix(n_class)
self.metrics = ConfusionMatrix(n_class)
self.n_class = n_class
self.size_g = size_g
self.size_p = size_p
self.sub_batch_size = sub_batch_size
self.mode = mode
self.test = test
if test:
self.flip_range = [False, True]
self.rotate_range = [0, 1, 2, 3]
else:
self.flip_range = [False]
self.rotate_range = [0]
def test(data,
weights=None,
batch_size=32,
imgsz=640,
conf_thres=-1, # not for NMS
iou_thres=0.25, # not for NMS
save_json=False,
single_cls=False,
augment=False,
verbose=False,
model=None,
dataloader=None,
save_dir=Path(''), # for saving images
save_txt=False, # for auto-labelling
save_hybrid=False, # for hybrid auto-labelling
save_conf=False, # save auto-label confidences
plots=True,
wandb_logger=None,
compute_loss=None,
half_precision=True, # dsv
is_coco=False,
max_by_class=True,
opt=None):
print_size, print_batches = 640, 3
log_errors = -1
if opt is not None:
print_size = opt.print_size
print_batches = opt.print_batches
max_by_class = opt.max_by_class
log_errors = opt.log_errors
ct = ast.literal_eval(opt.ct)
if len(ct) == 0:
ct = None
else:
ct = None
if print_batches < 0:
print_batches = 1000
# Initialize/load model and set device
training = model is not None
if training: # i.e. called by train.py
device = next(model.parameters()).device # get model device
else: # called directly
set_logging()
# device = select_device(opt.device, batch_size=batch_size)
device = select_device(opt.device, batch_size=2)
# Directories
save_dir = increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok) # increment run
(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
gs = max(int(model.stride.max()), 32) # grid size (max stride)
imgsz = check_img_size(imgsz, s=gs) # check img_size
# Half
half = device.type != 'cpu' and half_precision # half precision only supported on CUDA
if half:
model.half()
# Configure
model.eval()
if isinstance(data, str):
is_coco = data.endswith('coco.yaml')
with open(data) as f:
data = yaml.safe_load(f)
check_dataset(data) # check
nc = 1 if single_cls else int(data['nc']) # number of classes
# iouv = torch.linspace(0.5, 0.95, 10).to(device) # iou vector for [email protected]:0.95
iouv = torch.arange(iou_thres, 1, 0.05).to(device) # iou_thres : 0.95 : 0.05
niou = iouv.numel()
# Logging
log_imgs = 0
if wandb_logger and wandb_logger.wandb:
log_imgs = min(wandb_logger.log_imgs, 100)
# Dataloader
if not training:
if device.type != 'cpu':
model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once
task = opt.task if opt.task in ('train', 'val', 'test') else 'val' # path to train/val/test images
dataloader = create_dataloader(data[task], imgsz, batch_size, gs, opt, pad=0.5, rect=True, training=False,
prefix=colorstr(f'{task}: '))[0]
seen = 0
confusion_matrix = ConfusionMatrix(nc=nc)
names = data['names'] # names = {k: v for k, v in enumerate(model.names if hasattr(model, 'names') else model.module.names)}
names_dict = {k: v for k, v in enumerate(names)}
coco91class = coco80_to_coco91_class()
fmt = '%{}s'.format(2 + max([len(s) for s in names]))
s = (fmt + '%12s' * 7) % ('Class', 'Images', 'Targets', 'P', 'R', 'F1', '[email protected]', '[email protected]:.95')
p, r, f1, mp, mr, map50, map, t0, t1, mf1 = 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.
loss = torch.zeros(3, device=device)
jdict, stats, ap, ap_class, wandb_images = [], [], [], [], []
error_log = []
for batch_i, (imgs, targets, paths, shapes) in enumerate(dataloader):
imgs = imgs.to(device, non_blocking=True)
imgs = imgs.half() if half else imgs.float() # uint8 to fp16/32
imgs /= 255.0 # 0 - 255 to 0.0 - 1.0
targets = targets.to(device)
nb, _, height, width = imgs.shape # batch size, channels, height, width
# Run model
t = time_synchronized()
inf_out, train_out = model(imgs, augment=augment) # inference and training outputs
t0 += time_synchronized() - t
# Compute loss
if compute_loss:
loss += compute_loss([x.float() for x in train_out], targets)[1][:3] # box, obj, cls
# Run NMS
targets[:, 2:] *= torch.Tensor([width, height, width, height]).to(device) # to pixels
lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else [] # for autolabelling
t = time_synchronized()
output = non_max_suppression(inf_out, labels=lb, multi_label=False, agnostic=True) # , conf_thres=conf_thres, iou_thres=iou_thres)
t1 += time_synchronized() - t
# pfunc(f'test_batch_size=={len(output)}')
# Statistics per image
idx = []
fn,fp = 0,0
for si, pred in enumerate(output):
labels = targets[targets[:, 0] == si, 1:]
# Dims of all target boxes (in pixels)
target_dims = targets[targets[:, 0] == si, -2:]
nl = len(labels)
tcls = labels[:, 0].tolist() if nl else [] # target class
path = Path(paths[si])
seen += 1
if len(pred) == 0:
idx.append(None)
if nl:
stats.append((torch.zeros(0, niou, dtype=torch.bool), torch.Tensor(), torch.Tensor(), tcls))
continue
# Filter out-of-frame predictions (in padding)
gain, img_box = shapes[si][1][0][0], None
if gain == 1:
img_box = torch.zeros([1, 4])
img_box[0, :2] = torch.FloatTensor(shapes[si][1][1])
img_box[0, 2:] = torch.FloatTensor(shapes[si][0]) + torch.FloatTensor(shapes[si][1][1])
img_box = img_box.to(device)
io2s = box_io2(img_box, pred[:, :4])
k = (io2s > 0.95).nonzero(as_tuple=True)[1]
pred = pred[k, :]
idx.append(k)
else:
idx.append(None)
# Predictions
if single_cls:
pred[:, 5] = 0
predn = pred.clone()
scale_coords(imgs[si].shape[1:], predn[:, :4], shapes[si][0], shapes[si][1]) # native-space pred
# Dims of all pred boxes (in pixels)
pred_target_dims = torch.zeros(pred.shape[0], 4, dtype=torch.float32, device=device)
pred_target_dims[:, :2] = pred[:, 2:4] - pred[:, :2]
# Append to text file
if save_txt:
gn = torch.tensor(shapes[si][0])[[1, 0, 1, 0]] # normalization gain whwh
for *xyxy, conf, cls in predn.tolist():
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # label format
with open(save_dir / 'labels' / (path.stem + '.txt'), 'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n')
# W&B logging - Media Panel Plots
if len(wandb_images) < log_imgs and wandb_logger.current_epoch > 0: # Check for test operation
if wandb_logger.current_epoch % wandb_logger.bbox_interval == 0:
box_data = [{"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]},
"class_id": int(cls),
"box_caption": "%s %.3f" % (names[int(cls)], conf),
"scores": {"class_score": conf},
"domain": "pixel"} for *xyxy, conf, cls in pred.tolist()]
boxes = {"predictions": {"box_data": box_data, "class_labels": names_dict}} # inference-space
wandb_images.append(wandb_logger.wandb.Image(imgs[si], boxes=boxes, caption=path.name))
wandb_logger.log_training_progress(predn, path, names) if wandb_logger and wandb_logger.wandb_run else None
# Append to pycocotools JSON dictionary
if save_json:
# [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ...
image_id = int(path.stem) if path.stem.isnumeric() else path.stem
box = xyxy2xywh(predn[:, :4]) # xywh
box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner
for p, b in zip(pred.tolist(), box.tolist()):
jdict.append({'image_id': image_id,
'category_id': coco91class[int(p[5])] if is_coco else int(p[5]),
'bbox': [round(x, 3) for x in b],
'score': round(p[4], 5)})
# Assign all predictions as incorrect
correct = torch.zeros(pred.shape[0], niou, dtype=torch.bool, device=device)
if nl:
detected = [] # target indices
tcls_tensor = labels[:, 0]
# target boxes
tbox = xywh2xyxy(labels[:, 1:5])
scale_coords(imgs[si].shape[1:], tbox, shapes[si][0], shapes[si][1]) # native-space labels
if plots:
confusion_matrix.process_batch(predn, torch.cat((labels[:, 0:1], tbox), 1))
# Per target class
for cls in torch.unique(tcls_tensor):
ti = (cls == tcls_tensor).nonzero(as_tuple=False).view(-1) # target indices
pi = (cls == pred[:, 5]).nonzero(as_tuple=False).view(-1) # prediction indices
# Search for detections
if pi.shape[0]:
# Prediction to target ious
ious, i = box_iou(predn[pi, :4], tbox[ti]).max(1) # best ious, indices
# Append detections
detected_set = set()
for j in (ious > iouv[0]).nonzero(as_tuple=False):
d = ti[i[j]] # detected target... index into target array....
if d.item() not in detected_set:
detected_set.add(d.item())
detected.append(d)
correct[pi[j]] = ious[j] > iouv # iou_thres is 1xn
k = pi[j] # index into pred array.....
pred_target_dims[k, 2:] = target_dims[d]
if len(detected) == nl: # all targets already located in image
break
# Append statistics (correct, conf, pcls, tcls)
stats.append((correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls))
# stats.append((correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls, pred_target_dims.cpu(), target_dims.cpu()))
# FP/FN counts (per image)
if log_errors > -1:
corr = correct.cpu().numpy()
idx_conf = pred[:, 4].cpu() > conf_thres
if idx_conf.cpu().numpy().mean() < 1:
corr = corr[idx_conf]
tp = corr[:, 0].sum()
fp, fn = len(corr) - tp, len(tcls) - tp
if fp + fn > log_errors:
error_log.append(path.stem)
# Plot images
if plots and batch_i < print_batches:
prefix = Path(paths[0]).stem if batch_size == 1 else f'test_batch{batch_i}'
f1 = save_dir / f'{prefix}_labels.jpg' # labels
thread1 = Thread(target=plot_images, args=(imgs, targets, paths, f1, names, print_size), daemon=True)
f2 = save_dir / f'{prefix}_pred.jpg' # predictions
thread2 = Thread(target=plot_images, args=(imgs, output_to_target(output, idx), paths, f2, names, print_size), daemon=True)
thread1.start()
thread1.join()
thread2.start()
thread2.join()
##################################
# show FP and FN detections...
if log_errors > -1 and batch_size == 1:
# fn boxes...
if fn > 0:
idx_fn = np.ones(len(targets))
for d in detected:
idx_fn[d.item()] = 0
idx_fn = np.nonzero(idx_fn)[0]
f = save_dir / f'{prefix}_fn.jpg' # labels
thread = Thread(target=plot_images, args=(imgs, targets[idx_fn], paths, f, names, print_size, 16, True), daemon=True)
thread.start()
thread.join()
# fp boxes...
if fp > 0:
idx_fp = ~corr[:, 0]
f = save_dir / f'{prefix}_fp.jpg' # labels
thread = Thread(target=plot_images, args=(imgs, output_to_target(output, idx, idx_fp, idx_conf), paths, f, names, print_size, 16, True), daemon=True)
thread.start()
thread.join()
# Compute statistics
stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy
conf_best = -1
if len(stats) and stats[0].any():
mp, mr, map50, map, mf1, ap_class, conf_best, nt, (p, r, ap50, ap, f1, cc) = ap_per_class(*stats, plot=plots, save_dir=save_dir, names=names,
ct=ct, max_by_class=max_by_class, conf_thres=conf_thres)
else:
nt = torch.zeros(1)
# Print results
pfunc('------------------------------------------- Validation Set -----------------------------------------------')
pfunc(s)
pf = fmt + '%12.3g' * 7 # print format
# Print results per class
if (verbose or nc < 50) and nc > 1 and len(stats):
for i, c in enumerate(ap_class):
pfunc(pf % (names[c], seen, nt[c], p[i], r[i], f1[i], ap50[i], ap[i])) # log
# Print averages
if nc > 1:
pfunc('')
pfunc(pf % ('AVG', seen, nt.sum(), p.mean(), r.mean(), f1.mean(), ap50.mean(), ap.mean())) # unweighted average
ss = 'WEIGHTED AVG' if nc > 1 else names[0]
pfunc(pf % (ss, seen, nt.sum(), mp, mr, mf1, map50, map)) # weighted average (if nc>1)
if conf_best > -1:
pfunc('\nOptimal Confidence Threshold: {0:0.3f}'.format(conf_best)) # log
if max_by_class:
pfunc('Optimal Confidence Thresholds (Per-Class): {}'.format(list(cc.round(3)))) # log
# Print speeds
t = tuple(x / seen * 1E3 for x in (t0, t1, t0 + t1)) + (imgsz, imgsz, batch_size) # tuple
if not training:
print('Speed: %.1f/%.1f/%.1f ms inference/NMS/total per %gx%g image at batch-size %g' % t)
# Plots
if plots:
confusion_matrix.plot(save_dir=save_dir, names=names)
if wandb_logger and wandb_logger.wandb:
val_batches = [wandb_logger.wandb.Image(str(f), caption=f.name) for f in sorted(save_dir.glob('test*.jpg'))]
wandb_logger.log({"Validation": val_batches})
if wandb_images:
wandb_logger.log({"Bounding Box Debugger/Images": wandb_images})
# Save JSON
if save_json and len(jdict):
w = Path(weights[0] if isinstance(weights, list) else weights).stem if weights is not None else '' # weights
anno_json = '../coco/annotations/instances_val2017.json' # annotations json
pred_json = str(save_dir / f"{w}_predictions.json") # predictions json
print('\nEvaluating pycocotools mAP... saving %s...' % pred_json)
with open(pred_json, 'w') as f:
json.dump(jdict, f)
try: # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
anno = COCO(anno_json) # init annotations api
pred = anno.loadRes(pred_json) # init predictions api
eval = COCOeval(anno, pred, 'bbox')
if is_coco:
eval.params.imgIds = [int(Path(x).stem) for x in dataloader.dataset.img_files] # image IDs to evaluate
eval.evaluate()
eval.accumulate()
eval.summarize()
map, map50 = eval.stats[:2] # update results ([email protected]:0.95, [email protected])
except Exception as e:
print(f'pycocotools unable to run: {e}')
# Return results
model.float() # for training
if not training:
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
print(f"Results saved to {save_dir}{s}")
# if len(error_log)>0:
# fn = f'{save_dir}/error_log.txt'
# save_list(error_log, fn)
maps = np.zeros(nc) + map
for i, c in enumerate(ap_class):
maps[c] = ap[i]
return (mp, mr, mf1, map50, map, *(loss.cpu() / len(dataloader)).tolist()), maps, t
class Evaluator:
default_metrics = [
"False Positive Rate",
"Dice",
"Jaccard",
"Precision",
"Recall",
"Accuracy",
"False Omission Rate",
"Negative Predictive Value",
"False Negative Rate",
"True Negative Rate",
"False Discovery Rate",
"Total Positives Test",
"Total Positives Reference"
]
default_advanced_metrics = [
"Hausdorff Distance",
"Hausdorff Distance 95",
"Avg. Surface Distance",
"Avg. Symmetric Surface Distance"
]
def __init__(self, cf,
test=None,
reference=None,
labels=None,
metrics=None,
advanced_metrics=None,
nan_for_nonexisting=True):
self.cf = cf
self.test = None
self.reference = None
self.confusion_matrix = ConfusionMatrix()
self.labels = None
self.nan_for_nonexisting = nan_for_nonexisting
self.result = None
self.metrics = []
if metrics is None:
for m in self.default_metrics:
self.metrics.append(m)
else:
for m in metrics:
self.metrics.append(m)
self.advanced_metrics = []
if advanced_metrics is None:
for m in self.default_advanced_metrics:
self.advanced_metrics.append(m)
else:
for m in advanced_metrics:
self.advanced_metrics.append(m)
self.set_reference(reference)
self.set_test(test)
if labels is not None:
self.set_labels(labels)
else:
if test is not None and reference is not None:
self.gen_labels()
def set_test(self, test):
"""Set the test segmentation."""
self.test = test
def set_reference(self, reference):
"""Set the reference segmentation."""
self.reference = reference
def set_labels(self, labels):
"""Set the labels.
:param labels= may be a dictionary (int->str), a set (of ints), a tuple (of ints) or a list (of ints). Labels
will only have names if you pass a dictionary"""
if isinstance(labels, dict):
self.labels = collections.OrderedDict(labels)
elif isinstance(labels, set):
self.labels = list(labels)
elif isinstance(labels, np.ndarray):
self.labels = [i for i in labels]
elif isinstance(labels, (list, tuple)):
self.labels = labels
else:
raise TypeError("Can only handle dict, list, tuple, set & numpy array, but input is of type {}".format(type(labels)))
def gen_labels(self):
"""Construct label set from unique entries in segmentations."""
if self.test is None and self.reference is None:
raise ValueError("No test or reference segmentations.")
elif self.test is None:
labels = np.unique(self.reference)
else:
labels = np.union1d(np.unique(self.test),
np.unique(self.reference))
self.labels = list(map(lambda x: int(x), labels))
def set_metrics(self, metrics):
"""Set evaluation metrics"""
if isinstance(metrics, set):
self.metrics = list(metrics)
elif isinstance(metrics, (list, tuple, np.ndarray)):
self.metrics = metrics
else:
raise TypeError("Can only handle list, tuple, set & numpy array, but input is of type {}".format(type(metrics)))
def add_metric(self, metric):
if metric not in self.metrics:
self.metrics.append(metric)
def evaluate(self, test=None, reference=None, advanced=False, **metric_kwargs):
"""Compute metrics for segmentations."""
if test is not None:
self.set_test(test)
if reference is not None:
self.set_reference(reference)
if self.test is None or self.reference is None:
raise ValueError("Need both test and reference segmentations.")
if self.labels is None:
self.gen_labels()
self.metrics.sort()
# get functions for evaluation
_funcs = {m: ALL_METRICS[m] for m in self.metrics + self.advanced_metrics}
frames = inspect.getouterframes(inspect.currentframe())
for metric in self.metrics:
for f in frames:
if metric in f[0].f_locals:
_funcs[metric] = f[0].f_locals[metric]
break
else:
if metric in _funcs:
continue
else:
raise NotImplementedError(
"Metric {} not implemented.".format(metric))
# get results
self.result = OrderedDict()
eval_metrics = self.metrics
if advanced:
eval_metrics += self.advanced_metrics
if isinstance(self.labels, dict):
for label, name in self.labels.items():
k = str(name)
self.result[k] = OrderedDict()
if not hasattr(label, "__iter__"):
self.confusion_matrix.set_test(self.test == label)
self.confusion_matrix.set_reference(self.reference == label)
else:
current_test = 0
current_reference = 0
for l in label:
current_test += (self.test == l)
current_reference += (self.reference == l)
self.confusion_matrix.set_test(current_test)
self.confusion_matrix.set_reference(current_reference)
for metric in eval_metrics:
self.result[k][metric] = _funcs[metric](confusion_matrix=self.confusion_matrix,
nan_for_nonexisting=self.nan_for_nonexisting,
**metric_kwargs)
else:
for i, l in enumerate(self.labels):
k = str(l)
self.result[k] = OrderedDict()
self.confusion_matrix.set_test(self.test == l)
self.confusion_matrix.set_reference(self.reference == l)
for metric in eval_metrics:
self.result[k][metric] = _funcs[metric](confusion_matrix=self.confusion_matrix,
nan_for_nonexisting=self.nan_for_nonexisting,
**metric_kwargs)
return self.result
def to_dict(self):
if self.result is None:
self.evaluate()
return self.result
def to_array(self):
"""Return result as numpy array (labels x metrics)."""
if self.result is None:
self.evaluate
result_metrics = sorted(self.result[list(self.result.keys())[0]].keys())
a = np.zeros((len(self.labels), len(result_metrics)), dtype=np.float32)
if isinstance(self.labels, dict):
for i, label in enumerate(self.labels.keys()):
for j, metric in enumerate(result_metrics):
a[i][j] = self.result[self.labels[label]][metric]
else:
for i, label in enumerate(self.labels):
for j, metric in enumerate(result_metrics):
a[i][j] = self.result[label][metric]
return a
def to_pandas(self):
"""Return result as pandas DataFrame."""
a = self.to_array()
if isinstance(self.labels, dict):
labels = list(self.labels.values())
else:
labels = self.labels
result_metrics = sorted(self.result[list(self.result.keys())[0]].keys())
return pd.DataFrame(a, index=labels, columns=result_metrics)
def test(
data,
weights=None,
batch_size=32,
imgsz=640,
conf_thres=0.001,
iou_thres=0.6, # for NMS
save_json=False,
single_cls=False,
augment=False,
verbose=False,
model=None,
dataloader=None,
save_dir=Path(''), # for saving images
save_txt=False, # for auto-labelling
save_hybrid=False, # for hybrid auto-labelling
save_conf=False, # save auto-label confidences
plots=False,
log_imgs=0, # number of logged images
compute_loss=None):
# Initialize/load model and set device
logger = setup_logger('Test', './')
write_info(logger, True)
training = model is not None
if training: # called by train.py
device = next(model.parameters()).device # get model device
else: # called directly
set_logging()
device = select_device(opt.device, batch_size=batch_size)
# Directories
save_dir = Path(
increment_path(Path(opt.project) / opt.name,
exist_ok=opt.exist_ok)) # increment run
# (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
# Half
half = device.type != 'cpu' # half precision only supported on CUDA
if half:
model.half()
# Configure
model.eval()
with open(data) as f:
data = yaml.load(f, Loader=yaml.FullLoader) # model dict
check_dataset(data) # check
nc = 1 if single_cls else int(data['nc']) # number of classes
iouv = torch.linspace(0.5, 0.95,
10).to(device) # iou vector for [email protected]:0.95
niou = iouv.numel()
# Logging
log_imgs, wandb = min(log_imgs, 100), None # ceil
try:
import wandb # Weights & Biases
except ImportError:
log_imgs = 0
# Dataloader
if not training:
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
path = data['test'] if opt.task == 'test' else data[
'val'] # path to val/test images
dataloader = create_dataloader(
path,
imgsz,
batch_size,
model.stride.max(),
opt,
pad=0.5,
rect=True,
prefix=colorstr('test: ' if opt.task == 'test' else 'val: '))[0]
# write_imglist(path, logger)
seen = 0
confusion_matrix = ConfusionMatrix(nc=nc)
names = {
k: v
for k, v in enumerate(
model.names if hasattr(model, 'names') else model.module.names)
}
p, r, f1, mp, mr, map50, map, t0, t1 = 0., 0., 0., 0., 0., 0., 0., 0., 0.
loss = torch.zeros(3, device=device)
jdict, stats, ap, ap_class, wandb_images = [], [], [], [], []
for batch_i, (img, targets, paths, shapes) in enumerate(dataloader):
stats_perimg = []
img = img.to(device, non_blocking=True)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
targets = targets.to(device)
nb, _, height, width = img.shape # batch size, channels, height, width
with torch.no_grad():
# Run model
t = time_synchronized()
inf_out, train_out = model(
img, augment=augment) # inference and training outputs
t0 += time_synchronized() - t
# Compute loss
if compute_loss:
loss += compute_loss([x.float() for x in train_out],
targets)[1][:3] # box, obj, cls
# Run NMS
targets[:, 2:] *= torch.Tensor([width, height, width,
height]).to(device) # to pixels
lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)
] if save_hybrid else [] # for autolabelling
t = time_synchronized()
output = non_max_suppression(inf_out,
conf_thres=conf_thres,
iou_thres=iou_thres,
labels=lb)
t1 += time_synchronized() - t
# Statistics per image
for si, pred in enumerate(output):
labels = targets[targets[:, 0] == si, 1:]
nl = len(labels)
tcls = labels[:, 0].tolist() if nl else [] # target class
path = Path(paths[si])
seen += 1
if len(pred) == 0:
if nl:
stats.append((torch.zeros(0, niou, dtype=torch.bool),
torch.Tensor(), torch.Tensor(), tcls))
stats_perimg.append((torch.zeros(0, niou,
dtype=torch.bool),
torch.Tensor(), torch.Tensor(), tcls))
continue
# Predictions
predn = pred.clone()
scale_coords(img[si].shape[1:], predn[:, :4], shapes[si][0],
shapes[si][1]) # native-space pred
# Assign all predictions as incorrect
correct = torch.zeros(pred.shape[0],
niou,
dtype=torch.bool,
device=device)
if nl:
detected = [] # target indices
tcls_tensor = labels[:, 0]
# target boxes
tbox = xywh2xyxy(labels[:, 1:5])
scale_coords(img[si].shape[1:], tbox, shapes[si][0],
shapes[si][1]) # native-space labels
if plots:
confusion_matrix.process_batch(
pred, torch.cat((labels[:, 0:1], tbox), 1))
# Per target class
for cls in torch.unique(tcls_tensor):
ti = (cls == tcls_tensor).nonzero(as_tuple=False).view(
-1) # prediction indices
pi = (cls == pred[:, 5]).nonzero(as_tuple=False).view(
-1) # target indices
# Search for detections
if pi.shape[0]:
# Prediction to target ious
ious, i = box_iou(predn[pi, :4], tbox[ti]).max(
1) # best ious, indices
# Append detections
detected_set = set()
for j in (ious > iouv[0]).nonzero(as_tuple=False):
d = ti[i[j]] # detected target
if d.item() not in detected_set:
detected_set.add(d.item())
detected.append(d)
correct[
pi[j]] = ious[j] > iouv # iou_thres is 1xn
if len(
detected
) == nl: # all targets already located in image
break
# Append statistics (correct, conf, pcls, tcls)
stats.append(
(correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls))
stats_perimg.append(
(correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls))
# # f1 per image
# mf1_perimg = 0
# mp_perimg = 0
# mr_perimg = 0
# stats_perimg = [np.concatenate(x, 0) for x in zip(*stats_perimg)]
# if len(stats_perimg) and stats_perimg[0].any():
# p_perimg, r_perimg, _, f1_perimg, ap_class_perimg = ap_per_class(*stats_perimg, plot=plots, save_dir=save_dir, names=names)
# p_perimg, r_perimg, f1_perimg = p_perimg[:, 0], r_perimg[:, 0], f1_perimg[:, 0] # [P, R, [email protected], [email protected]:0.95]
# nt_perimg = np.bincount(stats_perimg[3].astype(np.int64), minlength=nc) # number of targets per class
# for ind, c in enumerate(ap_class_perimg):
# mf1_perimg += f1_perimg[ind] * nt_perimg[c]
# mp_perimg += p_perimg[ind] * nt_perimg[c]
# mr_perimg += r_perimg[ind] * nt_perimg[c]
# mf1_perimg = mf1_perimg/nt_perimg.sum()
# mp_perimg = mp_perimg / nt_perimg.sum()
# mr_perimg = mr_perimg / nt_perimg.sum()
#
# logger.info('[{}] {} [F1 score:{:4f} (Prec: {:4f}, Rec: {:4f})]'.format(str(batch_i + 1), paths[0].split('/')[-1], mf1_perimg, mp_perimg, mr_perimg))
# save GPS log
# Logger_System('./xmls', './Logger', output, paths, names)
# Compute statistics
stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy
if len(stats) and stats[0].any():
p, r, ap, f1, ap_class = ap_per_class(*stats,
plot=plots,
save_dir=save_dir,
names=names)
p, r, ap50, ap, f1 = p[:, 0], r[:, 0], ap[:, 0], ap.mean(
1), f1[:, 0] # [P, R, [email protected], [email protected]:0.95]
mp, mr, map50, map, mf1 = p.mean(), r.mean(), ap50.mean(), ap.mean(
), f1.mean()
nt = np.bincount(stats[3].astype(np.int64),
minlength=nc) # number of targets per class
else:
nt = torch.zeros(1)
meanf1 = 0
meanp = 0
meanr = 0
meanap = 0
# Print results per class
if (verbose or (nc <= 20)) and nc > 1 and len(stats):
for i, c in enumerate(ap_class):
meanf1 += f1[i] * nt[c]
meanp += p[i] * nt[c]
meanr += r[i] * nt[c]
meanap += ap50[i] * nt[c]
meanf1 = meanf1 / nt.sum()
meanp = meanp / nt.sum()
meanr = meanr / nt.sum()
meanap = meanap / nt.sum()
logger.info('[Final] F1 score:{:4f} (Prec: {:4f}, Rec: {:4f})\n'.format(
meanf1, meanp, meanr))
# Print speeds
t = tuple(x / seen * 1E3
for x in (t0, t1, t0 + t1)) + (imgsz, imgsz, batch_size) # tuple
if not training:
print(
'Speed: %.1f/%.1f/%.1f ms inference/NMS/total per %gx%g image at batch-size %g'
% t)
# Return results
model.float() # for training
maps = np.zeros(nc) + map
for i, c in enumerate(ap_class):
maps[c] = ap[i]
write_info(logger, False)
return (meanf1, mp, mr, map50, meanap,
*(loss.cpu() / len(dataloader)).tolist()), maps, t
class Trainer(object):
def __init__(self, criterion, optimizer, n_class, sub_batchsize, mode=1, fmreg=0.15):
self.criterion = criterion
self.optimizer = optimizer
self.metrics_global = ConfusionMatrix(n_class)
self.metrics_local = ConfusionMatrix(n_class)
self.metrics = ConfusionMatrix(n_class)
self.n_class = n_class
self.sub_batchsize = sub_batchsize
self.mode = mode
self.fmreg = fmreg # regulization item
def get_scores(self):
score_train = self.metrics.get_scores()
score_train_local = self.metrics_local.get_scores()
score_train_global = self.metrics_global.get_scores()
return score_train, score_train_global, score_train_local
def reset_metrics(self):
self.metrics.reset()
self.metrics_local.reset()
self.metrics_global.reset()
def train(self, sample, model):
model.train()
labels = sample['label'].squeeze(1).long()
labels_npy = np.array(labels)
labels_torch = labels.cuda()
h, w = sample['output_size'][0]
# print(labels[0].size)
if self.mode == 1: # global
img_g = sample['image_g'].cuda()
outputs_g = model.forward(img_g)
outputs_g = F.interpolate(outputs_g, size=(h, w), mode='bilinear')
# print(outputs_g.size(), labels_torch.size())
loss = self.criterion(outputs_g, labels_torch)
loss.backward()
self.optimizer.step()
self.optimizer.zero_grad()
if self.mode == 2: # local
img_l = sample['image_l'].cuda()
batch_size = img_l.size(0)
idx = 0
outputs_l = []
while idx+self.sub_batchsize <= batch_size:
output_l = model.forward(img_l[idx:idx+self.sub_batchsize])
output_l = F.interpolate(output_l, size=(h, w), mode='bilinear')
outputs_l.append(output_l)
loss = self.criterion(output_l, labels_torch[idx:idx+self.sub_batchsize])
loss.backward()
idx += self.sub_batchsize
outputs_l = torch.cat(outputs_l, dim=0)
self.optimizer.step()
self.optimizer.zero_grad()
if self.mode == 3: # global&local
img_g = sample['image_g'].cuda()
img_l = sample['image_l'].cuda()
batch_size = img_l.size(0)
idx = 0
outputs = []; outputs_g = []; outputs_l = []
while idx+self.sub_batchsize <= batch_size:
output, output_g, output_l, mse = model.forward(img_g[idx:idx+self.sub_batchsize], img_l[idx:idx+self.sub_batchsize], target=labels_torch[idx:idx+self.sub_batchsize])
outputs.append(output); outputs_g.append(output_g); outputs_l.append(output_l)
loss = 2* self.criterion(output, labels_torch[idx:idx+self.sub_batchsize]) + self.criterion(output_g, labels_torch[idx:idx+self.sub_batchsize]) + \
self.criterion(output_l, labels_torch[idx:idx+self.sub_batchsize]) + self.fmreg * mse
loss.backward()
idx += self.sub_batchsize
outputs = torch.cat(outputs, dim=0); outputs_g = torch.cat(outputs_g, dim=0); outputs_l = torch.cat(outputs_l, dim=0)
self.optimizer.step()
self.optimizer.zero_grad()
# predictions
if self.mode == 1:
outputs_g = outputs_g.cpu()
predictions_global = [outputs_g[i:i+1].argmax(1).detach().numpy() for i in range(len(labels))]
self.metrics_global.update(labels_npy, predictions_global)
if self.mode == 2:
outputs_l = outputs_l.cpu()
predictions_local = [outputs_l[i:i+1].argmax(1).detach().numpy() for i in range(len(labels))]
self.metrics_local.update(labels_npy, predictions_local)
if self.mode == 3:
outputs_g = outputs_g.cpu(); outputs_l = outputs_l.cpu(); outputs = outputs.cpu()
predictions_global = [outputs_g[i:i+1].argmax(1).detach().numpy() for i in range(len(labels))]
predictions_local = [outputs_l[i:i+1].argmax(1).detach().numpy() for i in range(len(labels))]
predictions = [outputs[i:i+1].argmax(1).detach().numpy() for i in range(len(labels))]
self.metrics_global.update(labels_npy, predictions_global)
self.metrics_local.update(labels_npy, predictions_local)
self.metrics.update(labels_npy, predictions)
return loss
"R10",
"M10",
"M20",
"S20",
"R20",
"S30",
"S40",
]
if __name__ == '__main__':
if not os.path.isdir(OUTDIR):
os.makedirs(OUTDIR)
iouv = torch.linspace(0.01, 0.5, 10).to(device) # iou vector for [email protected]:0.95
stats = []
confusion_matrix = ConfusionMatrix(nc=len(names))
names = {k: v for k, v in enumerate(names)}
with open(images_file) as f:
image_names = [line.strip() for line in f]
# bbox_labels = get_labels(labels_dir, labels)
# xyxy = xywh2xyxy(np.array(bbox_labels)[:, 1:])
# l = np.array(bbox_labels)[:, 0:1]
# labels = np.concatenate((l, xyxy), axis=1)
seen = 0
for img_name in tqdm(image_names):
seen += 1
labels = torch.from_numpy(get_labels(img_name)).to(device)
dets = get_detections(img_name)
if dets is None: continue
def evaluate(self, eval_reader, batch_size=1, epoch_id=None):
"""评估。
Args:
eval_reader (reader): 评估数据读取器。
batch_size (int): 评估时的batch大小。默认1。
epoch_id (int): 当前评估模型所在的训练轮数。
return_details (bool): 是否返回详细信息。默认False。
Returns:
dict: 当return_details为False时，返回dict。包含关键字：'miou'、'category_iou'、'macc'、
'category_acc'和'kappa'，分别表示平均iou、各类别iou、平均准确率、各类别准确率和kappa系数。
tuple (metrics, eval_details)：当return_details为True时，增加返回dict (eval_details)，
包含关键字：'confusion_matrix'，表示评估的混淆矩阵。
"""
self.arrange_transform(transforms=eval_reader.transforms, mode='train')
total_steps = math.ceil(eval_reader.num_samples * 1.0 / batch_size)
conf_mat = ConfusionMatrix(self.num_classes, streaming=True)
data_generator = eval_reader.generator(batch_size=batch_size,
drop_last=False)
if not hasattr(self, 'parallel_test_prog'):
self.parallel_test_prog = fluid.CompiledProgram(
self.test_prog).with_data_parallel(
share_vars_from=self.parallel_train_prog)
logging.info(
"Start to evaluating(total_samples={}, total_steps={})...".format(
eval_reader.num_samples, total_steps))
for step, data in tqdm.tqdm(enumerate(data_generator()),
total=total_steps):
images = np.array([d[0] for d in data])
images = images.astype(np.float32)
labels = np.array([d[1] for d in data])
num_samples = images.shape[0]
if num_samples < batch_size:
num_pad_samples = batch_size - num_samples
pad_images = np.tile(images[0:1], (num_pad_samples, 1, 1, 1))
images = np.concatenate([images, pad_images])
feed_data = {'image': images}
outputs = self.exe.run(self.parallel_test_prog,
feed=feed_data,
fetch_list=list(self.test_outputs.values()),
return_numpy=True)
pred = outputs[0]
if num_samples < batch_size:
pred = pred[0:num_samples]
mask = labels != self.ignore_index
conf_mat.calculate(pred=pred, label=labels, ignore=mask)
_, iou = conf_mat.mean_iou()
logging.debug("[EVAL] Epoch={}, Step={}/{}, iou={}".format(
epoch_id, step + 1, total_steps, iou))
category_iou, miou = conf_mat.mean_iou()
category_acc, macc = conf_mat.accuracy()
precision, recall = conf_mat.precision_recall()
metrics = OrderedDict(
zip([
'miou', 'category_iou', 'macc', 'category_acc', 'kappa',
'precision', 'recall'
], [
miou, category_iou, macc, category_acc,
conf_mat.kappa(), precision, recall
]))
logging.info('[EVAL] Finished, Epoch={}, {} .'.format(
epoch_id, dict2str(metrics)))
return metrics
class Trainer(object):
# def __init__(self, criterion, optimizer, n_class, size0, size_g, size_p, n, sub_batch_size=6, mode=1, lamb_fmreg=0.15):
def __init__(
self,
criterion,
optimizer,
n_class,
size_g,
size_p,
sub_batch_size=6,
mode=1,
lamb_fmreg=0.15,
):
self.criterion = criterion
self.optimizer = optimizer
self.metrics_global = ConfusionMatrix(n_class)
self.metrics_local = ConfusionMatrix(n_class)
self.metrics = ConfusionMatrix(n_class)
self.n_class = n_class
# self.size0 = size0
self.size_g = size_g
self.size_p = size_p
# self.n = n
self.sub_batch_size = sub_batch_size
self.mode = mode
self.lamb_fmreg = lamb_fmreg
# self.ratio = float(size_p[0]) / size0[0]
# self.step = (size0[0] - size_p[0]) // (n - 1)
# self.template, self.coordinates = template_patch2global(size0, size_p, n, self.step)
def set_train(self, model):
model.module.ensemble_conv.train()
if self.mode == 1 or self.mode == 3:
model.module.resnet_global.train()
model.module.fpn_global.train()
else:
model.module.resnet_local.train()
model.module.fpn_local.train()
def get_scores(self):
score_train = self.metrics.get_scores()
score_train_local = self.metrics_local.get_scores()
score_train_global = self.metrics_global.get_scores()
return score_train, score_train_global, score_train_local
def reset_metrics(self):
self.metrics.reset()
self.metrics_local.reset()
self.metrics_global.reset()
def train(self, sample, model, global_fixed):
images, labels = sample["image"], sample["label"] # PIL images
#lbls = [RGB_mapping_to_class(np.array(label)) for label in labels]
#labels = [Image.fromarray(lbl) for lbl in lbls]
# del(lbls)
labels_npy = masks_transform(
labels, numpy=True
) # label of origin size in numpy
images_glb = resize(images, self.size_g) # list of resized PIL images
images_glb = images_transform(images_glb)
labels_glb = resize(
labels, (self.size_g[0] // 4, self.size_g[1] // 4), label=True
) # down 1/4 for loss
# labels_glb = resize(labels, self.size_g, label=True) # must downsample image for reduced GPU memory
labels_glb = masks_transform(labels_glb) # 127 * 127 * 8 = 129032
if self.mode == 2 or self.mode == 3:
patches, coordinates, templates, sizes, ratios = global2patch(
# patches, coordinates, _, sizes, ratios = global2patch(
images, self.size_p
)
label_patches, _, _, _, _ = global2patch(labels, self.size_p)
# patches, label_patches = global2patch(images, self.n, self.step, self.size_p), global2patch(labels, self.n, self.step, self.size_p)
# predicted_patches = [ np.zeros((self.n**2, self.n_class, self.size_p[0], self.size_p[1])) for i in range(len(images)) ]
# predicted_ensembles = [ np.zeros((self.n**2, self.n_class, self.size_p[0], self.size_p[1])) for i in range(len(images)) ]
predicted_patches = [
np.zeros(
(len(coordinates[i]), self.n_class, self.size_p[0], self.size_p[1])
)
for i in range(len(images))
]
predicted_ensembles = [
np.zeros(
(len(coordinates[i]), self.n_class, self.size_p[0], self.size_p[1])
)
for i in range(len(images))
]
outputs_global = [None for i in range(len(images))]
if self.mode == 1:
# training with only (resized) global image
outputs_global, _ = model.forward(images_glb, None, None, None)
loss = self.criterion(outputs_global, labels_glb)
loss.backward()
self.optimizer.step()
self.optimizer.zero_grad()
if self.mode == 2:
# training with patches
for i in range(len(images)):
j = 0
# while j < self.n**2:
while j < len(coordinates[i]):
patches_var = images_transform(
patches[i][j : j + self.sub_batch_size]
) # b, c, h, w
label_patches_var = masks_transform(
resize(
label_patches[i][j : j + self.sub_batch_size],
(self.size_p[0] // 4, self.size_p[1] // 4),
label=True,
)
) # down 1/4 for loss
# label_patches_var = masks_transform(label_patches[i][j : j+self.sub_batch_size])
# output_ensembles, output_global, output_patches, fmreg_l2 = model.forward(images_glb[i:i+1], patches_var, self.coordinates[j : j+self.sub_batch_size], self.ratio, mode=self.mode, n_patch=self.n**2) # include cordinates
output_ensembles, output_global, output_patches, fmreg_l2 = model.forward(
images_glb[i : i + 1],
patches_var,
coordinates[i][j : j + self.sub_batch_size],
ratios[i],
mode=self.mode,
n_patch=len(coordinates[i]),
)
loss = (
self.criterion(output_patches, label_patches_var)
+ self.criterion(output_ensembles, label_patches_var)
+ self.lamb_fmreg * fmreg_l2
)
loss.backward()
# patch predictions
predicted_patches[i][j : j + output_patches.size()[0]] = (
F.interpolate(output_patches, size=self.size_p, mode="nearest")
.data.cpu()
.numpy()
)
predicted_ensembles[i][j : j + output_ensembles.size()[0]] = (
F.interpolate(
output_ensembles, size=self.size_p, mode="nearest"
)
.data.cpu()
.numpy()
)
j += self.sub_batch_size
outputs_global[i] = output_global
outputs_global = torch.cat(outputs_global, dim=0)
self.optimizer.step()
self.optimizer.zero_grad()
if self.mode == 3:
# train global with help from patches
# go through local patches to collect feature maps
# collect predictions from patches
for i in range(len(images)):
j = 0
# while j < self.n**2:
while j < len(coordinates[i]):
patches_var = images_transform(
patches[i][j : j + self.sub_batch_size]
) # b, c, h, w
# fm_patches, output_patches = model.module.collect_local_fm(images_glb[i:i+1], patches_var, self.ratio, self.coordinates, [j, j+self.sub_batch_size], len(images), global_model=global_fixed, template=self.template, n_patch_all=self.n**2) # include cordinates
fm_patches, output_patches = model.module.collect_local_fm(
images_glb[i : i + 1],
patches_var,
ratios[i],
coordinates[i],
[j, j + self.sub_batch_size],
len(images),
global_model=global_fixed,
template=templates[0],
n_patch_all=len(coordinates[i]),
)
predicted_patches[i][j : j + output_patches.size()[0]] = (
F.interpolate(output_patches, size=self.size_p, mode="nearest")
.data.cpu()
.numpy()
)
j += self.sub_batch_size
# train on global image
outputs_global, fm_global = model.forward(
images_glb, None, None, None, mode=self.mode
)
loss = self.criterion(outputs_global, labels_glb)
loss.backward(retain_graph=True)
# fmreg loss
# generate ensembles & calc loss
for i in range(len(images)):
j = 0
# while j < self.n**2:
while j < len(coordinates[i]):
label_patches_var = masks_transform(
resize(
label_patches[i][j : j + self.sub_batch_size],
(self.size_p[0] // 4, self.size_p[1] // 4),
label=True,
)
)
# label_patches_var = masks_transform(resize(label_patches[i][j : j+self.sub_batch_size], self.size_p, label=True))
fl = fm_patches[i][j : j + self.sub_batch_size].cuda()
# fg = model.module._crop_global(fm_global[i:i+1], self.coordinates[j:j+self.sub_batch_size], self.ratio)[0]
fg = model.module._crop_global(
fm_global[i : i + 1],
coordinates[i][j : j + self.sub_batch_size],
ratios[i],
)[0]
fg = F.interpolate(fg, size=fl.size()[2:], mode="bilinear")
output_ensembles = model.module.ensemble(fl, fg)
# output_ensembles = F.interpolate(model.module.ensemble(fl, fg), self.size_p, **model.module._up_kwargs)
loss = self.criterion(
output_ensembles, label_patches_var
) # + 0.15 * mse(fl, fg)
# if i == len(images) - 1 and j + self.sub_batch_size >= self.n**2:
if i == len(images) - 1 and j + self.sub_batch_size >= len(
coordinates[i]
):
loss.backward()
else:
loss.backward(retain_graph=True)
# ensemble predictions
predicted_ensembles[i][j : j + output_ensembles.size()[0]] = (
F.interpolate(
output_ensembles, size=self.size_p, mode="nearest"
)
.data.cpu()
.numpy()
)
j += self.sub_batch_size
self.optimizer.step()
self.optimizer.zero_grad()
# global predictions
# predictions_global = F.interpolate(outputs_global.cpu(), self.size0, mode='nearest').argmax(1).detach().numpy()
outputs_global = outputs_global.cpu()
predictions_global = [
F.interpolate(
outputs_global[i : i + 1], images[i].size[::-1], mode="nearest"
)
.argmax(1)
.detach()
.numpy()
for i in range(len(images))
]
self.metrics_global.update(labels_npy, predictions_global)
if self.mode == 2 or self.mode == 3:
# patch predictions
# scores_local = np.array(patch2global(predicted_patches, self.n_class, self.n, self.step, self.size0, self.size_p, len(images))) # merge softmax scores from patches (overlaps)
scores_local = np.array(
patch2global(
predicted_patches, self.n_class, sizes, coordinates, self.size_p
)
) # merge softmax scores from patches (overlaps)
predictions_local = scores_local.argmax(1) # b, h, w
self.metrics_local.update(labels_npy, predictions_local)
# combined/ensemble predictions
# scores = np.array(patch2global(predicted_ensembles, self.n_class, self.n, self.step, self.size0, self.size_p, len(images))) # merge softmax scores from patches (overlaps)
scores = np.array(
patch2global(
predicted_ensembles, self.n_class, sizes, coordinates, self.size_p
)
) # merge softmax scores from patches (overlaps)
predictions = scores.argmax(1) # b, h, w
self.metrics.update(labels_npy, predictions)
return loss
class Evaluator(object):
def __init__(self,
n_class,
size_g,
size_p,
sub_batch_size=6,
mode=1,
test=False):
self.metrics_global = ConfusionMatrix(n_class)
self.metrics_local = ConfusionMatrix(n_class)
self.metrics = ConfusionMatrix(n_class)
self.n_class = n_class
self.size_g = size_g
self.size_p = size_p
self.sub_batch_size = sub_batch_size
self.mode = mode
self.test = test
if test:
self.flip_range = [False, True]
self.rotate_range = [0, 1, 2, 3]
else:
self.flip_range = [False]
self.rotate_range = [0]
def get_scores(self):
score_train = self.metrics.get_scores()
score_train_local = self.metrics_local.get_scores()
score_train_global = self.metrics_global.get_scores()
return score_train, score_train_global, score_train_local
def reset_metrics(self):
self.metrics.reset()
self.metrics_local.reset()
self.metrics_global.reset()
def eval_test(self, sample, model, global_fixed):
with torch.no_grad():
images = sample['image']
if not self.test:
labels = sample['label'] # PIL images
labels_npy = masks_transform(labels, numpy=True)
images_global = resize(images, self.size_g)
outputs_global = np.zeros(
(len(images), self.n_class, self.size_g[0] // 4,
self.size_g[1] // 4))
if self.mode == 2 or self.mode == 3:
images_local = [image.copy() for image in images]
scores_local = [
np.zeros((1, self.n_class, images[i].size[1],
images[i].size[0])) for i in range(len(images))
]
scores = [
np.zeros((1, self.n_class, images[i].size[1],
images[i].size[0])) for i in range(len(images))
]
for flip in self.flip_range:
if flip:
# we already rotated images for 270'
for b in range(len(images)):
images_global[b] = transforms.functional.rotate(
images_global[b], 90) # rotate back!
images_global[b] = transforms.functional.hflip(
images_global[b])
if self.mode == 2 or self.mode == 3:
images_local[b] = transforms.functional.rotate(
images_local[b], 90) # rotate back!
images_local[b] = transforms.functional.hflip(
images_local[b])
for angle in self.rotate_range:
if angle > 0:
for b in range(len(images)):
images_global[b] = transforms.functional.rotate(
images_global[b], 90)
if self.mode == 2 or self.mode == 3:
images_local[b] = transforms.functional.rotate(
images_local[b], 90)
# prepare global images onto cuda
images_glb = images_transform(images_global) # b, c, h, w
if self.mode == 2 or self.mode == 3:
patches, coordinates, templates, sizes, ratios = global2patch(
images, self.size_p)
predicted_patches = [
np.zeros((len(coordinates[i]), self.n_class,
self.size_p[0], self.size_p[1]))
for i in range(len(images))
]
predicted_ensembles = [
np.zeros((len(coordinates[i]), self.n_class,
self.size_p[0], self.size_p[1]))
for i in range(len(images))
]
if self.mode == 1:
# eval with only resized global image ##########################
if flip:
outputs_global += np.flip(np.rot90(model.forward(
images_glb, None, None,
None)[0].data.cpu().numpy(),
k=angle,
axes=(3, 2)),
axis=3)
else:
outputs_global += np.rot90(model.forward(
images_glb, None, None,
None)[0].data.cpu().numpy(),
k=angle,
axes=(3, 2))
################################################################
if self.mode == 2:
# eval with patches ###########################################
for i in range(len(images)):
j = 0
while j < len(coordinates[i]):
patches_var = images_transform(
patches[i]
[j:j + self.sub_batch_size]) # b, c, h, w
output_ensembles, output_global, output_patches, _ = model.forward(
images_glb[i:i + 1],
patches_var,
coordinates[i][j:j + self.sub_batch_size],
ratios[i],
mode=self.mode,
n_patch=len(coordinates[i]))
# patch predictions
predicted_patches[i][j:j + output_patches.size(
)[0]] += F.interpolate(
output_patches,
size=self.size_p,
mode='nearest').data.cpu().numpy()
predicted_ensembles[i][j:j +
output_ensembles.size(
)[0]] += F.interpolate(
output_ensembles,
size=self.size_p,
mode='nearest'
).data.cpu().numpy()
j += patches_var.size()[0]
if flip:
outputs_global[i] += np.flip(np.rot90(
output_global[0].data.cpu().numpy(),
k=angle,
axes=(2, 1)),
axis=2)
scores_local[i] += np.flip(
np.rot90(np.array(
patch2global(
predicted_patches[i:i + 1],
self.n_class, sizes[i:i + 1],
coordinates[i:i + 1],
self.size_p)),
k=angle,
axes=(3, 2)),
axis=3
) # merge softmax scores from patches (overlaps)
scores[i] += np.flip(
np.rot90(np.array(
patch2global(
predicted_ensembles[i:i + 1],
self.n_class, sizes[i:i + 1],
coordinates[i:i + 1],
self.size_p)),
k=angle,
axes=(3, 2)),
axis=3
) # merge softmax scores from patches (overlaps)
else:
outputs_global[i] += np.rot90(
output_global[0].data.cpu().numpy(),
k=angle,
axes=(2, 1))
scores_local[i] += np.rot90(
np.array(
patch2global(
predicted_patches[i:i + 1],
self.n_class, sizes[i:i + 1],
coordinates[i:i + 1],
self.size_p)),
k=angle,
axes=(3, 2)
) # merge softmax scores from patches (overlaps)
scores[i] += np.rot90(
np.array(
patch2global(
predicted_ensembles[i:i + 1],
self.n_class, sizes[i:i + 1],
coordinates[i:i + 1],
self.size_p)),
k=angle,
axes=(3, 2)
) # merge softmax scores from patches (overlaps)
###############################################################
if self.mode == 3:
# eval global with help from patches ##################################################
# go through local patches to collect feature maps
# collect predictions from patches
for i in range(len(images)):
j = 0
while j < len(coordinates[i]):
patches_var = images_transform(
patches[i]
[j:j + self.sub_batch_size]) # b, c, h, w
fm_patches, output_patches = model.module.collect_local_fm(
images_glb[i:i + 1],
patches_var,
ratios[i],
coordinates[i],
[j, j + self.sub_batch_size],
len(images),
global_model=global_fixed,
template=templates[i],
n_patch_all=len(coordinates[i]))
predicted_patches[i][j:j + output_patches.size(
)[0]] += F.interpolate(
output_patches,
size=self.size_p,
mode='nearest').data.cpu().numpy()
j += self.sub_batch_size
# go through global image
tmp, fm_global = model.forward(images_glb,
None,
None,
None,
mode=self.mode)
if flip:
outputs_global += np.flip(np.rot90(
tmp.data.cpu().numpy(), k=angle, axes=(3, 2)),
axis=3)
else:
outputs_global += np.rot90(tmp.data.cpu().numpy(),
k=angle,
axes=(3, 2))
# generate ensembles
for i in range(len(images)):
j = 0
while j < len(coordinates[i]):
fl = fm_patches[i][j:j +
self.sub_batch_size].cuda()
fg = model.module._crop_global(
fm_global[i:i + 1],
coordinates[i][j:j + self.sub_batch_size],
ratios[i])[0]
fg = F.interpolate(fg,
size=fl.size()[2:],
mode='bilinear')
output_ensembles = model.module.ensemble(
fl, fg) # include cordinates
# ensemble predictions
predicted_ensembles[i][j:j +
output_ensembles.size(
)[0]] += F.interpolate(
output_ensembles,
size=self.size_p,
mode='nearest'
).data.cpu().numpy()
j += self.sub_batch_size
if flip:
scores_local[i] += np.flip(
np.rot90(np.array(
patch2global(
predicted_patches[i:i + 1],
self.n_class, sizes[i:i + 1],
coordinates[i:i + 1],
self.size_p)),
k=angle,
axes=(3, 2)),
axis=3
)[0] # merge softmax scores from patches (overlaps)
scores[i] += np.flip(
np.rot90(np.array(
patch2global(
predicted_ensembles[i:i + 1],
self.n_class, sizes[i:i + 1],
coordinates[i:i + 1],
self.size_p)),
k=angle,
axes=(3, 2)),
axis=3
)[0] # merge softmax scores from patches (overlaps)
else:
scores_local[i] += np.rot90(
np.array(
patch2global(
predicted_patches[i:i + 1],
self.n_class, sizes[i:i + 1],
coordinates[i:i + 1],
self.size_p)),
k=angle,
axes=(3, 2)
) # merge softmax scores from patches (overlaps)
scores[i] += np.rot90(
np.array(
patch2global(
predicted_ensembles[i:i + 1],
self.n_class, sizes[i:i + 1],
coordinates[i:i + 1],
self.size_p)),
k=angle,
axes=(3, 2)
) # merge softmax scores from patches (overlaps)
###################################################
# global predictions ###########################
outputs_global = torch.Tensor(outputs_global)
predictions_global = [
F.interpolate(outputs_global[i:i + 1],
images[i].size[::-1],
mode='nearest').argmax(1).detach().numpy()[0]
for i in range(len(images))
]
if not self.test:
self.metrics_global.update(labels_npy, predictions_global)
if self.mode == 2 or self.mode == 3:
# patch predictions ###########################
predictions_local = [
score.argmax(1)[0] for score in scores_local
]
if not self.test:
self.metrics_local.update(labels_npy, predictions_local)
###################################################
# combined/ensemble predictions ###########################
predictions = [score.argmax(1)[0] for score in scores]
if not self.test:
self.metrics.update(labels_npy, predictions)
return predictions, predictions_global, predictions_local
else:
return None, predictions_global, None
def test(
data,
weights=None,
batch_size=32,
imgsz=640,
conf_thres=0.001,
iou_thres=0.6, # for NMS
save_json=False,
single_cls=False,
augment=False,
verbose=False,
model=None,
dataloader=None,
save_dir=Path(''), # for saving images
save_txt=False, # for auto-labelling
save_hybrid=False, # for hybrid auto-labelling
save_conf=False, # save auto-label confidences
plots=True,
log_imgs=0): # number of logged images
# Initialize/load model and set device
# 判断是否在训练时调用test，如果是则获取训练时的设备
training = model is not None
if training: # called by train.py
device = next(model.parameters()).device # get model device
else: # called directly
set_logging()
device = select_device(opt.device, batch_size=batch_size)
# Directories
save_dir = Path(
increment_path(Path(opt.project) / opt.name,
exist_ok=opt.exist_ok)) # increment run
(save_dir / 'labels' if save_txt else save_dir).mkdir(
parents=True, exist_ok=True) # make dir
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
# 检查输入图片分辨率是否能被模型的最大步长（默认32）整除
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
# Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99
# if device.type != 'cpu' and torch.cuda.device_count() > 1:
# model = nn.DataParallel(model)
# Half
# 如果设备不是cpu并且gpu数目为1，则将模型由Float32转为Float16，提高前向传播的速度
half = device.type != 'cpu' # half precision only supported on CUDA
if half:
model.half() # to FP16
# Configure
# 加载数据配置信息
model.eval()
is_coco = data.endswith('coco.yaml') # is COCO dataset
with open(data) as f:
data = yaml.load(f, Loader=yaml.FullLoader) # model dict
check_dataset(data) # check
nc = 1 if single_cls else int(data['nc']) # number of classes
# 设置iou阈值，从0.5~0.95，每间隔0.05取一次
iouv = torch.linspace(0.5, 0.95,
10).to(device) # iou vector for [email protected]:0.95
# iou个数
niou = iouv.numel()
# Logging
log_imgs, wandb = min(log_imgs, 100), None # ceil
try:
import wandb # Weights & Biases
except ImportError:
log_imgs = 0
# Dataloader
if not training:
# 创建一个全0数组测试一下前向传播是否正常运行
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
# 获取图片路径
path = data['test'] if opt.task == 'test' else data[
'val'] # path to val/test images
# 创建dataloader
# 注意这里rect参数为True，yolov5的测试评估是基于矩形推理的
dataloader = create_dataloader(
path,
imgsz,
batch_size,
model.stride.max(),
opt,
pad=0.5,
rect=True,
prefix=colorstr('test: ' if opt.task == 'test' else 'val: '))[0]
seen = 0
confusion_matrix = ConfusionMatrix(nc=nc)
# 获取类别的名字
names = {
k: v
for k, v in enumerate(
model.names if hasattr(model, 'names') else model.module.names)
}
"""
获取coco数据集的类别索引
这里要说明一下，coco数据集有80个类别(索引范围应该为0~79)，
但是他的索引却属于0~90(笔者是通过查看coco数据测试集的json文件发现的，具体原因不知)
coco80_to_coco91_class()就是为了与上述索引对应起来，返回一个范围在0~90的索引数组
"""
coco91class = coco80_to_coco91_class()
# 设置tqdm进度条的显示信息
s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R',
'[email protected]', '[email protected]:.95')
# 初始化指标，时间
p, r, f1, mp, mr, map50, map, t0, t1 = 0., 0., 0., 0., 0., 0., 0., 0., 0.
# 初始化测试集的损失
loss = torch.zeros(3, device=device)
# 初始化json文件的字典，统计信息，ap
jdict, stats, ap, ap_class, wandb_images = [], [], [], [], []
for batch_i, (img, targets, paths,
shapes) in enumerate(tqdm(dataloader, desc=s)):
img = img.to(device, non_blocking=True)
# 图片也由Float32->Float16
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
targets = targets.to(device)
nb, _, height, width = img.shape # batch size, channels, height, width
with torch.no_grad():
# Run model
"""
time_synchronized()函数里面进行了torch.cuda.synchronize()，再返回的time.time()
torch.cuda.synchronize()等待gpu上完成所有的工作
总的来说就是这样测试时间会更准确
"""
t = time_synchronized()
# inf_out为预测结果, train_out训练结果
inf_out, train_out = model(
img, augment=augment) # inference and training outputs
t0 += time_synchronized() - t
# Compute loss
# 如果是在训练时进行的test，则通过训练结果计算并返回测试集的GIoU, obj, cls损失
if training:
loss += compute_loss([x.float() for x in train_out], targets,
model)[1][:3] # box, obj, cls
# Run NMS
targets[:, 2:] *= torch.Tensor([width, height, width,
height]).to(device) # to pixels
lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)
] if save_hybrid else [] # for autolabelling
t = time_synchronized()
"""
non_max_suppression进行非极大值抑制;
conf_thres为置信度阈值，iou_thres为iou阈值
merge为是否合并框
"""
output = non_max_suppression(inf_out,
conf_thres=conf_thres,
iou_thres=iou_thres,
labels=lb)
t1 += time_synchronized() - t
# Statistics per image
# 为每一张图片做统计, 写入预测信息到txt文件, 生成json文件字典, 统计tp等
for si, pred in enumerate(output):
# 获取第si张图片的标签信息, 包括class,x,y,w,h
# targets[:, 0]为标签属于哪一张图片的编号
labels = targets[targets[:, 0] == si, 1:]
nl = len(labels)
# 获取标签类别
tcls = labels[:, 0].tolist() if nl else [] # target class
path = Path(paths[si])
# 统计测试图片数量
seen += 1
# 如果预测为空，则添加空的信息到stats里
if len(pred) == 0:
if nl:
stats.append((torch.zeros(0, niou, dtype=torch.bool),
torch.Tensor(), torch.Tensor(), tcls))
continue
# Predictions
predn = pred.clone()
scale_coords(img[si].shape[1:], predn[:, :4], shapes[si][0],
shapes[si][1]) # native-space pred
# Append to text file
# 保存预测结果为txt文件
if save_txt:
# 获得对应图片的长和宽
gn = torch.tensor(shapes[si][0])[[1, 0, 1, 0
]] # normalization gain whwh
for *xyxy, conf, cls in predn.tolist():
# xyxy格式->xywh, 并对坐标进行归一化处理
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
line = (cls, *xywh,
conf) if save_conf else (cls,
*xywh) # label format
with open(save_dir / 'labels' / (path.stem + '.txt'),
'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n')
# W&B logging
if plots and len(wandb_images) < log_imgs:
box_data = [{
"position": {
"minX": xyxy[0],
"minY": xyxy[1],
"maxX": xyxy[2],
"maxY": xyxy[3]
},
"class_id": int(cls),
"box_caption": "%s %.3f" % (names[cls], conf),
"scores": {
"class_score": conf
},
"domain": "pixel"
} for *xyxy, conf, cls in pred.tolist()]
boxes = {
"predictions": {
"box_data": box_data,
"class_labels": names
}
} # inference-space
wandb_images.append(
wandb.Image(img[si], boxes=boxes, caption=path.name))
# Append to pycocotools JSON dictionary
if save_json:
# [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ...
image_id = int(
path.stem) if path.stem.isnumeric() else path.stem
box = xyxy2xywh(predn[:, :4]) # xywh
box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner
for p, b in zip(pred.tolist(), box.tolist()):
jdict.append({
'image_id':
image_id,
'category_id':
coco91class[int(p[5])] if is_coco else int(p[5]),
'bbox': [round(x, 3) for x in b],
'score':
round(p[4], 5)
})
# Assign all predictions as incorrect
# 初始化预测评定，niou为iou阈值的个数
correct = torch.zeros(pred.shape[0],
niou,
dtype=torch.bool,
device=device)
if nl:
# detected用来存放已检测到的目标
detected = [] # target indices
tcls_tensor = labels[:, 0]
# target boxes
# 获得xyxy格式的框并乘以wh
tbox = xywh2xyxy(labels[:, 1:5])
# 将预测框的坐标调整到基于其原本长宽的坐标
scale_coords(img[si].shape[1:], tbox, shapes[si][0],
shapes[si][1]) # native-space labels
if plots:
confusion_matrix.process_batch(
pred, torch.cat((labels[:, 0:1], tbox), 1))
# Per target class
# 对图片中的每个类单独处理
for cls in torch.unique(tcls_tensor):
# 标签框该类别的索引
ti = (cls == tcls_tensor).nonzero(as_tuple=False).view(
-1) # prediction indices
# 预测框该类别的索引
pi = (cls == pred[:, 5]).nonzero(as_tuple=False).view(
-1) # target indices
# Search for detections
if pi.shape[0]:
"""
Prediction to target ious
box_iou计算预测框与标签框的iou值，max(1)选出最大的ious值,i为对应的索引
pred shape[N, 4]
tbox shape[M, 4]
box_iou shape[N, M]
ious shape[N, 1]
i shape[N, 1], i里的值属于0~M
"""
ious, i = box_iou(predn[pi, :4], tbox[ti]).max(
1) # best ious, indices
# Append detections
detected_set = set()
for j in (ious > iouv[0]).nonzero(as_tuple=False):
# 获得检测到的目标
d = ti[i[j]] # detected target
if d.item() not in detected_set:
detected_set.add(d.item())
detected.append(d)
# iouv为以0.05为步长 0.5到0.95的序列
# 获得不同iou阈值下的true positive
correct[
pi[j]] = ious[j] > iouv # iou_thres is 1xn
if len(
detected
) == nl: # all targets already located in image
break
# Append statistics (correct, conf, pcls, tcls)
# 每张图片的结果统计到stats里
stats.append(
(correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls))
# Plot images
# 画出第1个batch的图片的ground truth和预测框并保存
if plots and batch_i < 3:
f = save_dir / f'test_batch{batch_i}_labels.jpg' # labels
Thread(target=plot_images,
args=(img, targets, paths, f, names),
daemon=True).start()
f = save_dir / f'test_batch{batch_i}_pred.jpg' # predictions
Thread(target=plot_images,
args=(img, output_to_target(output), paths, f, names),
daemon=True).start()
# Compute statistics
# 将stats列表的信息拼接到一起
stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy
if len(stats) and stats[0].any():
# 根据上面得到的tp等信息计算指标
# 精准度TP/TP+FP，召回率TP/P，map，f1分数，类别
p, r, ap, f1, ap_class = ap_per_class(*stats,
plot=plots,
save_dir=save_dir,
names=names)
p, r, ap50, ap = p[:, 0], r[:, 0], ap[:, 0], ap.mean(
1) # [P, R, [email protected], [email protected]:0.95]
mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
# nt是一个列表，测试集每个类别有多少个标签框
nt = np.bincount(stats[3].astype(np.int64),
minlength=nc) # number of targets per class
else:
nt = torch.zeros(1)
# Print results
pf = '%20s' + '%12.3g' * 6 # print format
print(pf % ('all', seen, nt.sum(), mp, mr, map50, map))
# Print results per class
if (verbose or (nc <= 20 and not training)) and nc > 1 and len(stats):
for i, c in enumerate(ap_class):
print(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i]))
# Print speeds
t = tuple(x / seen * 1E3
for x in (t0, t1, t0 + t1)) + (imgsz, imgsz, batch_size) # tuple
if not training:
print(
'Speed: %.1f/%.1f/%.1f ms inference/NMS/total per %gx%g image at batch-size %g'
% t)
# Plots
if plots:
confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))
if wandb and wandb.run:
wandb.log({"Images": wandb_images})
wandb.log({
"Validation": [
wandb.Image(str(f), caption=f.name)
for f in sorted(save_dir.glob('test*.jpg'))
]
})
# Save JSON
# 采用之前保存的json格式预测结果，通过cocoapi评估指标
# 需要注意的是 测试集的标签也需要转成coco的json格式
if save_json and len(jdict):
w = Path(weights[0] if isinstance(weights, list) else weights
).stem if weights is not None else '' # weights
anno_json = '../coco/annotations/instances_val2017.json' # annotations json
pred_json = str(save_dir / f"{w}_predictions.json") # predictions json
print('\nEvaluating pycocotools mAP... saving %s...' % pred_json)
with open(pred_json, 'w') as f:
json.dump(jdict, f)
try: # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
anno = COCO(anno_json) # init annotations api
pred = anno.loadRes(pred_json) # init predictions api
eval = COCOeval(anno, pred, 'bbox')
if is_coco:
eval.params.imgIds = [
int(Path(x).stem) for x in dataloader.dataset.img_files
] # image IDs to evaluate
eval.evaluate()
eval.accumulate()
eval.summarize()
map, map50 = eval.stats[:
2] # update results ([email protected]:0.95, [email protected])
except Exception as e:
print(f'pycocotools unable to run: {e}')
# Return results
if not training:
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
print(f"Results saved to {save_dir}{s}")
model.float() # for training
maps = np.zeros(nc) + map
for i, c in enumerate(ap_class):
maps[c] = ap[i]
return (mp, mr, map50, map,
*(loss.cpu() / len(dataloader)).tolist()), maps, t
class Trainer(object):
def __init__(self,
criterion,
optimizer,
n_class,
size_g,
size_p,
sub_batch_size=6,
mode=1,
lamb_fmreg=0.15):
self.criterion = criterion
self.optimizer = optimizer
self.metrics_global = ConfusionMatrix(n_class)
self.metrics_local = ConfusionMatrix(n_class)
self.metrics = ConfusionMatrix(n_class)
self.n_class = n_class
self.size_g = size_g
self.size_p = size_p
self.sub_batch_size = sub_batch_size
self.mode = mode
self.lamb_fmreg = lamb_fmreg
def set_train(self, model):
model.module.ensemble_conv.train()
if self.mode == 1 or self.mode == 3:
model.module.resnet_global.train()
model.module.fpn_global.train()
else:
model.module.resnet_local.train()
model.module.fpn_local.train()
def get_scores(self):
score_train = self.metrics.get_scores()
score_train_local = self.metrics_local.get_scores()
score_train_global = self.metrics_global.get_scores()
return score_train, score_train_global, score_train_local
def reset_metrics(self):
self.metrics.reset()
self.metrics_local.reset()
self.metrics_global.reset()
def train(self, sample, model, global_fixed):
images, labels = sample['image'], sample['label'] # PIL images
ids = sample['id']
width, height = images[0].size
if width != self.size_g[0] or height != self.size_g[1]:
images_glb = resize(images,
self.size_g) # list of resized PIL images
else:
images_glb = list(images)
images_glb = images_transform(images_glb)
labels_glb = masks_transform(labels)
if self.mode == 2 or self.mode == 3:
patches, coordinates, sizes, ratios, label_patches = global2patch(
images, labels_glb, self.size_p, ids)
predicted_patches = np.zeros((len(images), 4))
predicted_ensembles = np.zeros((len(images), 4))
outputs_global = [None for i in range(len(images))]
if self.mode == 1:
# training with only (resized) global image #########################################
outputs_global, _ = model.forward(images_glb, None, None, None)
loss = self.criterion(outputs_global, labels_glb)
loss.backward()
self.optimizer.step()
self.optimizer.zero_grad()
##############################################
if self.mode == 2:
# training with patches ###########################################
for i in range(len(images)):
j = 0
while j < len(coordinates[i]):
patches_var = images_transform(
patches[i][j:j + self.sub_batch_size]) # b, c, h, w
label_patches_var = masks_transform(
label_patches[i][j:j + self.sub_batch_size])
output_ensembles, output_global, output_patches, fmreg_l2 = model.forward(
images_glb[i:i + 1],
patches_var,
coordinates[i][j:j + self.sub_batch_size],
ratios[i],
mode=self.mode,
n_patch=len(coordinates[i]))
loss = self.criterion(
output_patches, label_patches_var) + self.criterion(
output_ensembles,
label_patches_var) + self.lamb_fmreg * fmreg_l2
loss.backward()
# patch predictions
for pred_patch, pred_ensemble in zip(
torch.max(output_patches.data, 1)[1].data,
torch.max(output_ensembles.data, 1)[1].data):
predicted_patches[i][int(pred_patch)] += 1
predicted_ensembles[i][int(pred_ensemble)] += 1
j += self.sub_batch_size
outputs_global[i] = output_global
outputs_global = torch.cat(outputs_global, dim=0)
self.optimizer.step()
self.optimizer.zero_grad()
#####################################################################################
if self.mode == 3:
# train global with help from patches ##################################################
# go through local patches to collect feature maps
# collect predictions from patches
for i in range(len(images)):
j = 0
while j < len(coordinates[i]):
patches_var = images_transform(
patches[i][j:j + self.sub_batch_size]) # b, c, h, w
_, output_patches = model.module.collect_local_fm(
images_glb[i:i + 1],
patches_var,
ratios[i],
coordinates[i], [j, j + self.sub_batch_size],
len(images),
global_model=global_fixed,
n_patch_all=len(coordinates[i]))
for pred_patch in torch.max(output_patches.data,
1)[1].data:
predicted_patches[i][int(pred_patch)] += 1
j += self.sub_batch_size
# train on global image
outputs_global, fm_global = model.forward(images_glb,
None,
None,
None,
mode=self.mode)
loss = self.criterion(outputs_global, labels_glb)
loss.backward(retain_graph=True)
# fmreg loss
# generate ensembles & calc loss
for i in range(len(images)):
j = 0
while j < len(coordinates[i]):
label_patches_var = masks_transform(
label_patches[i][j:j + self.sub_batch_size])
patches_var = images_transform(
patches[i][j:j + self.sub_batch_size]) # b, c, h, w
fl = model.module.generate_local_fm(
images_glb[i:i + 1],
patches_var,
ratios[i],
coordinates[i], [j, j + self.sub_batch_size],
len(images),
global_model=global_fixed,
n_patch_all=len(coordinates[i]))
fg = model.module._crop_global(
fm_global[i:i + 1],
coordinates[i][j:j + self.sub_batch_size],
ratios[i])[0]
fg = F.interpolate(fg, size=fl.size()[2:], mode='bilinear')
output_ensembles = model.module.ensemble(fl, fg)
loss = self.criterion(
output_ensembles,
label_patches_var) # + 0.15 * mse(fl, fg)
if i == len(images) - 1 and j + self.sub_batch_size >= len(
coordinates[i]):
loss.backward()
else:
loss.backward(retain_graph=True)
# ensemble predictions
for pred_ensemble in torch.max(output_ensembles.data,
1)[1].data:
predicted_ensembles[i][int(pred_ensemble)] += 1
j += self.sub_batch_size
self.optimizer.step()
self.optimizer.zero_grad()
# global predictions ###########################
_, predictions_global = torch.max(outputs_global.data, 1)
self.metrics_global.update(labels_glb, predictions_global)
if self.mode == 2 or self.mode == 3:
# patch predictions ###########################
predictions_local = predicted_patches.argmax(1)
#self.metrics_local.update(labels_npy, predictions_local)
self.metrics_local.update(labels_glb, predictions_local)
###################################################
# combined/ensemble predictions ###########################
predictions = predicted_ensembles.argmax(1)
self.metrics.update(labels_glb, predictions)
return loss
class Trainer(object):
def __init__(self, criterion, optimizer, n_class, size_g, size_p, sub_batch_size=6, mode=1, lamb_fmreg=0.15):
self.criterion = criterion
self.optimizer = optimizer
self.metrics_global = ConfusionMatrix(n_class)
self.metrics_local = ConfusionMatrix(n_class)
self.metrics = ConfusionMatrix(n_class)
self.n_class = n_class
self.size_g = size_g
self.size_p = size_p
self.sub_batch_size = sub_batch_size
self.mode = mode
self.lamb_fmreg = lamb_fmreg
def set_train(self, model):
model.module.ensemble_conv.train()
if self.mode == 1 or self.mode == 3:
model.module.resnet_global.train()
model.module.fpn_global.train()
else:
model.module.resnet_local.train()
model.module.fpn_local.train()
def get_scores(self):
score_train = self.metrics.get_scores()
score_train_local = self.metrics_local.get_scores()
score_train_global = self.metrics_global.get_scores()
return score_train, score_train_global, score_train_local
def reset_metrics(self):
self.metrics.reset()
self.metrics_local.reset()
self.metrics_global.reset()
def train(self, sample, model, global_fixed):
images, labels, labels_npy, images_glb = sample['image'], sample['label'], sample['label_npy'], sample['image_glb'] # PIL images
#labels_npy = masks_transform(labels, numpy=True) # label of origin size in numpy
#images_glb = resize(images, self.size_g) # list of resized PIL images
#images_glb = images_transform(images_glb)
labels_glb = resize(labels, (self.size_g[0] // 4, self.size_g[1] // 4), label=True) # FPN down 1/4, for loss
labels_glb = masks_transform(labels_glb)
if self.mode == 2 or self.mode == 3:
patches, coordinates, templates, sizes, ratios = global2patch(images, self.size_p)
label_patches, _, _, _, _ = global2patch(labels, self.size_p)
#predicted_patches = [ np.zeros((len(coordinates[i]), self.n_class, self.size_p[0], self.size_p[1])) for i in range(len(images)) ]
#predicted_ensembles = [ np.zeros((len(coordinates[i]), self.n_class, self.size_p[0], self.size_p[1])) for i in range(len(images)) ]
#outputs_global = [ None for i in range(len(images)) ]
if self.mode == 1:
# training with only (resized) global image #########################################
outputs_global, _ = model.forward(images_glb, None, None, None)
loss = self.criterion(outputs_global, labels_glb)
loss.backward()
self.optimizer.step()
self.optimizer.zero_grad()
##############################################
if self.mode == 2:
# training with patches ###########################################
subdataset = AerialSubdatasetMode2(images_glb, ratios, coordinates, patches, label_patches, (self.size_p[0] // 4, self.size_p[1] // 4))
data_loader = torch.utils.data.DataLoader(dataset=subdataset, \
batch_size=self.sub_batch_size, \
num_workers=20, \
collate_fn=collate_mode2, \
shuffle=False, pin_memory=True)
for batch_sample in data_loader:
for sub_batch_id in range(len(batch_sample['n_patch'])):
patches_var = batch_sample['patches'][sub_batch_id].cuda()
label_patches_var = batch_sample['labels'][sub_batch_id].cuda()
output_ensembles, output_global, output_patches, fmreg_l2 = model.forward(batch_sample['images_glob'][sub_batch_id].cuda(), \
patches_var, \
batch_sample['coords'][sub_batch_id], \
batch_sample['ratio'][sub_batch_id], mode=self.mode, \
n_patch=batch_sample['n_patch'][sub_batch_id])
loss = self.criterion(output_patches, label_patches_var) + self.criterion(output_ensembles, label_patches_var) + self.lamb_fmreg * fmreg_l2
loss.backward()
'''
for i in range(len(images)):
j = 0
print("LEN", len(coordinates[i]))
while j < len(coordinates[i]):
track.start("transform_internal")
patches_var = images_transform(patches[i][j : j+self.sub_batch_size]) # b, c, h, w
label_patches_var = masks_transform(resize(label_patches[i][j : j+self.sub_batch_size], (self.size_p[0] // 4, self.size_p[1] // 4), label=True)) # down 1/4 for loss
track.end("transform_internal")
track.start("ff_internal")
output_ensembles, output_global, output_patches, fmreg_l2 = model.forward(images_glb[i:i+1], patches_var, coordinates[i][j : j+self.sub_batch_size], ratios[i], mode=self.mode, n_patch=len(coordinates[i]))
track.end("ff_internal")
loss = self.criterion(output_patches, label_patches_var) + self.criterion(output_ensembles, label_patches_var) + self.lamb_fmreg * fmreg_l2
loss.backward()
# patch predictions
#predicted_patches[i][j:j+output_patches.size()[0]] = F.interpolate(output_patches, size=self.size_p, mode='nearest').data.cpu().numpy()
#predicted_ensembles[i][j:j+output_ensembles.size()[0]] = F.interpolate(output_ensembles, size=self.size_p, mode='nearest').data.cpu().numpy()
j += self.sub_batch_size
#outputs_global[i] = output_global
#outputs_global = torch.cat(outputs_global, dim=0)
'''
self.optimizer.step()
self.optimizer.zero_grad()
#####################################################################################
if self.mode == 3:
# train global with help from patches ##################################################
# go through local patches to collect feature maps
# collect predictions from patches
track.start("Collect patches")
# import pdb; pdb.set_trace();
subdataset = AerialSubdatasetMode3a(patches, coordinates, images_glb, ratios, templates)
data_loader = torch.utils.data.DataLoader(dataset=subdataset, \
batch_size=self.sub_batch_size, \
num_workers=20, \
collate_fn=collate_mode3a, \
shuffle=False, pin_memory=True)
for batch_sample in data_loader:
for sub_batch_id in range(len(batch_sample['ratios'])):
patches_var = batch_sample['patches'][sub_batch_id].cuda()
coord = batch_sample['coords'][sub_batch_id]
j = batch_sample['coord_ids'][sub_batch_id]
fm_patches, _ = model.module.collect_local_fm(batch_sample['images_glb'][sub_batch_id].cuda(), \
patches_var, \
batch_sample['ratios'][sub_batch_id], \
coord, \
[min(j), max(j) + 1], \
len(images), \
global_model=global_fixed, \
template= batch_sample['templates'][sub_batch_id].cuda(), \
n_patch_all=len(coord))
# for i in range(len(images)):
# j = 0
# while j < len(coordinates[i]):
# patches_var = images_transform(patches[i][j : j+self.sub_batch_size]) # b, c, h, w
# fm_patches, _ = model.module.collect_local_fm(images_glb[i:i+1], patches_var, ratios[i], coordinates[i], [j, j+self.sub_batch_size], len(images), global_model=global_fixed, template=templates[i], n_patch_all=len(coordinates[i]))
# j += self.sub_batch_size
track.end("Collect patches")
images_glb = images_glb.cuda()
# train on global image
outputs_global, fm_global = model.forward(images_glb, None, None, None, mode=self.mode)
loss = self.criterion(outputs_global, labels_glb)
loss.backward(retain_graph=True)
subdataset = AerialSubdatasetMode3b(label_patches, \
(self.size_p[0] // 4, self.size_p[1] // 4), \
fm_patches,\
coordinates, ratios)
data_loader = torch.utils.data.DataLoader(dataset=subdataset, \
batch_size=self.sub_batch_size, \
num_workers=20, \
collate_fn=collate_mode3b, \
shuffle=False, pin_memory=True)
track.start("load_mode_3b")
for batch_idx, batch_sample in enumerate(data_loader):
for sub_batch_id in range(len(batch_sample['ratios'])):
label_patches_var = batch_sample['label_patches'][sub_batch_id].cuda()
fl = batch_sample['fl'][sub_batch_id].cuda()
image_id = batch_sample['id'][sub_batch_id]
track.end("load_mode_3b")
fg = model.module._crop_global(fm_global[image_id: image_id+1], \
batch_sample['coords'][sub_batch_id], \
batch_sample['ratios'][sub_batch_id])[0]
fg = F.interpolate(fg, size=fl.size()[2:], mode='bilinear')
output_ensembles = model.module.ensemble(fl, fg)
loss = self.criterion(output_ensembles, label_patches_var)# + 0.15 * mse(fl, fg)
if batch_idx == len(data_loader) - 1 and sub_batch_id == len(batch_sample['ratios']) - 1:
loss.backward()
else:
loss.backward(retain_graph=True)
track.start("load_mode_3b")
# fmreg loss
# generate ensembles & calc loss
"""
track.start("load_mode_3b")
for i in range(len(images)):
j = 0
while j < len(coordinates[i]):
label_patches_var = masks_transform(resize(label_patches[i][j : j+self.sub_batch_size], (self.size_p[0] // 4, self.size_p[1] // 4), label=True))
fl = fm_patches[i][j : j+self.sub_batch_size].cuda()
track.end("load_mode_3b")
fg = model.module._crop_global(fm_global[i:i+1], coordinates[i][j:j+self.sub_batch_size], ratios[i])[0]
fg = F.interpolate(fg, size=fl.size()[2:], mode='bilinear')
output_ensembles = model.module.ensemble(fl, fg)
loss = self.criterion(output_ensembles, label_patches_var)# + 0.15 * mse(fl, fg)
if i == len(images) - 1 and j + self.sub_batch_size >= len(coordinates[i]):
loss.backward()
else:
loss.backward(retain_graph=True)
track.start("load_mode_3b")
# ensemble predictions
#predicted_ensembles[i][j:j+output_ensembles.size()[0]] = F.interpolate(output_ensembles, size=self.size_p, mode='nearest').data.cpu().numpy()
j += self.sub_batch_size
"""
self.optimizer.step()
self.optimizer.zero_grad()
'''
# global predictions ###########################
outputs_global = outputs_global.cpu()
predictions_global = [F.interpolate(outputs_global[i:i+1], images[i].size[::-1], mode='nearest').argmax(1).detach().numpy() for i in range(len(images))]
self.metrics_global.update(labels_npy, predictions_global)
if self.mode == 2 or self.mode == 3:
# patch predictions ###########################
scores_local = np.array(patch2global(predicted_patches, self.n_class, sizes, coordinates, self.size_p)) # merge softmax scores from patches (overlaps)
predictions_local = scores_local.argmax(1) # b, h, w
self.metrics_local.update(labels_npy, predictions_local)
###################################################
# combined/ensemble predictions ###########################
scores = np.array(patch2global(predicted_ensembles, self.n_class, sizes, coordinates, self.size_p)) # merge softmax scores from patches (overlaps)
predictions = scores.argmax(1) # b, h, w
self.metrics.update(labels_npy, predictions)
'''
return loss
class Evaluator(object):
def __init__(self,
n_class,
size_g,
size_p,
sub_batch_size=6,
mode=1,
test=False):
self.metrics_global = ConfusionMatrix(n_class)
self.metrics_local = ConfusionMatrix(n_class)
self.metrics = ConfusionMatrix(n_class)
self.n_class = n_class
self.size_g = size_g
self.size_p = size_p
self.sub_batch_size = sub_batch_size
self.mode = mode
self.test = test
if test:
self.flip_range = [False, True]
self.rotate_range = [0, 1, 2, 3]
else:
self.flip_range = [False]
self.rotate_range = [0]
def get_scores(self):
score_train = self.metrics.get_scores()
score_train_local = self.metrics_local.get_scores()
score_train_global = self.metrics_global.get_scores()
return score_train, score_train_global, score_train_local
def reset_metrics(self):
self.metrics.reset()
self.metrics_local.reset()
self.metrics_global.reset()
def eval_test(self, sample, model, global_fixed):
with torch.no_grad():
images = sample['image']
ids = sample['id']
if not self.test:
labels = sample['label'] # PIL images
labels_glb = masks_transform(labels)
width, height = images[0].size
if width > self.size_g[0] or height > self.size_g[1]:
images_global = resize(
images, self.size_g) # list of resized PIL images
else:
images_global = list(images)
if self.mode == 2 or self.mode == 3:
images_local = [image.copy() for image in images]
scores_local = [
np.zeros((1, self.n_class, images[i].size[1],
images[i].size[0])) for i in range(len(images))
]
scores = [
np.zeros((1, self.n_class, images[i].size[1],
images[i].size[0])) for i in range(len(images))
]
for flip in self.flip_range:
if flip:
# we already rotated images for 270'
for b in range(len(images)):
images_global[b] = transforms.functional.rotate(
images_global[b], 90) # rotate back!
images_global[b] = transforms.functional.hflip(
images_global[b])
if self.mode == 2 or self.mode == 3:
images_local[b] = transforms.functional.rotate(
images_local[b], 90) # rotate back!
images_local[b] = transforms.functional.hflip(
images_local[b])
for angle in self.rotate_range:
if angle > 0:
for b in range(len(images)):
images_global[b] = transforms.functional.rotate(
images_global[b], 90)
if self.mode == 2 or self.mode == 3:
images_local[b] = transforms.functional.rotate(
images_local[b], 90)
# prepare global images onto cuda
images_glb = images_transform(images_global) # b, c, h, w
if self.mode == 2 or self.mode == 3:
patches, coordinates, sizes, ratios, label_patches = global2patch(
images, labels_glb, self.size_p, ids)
predicted_patches = np.zeros((len(images), 4))
predicted_ensembles = np.zeros((len(images), 4))
outputs_global = [None for i in range(len(images))]
if self.mode == 1:
# eval with only resized global image ##########################
if flip:
outputs_global += np.flip(np.rot90(model.forward(
images_glb, None, None,
None)[0].data.cpu().numpy(),
k=angle,
axes=(3, 2)),
axis=3)
else:
outputs_global, _ = model.forward(
images_glb, None, None, None)
################################################################
if self.mode == 2:
# eval with patches ###########################################
for i in range(len(images)):
j = 0
while j < len(coordinates[i]):
patches_var = images_transform(
patches[i]
[j:j + self.sub_batch_size]) # b, c, h, w
output_ensembles, output_global, output_patches, _ = model.forward(
images_glb[i:i + 1],
patches_var,
coordinates[i][j:j + self.sub_batch_size],
ratios[i],
mode=self.mode,
n_patch=len(coordinates[i]))
# patch predictions
for pred_patch, pred_ensemble in zip(
torch.max(output_patches.data,
1)[1].data,
torch.max(output_ensembles.data,
1)[1].data):
predicted_patches[i][int(pred_patch)] += 1
predicted_ensembles[i][int(
pred_ensemble)] += 1
j += patches_var.size()[0]
outputs_global[i] = output_global
outputs_global = torch.cat(outputs_global, dim=0)
###############################################################
if self.mode == 3:
# eval global with help from patches ##################################################
# go through local patches to collect feature maps
# collect predictions from patches
for i in range(len(images)):
j = 0
while j < len(coordinates[i]):
patches_var = images_transform(
patches[i]
[j:j + self.sub_batch_size]) # b, c, h, w
_, output_patches = model.module.collect_local_fm(
images_glb[i:i + 1],
patches_var,
ratios[i],
coordinates[i],
[j, j + self.sub_batch_size],
len(images),
global_model=global_fixed,
n_patch_all=len(coordinates[i]))
for pred_patch in torch.max(
output_patches.data, 1)[1].data:
predicted_patches[i][int(pred_patch)] += 1
j += self.sub_batch_size
# go through global image
tmp, fm_global = model.forward(images_glb,
None,
None,
None,
mode=self.mode)
if flip:
outputs_global += np.flip(np.rot90(
tmp.data.cpu().numpy(), k=angle, axes=(3, 2)),
axis=3)
else:
outputs_global = tmp
# generate ensembles
for i in range(len(images)):
j = 0
while j < len(coordinates[i]):
patches_var = images_transform(
patches[i]
[j:j + self.sub_batch_size]) # b, c, h, w
fl = model.module.generate_local_fm(
images_glb[i:i + 1],
patches_var,
ratios[i],
coordinates[i],
[j, j + self.sub_batch_size],
len(images),
global_model=global_fixed,
n_patch_all=len(coordinates[i]))
fg = model.module._crop_global(
fm_global[i:i + 1],
coordinates[i][j:j + self.sub_batch_size],
ratios[i])[0]
fg = F.interpolate(fg,
size=fl.size()[2:],
mode='bilinear')
output_ensembles = model.module.ensemble(
fl, fg) # include cordinates
# ensemble predictions
for pred_ensemble in torch.max(
output_ensembles.data, 1)[1].data:
predicted_ensembles[i][int(
pred_ensemble)] += 1
j += self.sub_batch_size
###################################################
_, predictions_global = torch.max(outputs_global.data, 1)
if not self.test:
self.metrics_global.update(labels_glb, predictions_global)
if self.mode == 2 or self.mode == 3:
# patch predictions ###########################
predictions_local = predicted_patches.argmax(1)
if not self.test:
self.metrics_local.update(labels_glb, predictions_local)
###################################################
predictions = predicted_ensembles.argmax(1)
if not self.test:
self.metrics.update(labels_glb, predictions)
return predictions, predictions_global, predictions_local
else:
return None, predictions_global, None
class Trainer(object):
def __init__(self,
criterion,
optimizer,
n_class,
size_g,
size_p,
sub_batch_size=6,
mode=1,
lamb_fmreg=0.15):
self.criterion = criterion
self.optimizer = optimizer
self.metrics_global = ConfusionMatrix(n_class)
self.metrics_local = ConfusionMatrix(n_class)
self.metrics = ConfusionMatrix(n_class)
self.n_class = n_class
self.size_g = size_g
self.size_p = size_p
self.sub_batch_size = sub_batch_size
self.mode = mode
self.lamb_fmreg = lamb_fmreg
def set_train(self, model):
model.module.ensemble_conv.train()
if self.mode == 1 or self.mode == 3:
model.module.resnet_global.train()
model.module.fpn_global.train()
else:
model.module.resnet_local.train()
model.module.fpn_local.train()
def get_scores(self):
score_train = self.metrics.get_scores()
score_train_local = self.metrics_local.get_scores()
score_train_global = self.metrics_global.get_scores()
return score_train, score_train_global, score_train_local
def reset_metrics(self):
self.metrics.reset()
self.metrics_local.reset()
self.metrics_global.reset()
def train(self, sample, model, global_fixed):
images, labels = sample['image'], sample['label'] # PIL images
labels_npy = masks_transform(
labels, numpy=True) # label of origin size in numpy
images_glb = resize(images, self.size_g) # list of resized PIL images
images_glb = images_transform(images_glb)
labels_glb = resize(labels, (self.size_g[0] // 4, self.size_g[1] // 4),
label=True) # FPN down 1/4, for loss
labels_glb = masks_transform(labels_glb)
if self.mode == 2 or self.mode == 3:
patches, coordinates, templates, sizes, ratios = global2patch(
images, self.size_p)
label_patches, _, _, _, _ = global2patch(labels, self.size_p)
predicted_patches = [
np.zeros((len(coordinates[i]), self.n_class, self.size_p[0],
self.size_p[1])) for i in range(len(images))
]
predicted_ensembles = [
np.zeros((len(coordinates[i]), self.n_class, self.size_p[0],
self.size_p[1])) for i in range(len(images))
]
outputs_global = [None for i in range(len(images))]
if self.mode == 1:
# training with only (resized) global image #########################################
outputs_global, _ = model.forward(images_glb, None, None, None)
loss = self.criterion(outputs_global, labels_glb)
loss.backward()
self.optimizer.step()
self.optimizer.zero_grad()
##############################################
if self.mode == 2:
# training with patches ###########################################
for i in range(len(images)):
j = 0
while j < len(coordinates[i]):
patches_var = images_transform(
patches[i][j:j + self.sub_batch_size]) # b, c, h, w
label_patches_var = masks_transform(
resize(label_patches[i][j:j + self.sub_batch_size],
(self.size_p[0] // 4, self.size_p[1] // 4),
label=True)) # down 1/4 for loss
output_ensembles, output_global, output_patches, fmreg_l2 = model.forward(
images_glb[i:i + 1],
patches_var,
coordinates[i][j:j + self.sub_batch_size],
ratios[i],
mode=self.mode,
n_patch=len(coordinates[i]))
loss = self.criterion(
output_patches, label_patches_var) + self.criterion(
output_ensembles,
label_patches_var) + self.lamb_fmreg * fmreg_l2
loss.backward()
# patch predictions
predicted_patches[i][j:j + output_patches.size(
)[0]] = F.interpolate(output_patches,
size=self.size_p,
mode='nearest').data.cpu().numpy()
predicted_ensembles[i][j:j + output_ensembles.size(
)[0]] = F.interpolate(output_ensembles,
size=self.size_p,
mode='nearest').data.cpu().numpy()
j += self.sub_batch_size
outputs_global[i] = output_global
outputs_global = torch.cat(outputs_global, dim=0)
self.optimizer.step()
self.optimizer.zero_grad()
#####################################################################################
if self.mode == 3:
# train global with help from patches ##################################################
# go through local patches to collect feature maps
# collect predictions from patches
for i in range(len(images)):
j = 0
while j < len(coordinates[i]):
patches_var = images_transform(
patches[i][j:j + self.sub_batch_size]) # b, c, h, w
fm_patches, output_patches = model.module.collect_local_fm(
images_glb[i:i + 1],
patches_var,
ratios[i],
coordinates[i], [j, j + self.sub_batch_size],
len(images),
global_model=global_fixed,
template=templates[i],
n_patch_all=len(coordinates[i]))
predicted_patches[i][j:j + output_patches.size(
)[0]] = F.interpolate(output_patches,
size=self.size_p,
mode='nearest').data.cpu().numpy()
j += self.sub_batch_size
# train on global image
outputs_global, fm_global = model.forward(images_glb,
None,
None,
None,
mode=self.mode)
loss = self.criterion(outputs_global, labels_glb)
loss.backward(retain_graph=True)
# fmreg loss
# generate ensembles & calc loss
for i in range(len(images)):
j = 0
while j < len(coordinates[i]):
label_patches_var = masks_transform(
resize(label_patches[i][j:j + self.sub_batch_size],
(self.size_p[0] // 4, self.size_p[1] // 4),
label=True))
fl = fm_patches[i][j:j + self.sub_batch_size].cuda()
fg = model.module._crop_global(
fm_global[i:i + 1],
coordinates[i][j:j + self.sub_batch_size],
ratios[i])[0]
fg = F.interpolate(fg, size=fl.size()[2:], mode='bilinear')
output_ensembles = model.module.ensemble(fl, fg)
loss = self.criterion(
output_ensembles,
label_patches_var) # + 0.15 * mse(fl, fg)
if i == len(images) - 1 and j + self.sub_batch_size >= len(
coordinates[i]):
loss.backward()
else:
loss.backward(retain_graph=True)
# ensemble predictions
predicted_ensembles[i][j:j + output_ensembles.size(
)[0]] = F.interpolate(output_ensembles,
size=self.size_p,
mode='nearest').data.cpu().numpy()
j += self.sub_batch_size
self.optimizer.step()
self.optimizer.zero_grad()
# global predictions ###########################
outputs_global = outputs_global.cpu()
predictions_global = [
F.interpolate(outputs_global[i:i + 1],
images[i].size[::-1],
mode='nearest').argmax(1).detach().numpy()
for i in range(len(images))
]
self.metrics_global.update(labels_npy, predictions_global)
if self.mode == 2 or self.mode == 3:
# patch predictions ###########################
scores_local = np.array(
patch2global(predicted_patches, self.n_class, sizes,
coordinates, self.size_p)
) # merge softmax scores from patches (overlaps)
predictions_local = scores_local.argmax(1) # b, h, w
self.metrics_local.update(labels_npy, predictions_local)
###################################################
# combined/ensemble predictions ###########################
scores = np.array(
patch2global(predicted_ensembles, self.n_class, sizes,
coordinates, self.size_p)
) # merge softmax scores from patches (overlaps)
predictions = scores.argmax(1) # b, h, w
self.metrics.update(labels_npy, predictions)
return loss
def test(
data,
weights=None,
batch_size=16,
imgsz=640,
conf_thres=0.001,
iou_thres=0.6, # for NMS
save_json=False,
single_cls=False,
augment=False,
verbose=True,
model=None,
dataloader=None,
save_dir='',
merge=False,
save_txt=False,
plots=True,
log_imgs=16):
if not save_dir:
save_dir = Path(increment_dir('runs/test/exp/')) # increment run
save_dir.mkdir(parents=True, exist_ok=True) # make dir
print("save_dir:", save_dir)
else:
save_dir = Path(save_dir)
# Initialize/load model and set device
training = model is not None
with open(data) as f:
data = yaml.load(f, Loader=yaml.FullLoader) # model dict
nc = 1 if single_cls else int(data['nc']) # number of classes
if training: # called by train.py
device = next(model.parameters()).device # get model device
else: # called directly
device = torch_utils.select_device(opt.device, batch_size=batch_size)
merge, save_txt = opt.merge, opt.save_txt # use Merge NMS, save *.txt labels
if save_txt:
out = Path('inference/output')
if os.path.exists(out):
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
# Remove previous
for f in glob.glob(str(Path(save_dir) / 'test_batch*.jpg')):
os.remove(f)
# Load model
# try:
# model = attempt_load(weights, map_location=device) # load FP32 model
# except ModuleNotFoundError:
# model = Model(opt.cfg, nc=nc).to(device)
# model.load_state_dict(torch.load(weights))
model = attempt_load(weights, map_location=device) # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
# Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99
# if device.type != 'cpu' and torch.cuda.device_count() > 1:
# model = nn.DataParallel(model)
# Half
half = device.type != 'cpu' # half precision only supported on CUDA
if half:
model.half()
# Configure
model.eval()
iouv = torch.linspace(0.5, 0.95,
10).to(device) # iou vector for [email protected]:0.95
niou = iouv.numel()
# Dataloader
if not training:
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
path = data['test'] if opt.task == 'test' else data[
'val'] # path to val/test images
dataloader = create_dataloader(path,
imgsz,
batch_size,
model.stride.max(),
opt,
hyp=None,
augment=False,
cache=False,
pad=0.5,
rect=True)[0]
seen = 0
confusion_matrix = ConfusionMatrix(nc=nc)
names = model.names if hasattr(model, 'names') else model.module.names
wb_names = {
k: v
for k, v in enumerate(
model.names if hasattr(model, 'names') else model.module.names)
}
coco91class = coco80_to_coco91_class()
s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R',
'[email protected]', '[email protected]:.95')
p, r, f1, mp, mr, map50, map, t0, t1 = 0., 0., 0., 0., 0., 0., 0., 0., 0.
loss = torch.zeros(3, device=device)
jdict, stats, ap, ap_class, wandb_images = [], [], [], [], []
for batch_i, (img, targets, paths,
shapes) in enumerate(tqdm(dataloader, desc=s)):
img = img.to(device, non_blocking=True)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
targets = targets.to(device)
nb, _, height, width = img.shape # batch size, channels, height, width
whwh = torch.Tensor([width, height, width, height]).to(device)
# Disable gradients
with torch.no_grad():
# Run model
t = torch_utils.time_synchronized()
inf_out, train_out = model(
img, augment=augment) # inference and training outputs
t0 += torch_utils.time_synchronized() - t
# Compute loss
if training: # if model has loss hyperparameters
loss += compute_loss([x.float() for x in train_out], targets,
model)[1][:3] # GIoU, obj, cls
# Run NMS
t = torch_utils.time_synchronized()
output = non_max_suppression(inf_out,
conf_thres=conf_thres,
iou_thres=iou_thres,
merge=merge)
t1 += torch_utils.time_synchronized() - t
# Statistics per image
for si, pred in enumerate(output):
labels = targets[targets[:, 0] == si, 1:]
nl = len(labels)
tcls = labels[:, 0].tolist() if nl else [] # target class
seen += 1
if pred is None:
if nl:
stats.append((torch.zeros(0, niou, dtype=torch.bool),
torch.Tensor(), torch.Tensor(), tcls))
continue
# Append to text file
if save_txt:
gn = torch.tensor(shapes[si][0])[[1, 0, 1, 0
]] # normalization gain whwh
txt_path = str(out / Path(paths[si]).stem)
pred[:, :4] = scale_coords(img[si].shape[1:], pred[:, :4],
shapes[si][0],
shapes[si][1]) # to original
for *xyxy, conf, cls in pred:
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
with open(txt_path + '.txt', 'a') as f:
f.write(
('%g ' * 5 + '\n') % (cls, *xywh)) # label format
# W&B logging
if plots and len(wandb_images) < log_imgs:
path = Path(paths[si])
box_data = [{
"position": {
"minX": xyxy[0],
"minY": xyxy[1],
"maxX": xyxy[2],
"maxY": xyxy[3]
},
"class_id": int(cls),
"box_caption": "%s %.3f" % (wb_names[cls], conf),
"scores": {
"class_score": conf
},
"domain": "pixel"
} for *xyxy, conf, cls in pred.tolist()]
boxes = {
"predictions": {
"box_data": box_data,
"class_labels": wb_names
}
} # inference-space
wandb_images.append(
wandb.Image(img[si], boxes=boxes, caption=path.name))
# Clip boxes to image bounds
clip_coords(pred, (height, width))
# Append to pycocotools JSON dictionary
if save_json:
# [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ...
image_id = Path(paths[si]).stem
box = pred[:, :4].clone() # xyxy
scale_coords(img[si].shape[1:], box, shapes[si][0],
shapes[si][1]) # to original shape
box = xyxy2xywh(box) # xywh
box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner
for p, b in zip(pred.tolist(), box.tolist()):
jdict.append({
'image_id':
int(image_id) if image_id.isnumeric() else image_id,
'category_id':
coco91class[int(p[5])],
'bbox': [round(x, 3) for x in b],
'score':
round(p[4], 5)
})
# Assign all predictions as incorrect
correct = torch.zeros(pred.shape[0],
niou,
dtype=torch.bool,
device=device)
if nl:
detected = [] # target indices
tcls_tensor = labels[:, 0]
# target boxes
tbox = xywh2xyxy(labels[:, 1:5]) * whwh
if plots:
confusion_matrix.process_batch(
pred.clone(), torch.cat((labels[:, 0:1], tbox), 1))
# Per target class
for cls in torch.unique(tcls_tensor):
ti = (cls == tcls_tensor).nonzero(as_tuple=False).view(
-1) # prediction indices
pi = (cls == pred[:, 5]).nonzero(as_tuple=False).view(
-1) # target indices
# Search for detections
if pi.shape[0]:
# Prediction to target ious
ious, i = box_iou(pred[pi, :4], tbox[ti]).max(
1) # best ious, indices
# Append detections
for j in (ious > iouv[0]).nonzero(as_tuple=False):
d = ti[i[j]] # detected target
if d not in detected:
detected.append(d)
correct[
pi[j]] = ious[j] > iouv # iou_thres is 1xn
if len(
detected
) == nl: # all targets already located in image
break
# Append statistics (correct, conf, pcls, tcls)
stats.append(
(correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls))
# Plot images
if batch_i < 3:
f = save_dir / ('test_batch%g_gt.jpg' % batch_i) # filename
print("save batch gt", f)
plot_images(img, targets, paths, str(f), names) # ground truth
f = save_dir / ('test_batch%g_pred.jpg' % batch_i)
print("save batch label", f)
plot_images(img, output_to_target(output, width, height), paths,
str(f), names) # predictions
# Compute statistics
stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy
if len(stats) and stats[0].any():
p, r, ap, f1, ap_class = ap_per_class(*stats)
p, r, ap50, ap = p[:, 0], r[:, 0], ap[:, 0], ap.mean(
1) # [P, R, [email protected], [email protected]:0.95]
mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
nt = np.bincount(stats[3].astype(np.int64),
minlength=nc) # number of targets per class
else:
nt = torch.zeros(1)
# Print results
pf = '%20s' + '%12.3g' * 6 # print format
print(pf % ('all', seen, nt.sum(), mp, mr, map50, map))
# Print results per class
if (verbose or nc < 50) and nc > 1 and len(stats):
for i, c in enumerate(ap_class):
print(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i]))
# Print speeds
t = tuple(x / seen * 1E3
for x in (t0, t1, t0 + t1)) + (imgsz, imgsz, batch_size) # tuple
if not training:
print(
'Speed: %.1f/%.1f/%.1f ms inference/NMS/total per %gx%g image at batch-size %g'
% t)
# Plots
if plots:
confusion_matrix.plot(save_dir=save_dir, names=list(wb_names.values()))
if wandb and wandb.run:
val_batches = [
wandb.Image(str(f), caption=f.name)
for f in sorted(save_dir.glob('test*.jpg'))
]
wandb.log({
"Images": wandb_images,
"Validation": val_batches
},
commit=False)
# Save JSON
if save_json and len(jdict):
f = 'detections_val2017_%s_results.json' % \
(weights.split(os.sep)[-1].replace('.pt', '') if isinstance(weights, str) else '') # filename
print('\nCOCO mAP with pycocotools... saving %s...' % f)
with open(f, 'w') as file:
json.dump(jdict, file)
try: # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
imgIds = [int(Path(x).stem) for x in dataloader.dataset.img_files]
cocoGt = COCO(
glob.glob('../coco/annotations/instances_val*.json')
[0]) # initialize COCO ground truth api
cocoDt = cocoGt.loadRes(f) # initialize COCO pred api
cocoEval = COCOeval(cocoGt, cocoDt, 'bbox')
cocoEval.params.imgIds = imgIds # image IDs to evaluate
cocoEval.evaluate()
cocoEval.accumulate()
cocoEval.summarize()
map, map50 = cocoEval.stats[:
2] # update results ([email protected]:0.95, [email protected])
except Exception as e:
print('ERROR: pycocotools unable to run: %s' % e)
# Return results
model.float() # for training
maps = np.zeros(nc) + map
for i, c in enumerate(ap_class):
maps[c] = ap[i]
return (mp, mr, map50, map,
*(loss.cpu() / len(dataloader)).tolist()), maps, t
# Dataloader
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half()) # run once
path = data['test']
dataloader = create_dataloader(path,
imgsz,
batch_size,
model.stride.max(),
opt,
pad=0.5,
rect=True)[0]
"""
define metrics
"""
confusion_matrix = ConfusionMatrix(nc=nc)
names = {k: v for k, v in enumerate(model.names)}
s = ('%20s' + '%12s' * 7) % ('Class', 'Images', 'Targets', 'P', 'R',
'[email protected]', '[email protected]:.95', 'mean f1')
p, r, f1, mp, mr, map50, map, t0, t1 = 0., 0., 0., 0., 0., 0., 0., 0., 0.
stats = []
seen = 0
"""
# run inference
"""
for batch_i, (img, targets, paths,
shapes) in enumerate(tqdm(dataloader, desc=s)):
img = img.to(device, non_blocking=True)
img = img.half() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
targets = targets.to(device)
class Trainer(object):
# def __init__(self, criterion, optimizer, n_class, size0, size_g, size_p, n, sub_batch_size=6, mode=1, lamb_fmreg=0.15):
def __init__(self, criterion, optimizer, n_class, size_g, size_p, sub_batch_size=6, mode=1, lamb_fmreg=0.15):
self.criterion = criterion
self.optimizer = optimizer
self.metrics_global = ConfusionMatrix(n_class)
self.metrics_local = ConfusionMatrix(n_class)
self.metrics = ConfusionMatrix(n_class)
self.n_class = n_class
# self.size0 = size0
self.size_g = size_g
self.size_p = size_p
# self.n = n
self.sub_batch_size = sub_batch_size
self.mode = mode
self.lamb_fmreg = lamb_fmreg
# self.ratio = float(size_p[0]) / size0[0]
# self.step = (size0[0] - size_p[0]) // (n - 1)
# self.template, self.coordinates = template_patch2global(size0, size_p, n, self.step)
def set_train(self, model, parallel=True):
if not parallel:
model.ensemble_conv.train()
if self.mode == 1 or self.mode == 3:
model.resnet_global.train()
model.fpn_global.train()
else:
model.resnet_local.train()
model.fpn_local.train()
else:
model.module.ensemble_conv.train()
if self.mode == 1 or self.mode == 3:
model.module.resnet_global.train()
model.module.fpn_global.train()
else:
model.module.resnet_local.train()
model.module.fpn_local.train()
def get_scores(self):
score_train = self.metrics.get_scores()
score_train_local = self.metrics_local.get_scores()
score_train_global = self.metrics_global.get_scores()
return score_train, score_train_global, score_train_local
def reset_metrics(self):
self.metrics.reset()
self.metrics_local.reset()
self.metrics_global.reset()
def train(self, sample, model, global_fixed, parallel=True):
images, labels = sample['image'], sample['label'] # PIL images
labels_npy = masks_transform(labels, numpy=True) # label of origin size in numpy
images_glb = resize(images, self.size_g) # list of resized PIL images
images_glb = images_transform(images_glb)
labels_glb = resize(labels, (self.size_g[0] // 4, self.size_g[1] // 4), label=True) # down 1/4 for loss
# labels_glb = resize(labels, self.size_g, label=True) # must downsample image for reduced GPU memory
labels_glb = masks_transform(labels_glb)
if self.mode == 2 or self.mode == 3:
patches, coordinates, templates, sizes, ratios = global2patch(images, self.size_p)
label_patches, _, _, _, _ = global2patch(labels, self.size_p)
# patches, label_patches = global2patch(images, self.n, self.step, self.size_p), global2patch(labels, self.n, self.step, self.size_p)
# predicted_patches = [ np.zeros((self.n**2, self.n_class, self.size_p[0], self.size_p[1])) for i in range(len(images)) ]
# predicted_ensembles = [ np.zeros((self.n**2, self.n_class, self.size_p[0], self.size_p[1])) for i in range(len(images)) ]
predicted_patches = [ np.zeros((len(coordinates[i]), self.n_class, self.size_p[0], self.size_p[1])) for i in range(len(images)) ]
predicted_ensembles = [ np.zeros((len(coordinates[i]), self.n_class, self.size_p[0], self.size_p[1])) for i in range(len(images)) ]
outputs_global = [ None for i in range(len(images)) ]
if self.mode == 1:
# training with only (resized) global image #########################################
outputs_global, _ = model.forward(images_glb, None, None, None)
loss = self.criterion(outputs_global, labels_glb)
loss.backward()
self.optimizer.step()
self.optimizer.zero_grad()
##############################################
if self.mode == 2:
# training with patches ###########################################
for i in range(len(images)):
# sync the start for each global image
torch.distributed.barrier()
#print("Will start training global image:", i, " on rank :", torch.distributed.get_rank())
#group = torch.distributed.new_group(range(torch.distributed.get_world_size()))
coordinate_size = torch.tensor([len(coordinates[i])], dtype=torch.int32).cuda()
torch.distributed.all_reduce(coordinate_size, op=torch.distributed.ReduceOp.MAX)
#print("Will train:", coordinate_size.item(), "loops in global image.", "On rank :", torch.distributed.get_rank())
j = 0
_loop = 0
# while j < self.n**2:
#print("==== ==== len(coordinates[i]):", len(coordinates[i]), ", at rank:", torch.distributed.get_rank())
#while j < len(coordinates[i]):
#while j in range(6):
#while _ in range(coordinate_size.item()):
while _loop < coordinate_size.item():
patches_var = images_transform(patches[i][j : j+self.sub_batch_size]) # b, c, h, w
label_patches_var = masks_transform(resize(label_patches[i][j : j+self.sub_batch_size], (self.size_p[0] // 4, self.size_p[1] // 4), label=True)) # down 1/4 for loss
# label_patches_var = masks_transform(label_patches[i][j : j+self.sub_batch_size])
# output_ensembles, output_global, output_patches, fmreg_l2 = model.forward(images_glb[i:i+1], patches_var, self.coordinates[j : j+self.sub_batch_size], self.ratio, mode=self.mode, n_patch=self.n**2) # include cordinates
output_ensembles, output_global, output_patches, fmreg_l2 = model.forward(images_glb[i:i+1], patches_var, coordinates[i][j : j+self.sub_batch_size], ratios[i], mode=self.mode, n_patch=len(coordinates[i]))
loss = self.criterion(output_patches, label_patches_var) + self.criterion(output_ensembles, label_patches_var) + self.lamb_fmreg * fmreg_l2
loss.backward()
# patch predictions
predicted_patches[i][j:j+output_patches.size()[0]] = F.interpolate(output_patches, size=self.size_p, mode='nearest').data.cpu().numpy()
predicted_ensembles[i][j:j+output_ensembles.size()[0]] = F.interpolate(output_ensembles, size=self.size_p, mode='nearest').data.cpu().numpy()
# Because we choose loop the biggest coordinate_size in all ranks,
# make sure not cross the border for current rank
j = (j+self.sub_batch_size)%len(coordinates[i])
_loop += self.sub_batch_size
#print("==== ==== nested loop:", j, ", at rank:", torch.distributed.get_rank())
outputs_global[i] = output_global
outputs_global = torch.cat(outputs_global, dim=0)
self.optimizer.step()
self.optimizer.zero_grad()
#####################################################################################
if self.mode == 3:
# train global with help from patches ##################################################
# go through local patches to collect feature maps
# collect predictions from patches
for i in range(len(images)):
j = 0
# while j < self.n**2:
while j < len(coordinates[i]):
patches_var = images_transform(patches[i][j : j+self.sub_batch_size]) # b, c, h, w
if parallel:
# fm_patches, output_patches = model.module.collect_local_fm(images_glb[i:i+1], patches_var, self.ratio, self.coordinates, [j, j+self.sub_batch_size], len(images), global_model=global_fixed, template=self.template, n_patch_all=self.n**2) # include cordinates
fm_patches, output_patches = model.module.collect_local_fm(images_glb[i:i+1], patches_var, ratios[i], coordinates[i], [j, j+self.sub_batch_size], len(images), global_model=global_fixed, template=self.template, n_patch_all=len(coordinates[i]))
else:
# fm_patches, output_patches = model.module.collect_local_fm(images_glb[i:i+1], patches_var, self.ratio, self.coordinates, [j, j+self.sub_batch_size], len(images), global_model=global_fixed, template=self.template, n_patch_all=self.n**2) # include cordinates
fm_patches, output_patches = model.collect_local_fm(images_glb[i:i+1], patches_var, ratios[i], coordinates[i], [j, j+self.sub_batch_size], len(images), global_model=global_fixed, template=self.template, n_patch_all=len(coordinates[i]))
predicted_patches[i][j:j+output_patches.size()[0]] = F.interpolate(output_patches, size=self.size_p, mode='nearest').data.cpu().numpy()
j += self.sub_batch_size
# train on global image
outputs_global, fm_global = model.forward(images_glb, None, None, None, mode=self.mode)
loss = self.criterion(outputs_global, labels_glb)
loss.backward(retain_graph=True)
# fmreg loss
# generate ensembles & calc loss
for i in range(len(images)):
j = 0
# while j < self.n**2:
while j < len(coordinates[i]):
label_patches_var = masks_transform(resize(label_patches[i][j : j+self.sub_batch_size], (self.size_p[0] // 4, self.size_p[1] // 4), label=True))
# label_patches_var = masks_transform(resize(label_patches[i][j : j+self.sub_batch_size], self.size_p, label=True))
fl = fm_patches[i][j : j+self.sub_batch_size].cuda()
if parallel:
# fg = model.module._crop_global(fm_global[i:i+1], self.coordinates[j:j+self.sub_batch_size], self.ratio)[0]
fg = model.module._crop_global(fm_global[i:i+1], coordinates[i][j:j+self.sub_batch_size], ratios[i])[0]
else:
# fg = model.module._crop_global(fm_global[i:i+1], self.coordinates[j:j+self.sub_batch_size], self.ratio)[0]
fg = model._crop_global(fm_global[i:i+1], coordinates[i][j:j+self.sub_batch_size], ratios[i])[0]
fg = F.interpolate(fg, size=fl.size()[2:], mode='bilinear')
if parallel:
output_ensembles = model.module.ensemble(fl, fg)
# output_ensembles = F.interpolate(model.module.ensemble(fl, fg), self.size_p, **model.module._up_kwargs)
else:
output_ensembles = model.ensemble(fl, fg)
# output_ensembles = F.interpolate(model.module.ensemble(fl, fg), self.size_p, **model.module._up_kwargs)
loss = self.criterion(output_ensembles, label_patches_var)# + 0.15 * mse(fl, fg)
# if i == len(images) - 1 and j + self.sub_batch_size >= self.n**2:
if i == len(images) - 1 and j + self.sub_batch_size >= len(coordinates[i]):
loss.backward()
else:
loss.backward(retain_graph=True)
# ensemble predictions
predicted_ensembles[i][j:j+output_ensembles.size()[0]] = F.interpolate(output_ensembles, size=self.size_p, mode='nearest').data.cpu().numpy()
j += self.sub_batch_size
self.optimizer.step()
self.optimizer.zero_grad()
# global predictions ###########################
# predictions_global = F.interpolate(outputs_global.cpu(), self.size0, mode='nearest').argmax(1).detach().numpy()
outputs_global = outputs_global.cpu()
predictions_global = [F.interpolate(outputs_global[i:i+1], images[i].size[::-1], mode='nearest').argmax(1).detach().numpy() for i in range(len(images))]
self.metrics_global.update(labels_npy, predictions_global)
if self.mode == 2 or self.mode == 3:
# patch predictions ###########################
# scores_local = np.array(patch2global(predicted_patches, self.n_class, self.n, self.step, self.size0, self.size_p, len(images))) # merge softmax scores from patches (overlaps)
scores_local = np.array(patch2global(predicted_patches, self.n_class, sizes, coordinates, self.size_p)) # merge softmax scores from patches (overlaps)
predictions_local = scores_local.argmax(1) # b, h, w
self.metrics_local.update(labels_npy, predictions_local)
###################################################
# combined/ensemble predictions ###########################
# scores = np.array(patch2global(predicted_ensembles, self.n_class, self.n, self.step, self.size0, self.size_p, len(images))) # merge softmax scores from patches (overlaps)
scores = np.array(patch2global(predicted_ensembles, self.n_class, sizes, coordinates, self.size_p)) # merge softmax scores from patches (overlaps)
predictions = scores.argmax(1) # b, h, w
self.metrics.update(labels_npy, predictions)
return loss
class Evaluator(object):
# def __init__(self, n_class, size0, size_g, size_p, n, sub_batch_size=6, mode=1, test=False):
def __init__(self, n_class, size_g, size_p, sub_batch_size=6, mode=1, test=False):
self.metrics_global = ConfusionMatrix(n_class)
self.metrics_local = ConfusionMatrix(n_class)
self.metrics = ConfusionMatrix(n_class)
self.n_class = n_class
# self.size0 = size0
self.size_g = size_g
self.size_p = size_p
# self.n = n
self.sub_batch_size = sub_batch_size
self.mode = mode
self.test = test
# self.ratio = float(size_p[0]) / size0[0]
# self.step = (size0[0] - size_p[0]) // (n - 1)
# self.template, self.coordinates = template_patch2global(size0, size_p, n, self.step)
if test:
self.flip_range = [False, True]
self.rotate_range = [0, 1, 2, 3]
else:
self.flip_range = [False]
self.rotate_range = [0]
def get_scores(self):
score_train = self.metrics.get_scores()
score_train_local = self.metrics_local.get_scores()
score_train_global = self.metrics_global.get_scores()
return score_train, score_train_global, score_train_local
def reset_metrics(self):
self.metrics.reset()
self.metrics_local.reset()
self.metrics_global.reset()
def eval_test(self, sample, model, global_fixed):
with torch.no_grad():
images = sample["image"]
if not self.test:
labels = sample["label"] # PIL images
#lbls = [RGB_mapping_to_class(np.array(label)) for label in labels]
#labels = [Image.fromarray(lbl) for lbl in lbls]
labels_npy = masks_transform(labels, numpy=True)
images_global = resize(images, self.size_g)
outputs_global = np.zeros(
(len(images), self.n_class, self.size_g[0] // 4, self.size_g[1] // 4)
)
# outputs_global = np.zeros((len(images), self.n_class, self.size_g[0], self.size_g[1]))
if self.mode == 2 or self.mode == 3:
images_local = [image.copy() for image in images]
# scores_local = np.zeros((len(images), self.n_class, self.size0[0], self.size0[1]))
# scores = np.zeros((len(images), self.n_class, self.size0[0], self.size0[1]))
scores_local = [
np.zeros((1, self.n_class, images[i].size[1], images[i].size[0]))
for i in range(len(images))
]
scores = [
np.zeros((1, self.n_class, images[i].size[1], images[i].size[0]))
for i in range(len(images))
]
for flip in self.flip_range:
if flip:
# we already rotated images for 270'
for b in range(len(images)):
images_global[b] = transforms.functional.rotate(
images_global[b], 90
) # rotate back!
images_global[b] = transforms.functional.hflip(images_global[b])
if self.mode == 2 or self.mode == 3:
images_local[b] = transforms.functional.rotate(
images_local[b], 90
) # rotate back!
images_local[b] = transforms.functional.hflip(
images_local[b]
)
for angle in self.rotate_range:
if angle > 0:
for b in range(len(images)):
images_global[b] = transforms.functional.rotate(
images_global[b], 90
)
if self.mode == 2 or self.mode == 3:
images_local[b] = transforms.functional.rotate(
images_local[b], 90
)
# prepare global images onto cuda
images_glb = images_transform(images_global) # b, c, h, w
if self.mode == 2 or self.mode == 3:
# patches = global2patch(images_local, self.n, self.step, self.size_p)
patches, coordinates, templates, sizes, ratios = global2patch(
# patches, coordinates, _, sizes, ratios = global2patch(
images, self.size_p
)
# predicted_patches = [ np.zeros((self.n**2, self.n_class, self.size_p[0], self.size_p[1])) for i in range(len(images)) ]
# predicted_ensembles = [ np.zeros((self.n**2, self.n_class, self.size_p[0], self.size_p[1])) for i in range(len(images)) ]
predicted_patches = [
np.zeros(
(
len(coordinates[i]),
self.n_class,
self.size_p[0],
self.size_p[1],
)
)
for i in range(len(images))
]
predicted_ensembles = [
np.zeros(
(
len(coordinates[i]),
self.n_class,
self.size_p[0],
self.size_p[1],
)
)
for i in range(len(images))
]
if self.mode == 1:
# eval with only resized global image
if flip:
outputs_global += np.flip(
np.rot90(
model.forward(images_glb, None, None, None)[0]
.data.cpu()
.numpy(),
k=angle,
axes=(3, 2),
),
axis=3,
)
else:
outputs_global += np.rot90(
model.forward(images_glb, None, None, None)[0]
.data.cpu()
.numpy(),
k=angle,
axes=(3, 2),
)
if self.mode == 2:
# eval with patches
for i in range(len(images)):
j = 0
# while j < self.n**2:
while j < len(coordinates[i]):
patches_var = images_transform(
patches[i][j : j + self.sub_batch_size]
) # b, c, h, w
# output_ensembles, output_global, output_patches, _ = model.forward(images_glb[i:i+1], patches_var, self.coordinates[j : j+self.sub_batch_size], self.ratio, mode=self.mode, n_patch=self.n**2) # include cordinates
output_ensembles, output_global, output_patches, _ = model.forward(
images_glb[i : i + 1],
patches_var,
coordinates[i][j : j + self.sub_batch_size],
ratios[i],
mode=self.mode,
n_patch=len(coordinates[i]),
)
# patch predictions
predicted_patches[i][
j : j + output_patches.size()[0]
] += (
F.interpolate(
output_patches, size=self.size_p, mode="nearest"
)
.data.cpu()
.numpy()
)
predicted_ensembles[i][
j : j + output_ensembles.size()[0]
] += (
F.interpolate(
output_ensembles,
size=self.size_p,
mode="nearest",
)
.data.cpu()
.numpy()
)
j += patches_var.size()[0]
if flip:
outputs_global[i] += np.flip(
np.rot90(
output_global[0].data.cpu().numpy(),
k=angle,
axes=(2, 1),
),
axis=2,
)
# scores_local[i] += np.flip(np.rot90(np.array(patch2global(predicted_patches[i:i+1], self.n_class, self.n, self.step, self.size0, self.size_p, len(images))), k=angle, axes=(3, 2)), axis=3) # merge softmax scores from patches (overlaps)
# scores[i] += np.flip(np.rot90(np.array(patch2global(predicted_ensembles[i:i+1], self.n_class, self.n, self.step, self.size0, self.size_p, len(images))), k=angle, axes=(3, 2)), axis=3) # merge softmax scores from patches (overlaps)
scores_local[i] += np.flip(
np.rot90(
np.array(
patch2global(
predicted_patches[i : i + 1],
self.n_class,
sizes[i : i + 1],
coordinates[i : i + 1],
self.size_p,
)
),
k=angle,
axes=(3, 2),
),
axis=3,
) # merge softmax scores from patches (overlaps)
scores[i] += np.flip(
np.rot90(
np.array(
patch2global(
predicted_ensembles[i : i + 1],
self.n_class,
sizes[i : i + 1],
coordinates[i : i + 1],
self.size_p,
)
),
k=angle,
axes=(3, 2),
),
axis=3,
) # merge softmax scores from patches (overlaps)
else:
outputs_global[i] += np.rot90(
output_global[0].data.cpu().numpy(),
k=angle,
axes=(2, 1),
)
# scores_local[i] += np.rot90(np.array(patch2global(predicted_patches[i:i+1], self.n_class, self.n, self.step, self.size0, self.size_p, len(images))), k=angle, axes=(3, 2)) # merge softmax scores from patches (overlaps)
# scores[i] += np.rot90(np.array(patch2global(predicted_ensembles[i:i+1], self.n_class, self.n, self.step, self.size0, self.size_p, len(images))), k=angle, axes=(3, 2)) # merge softmax scores from patches (overlaps)
scores_local[i] += np.rot90(
np.array(
patch2global(
predicted_patches[i : i + 1],
self.n_class,
sizes[i : i + 1],
coordinates[i : i + 1],
self.size_p,
)
),
k=angle,
axes=(3, 2),
) # merge softmax scores from patches (overlaps)
scores[i] += np.rot90(
np.array(
patch2global(
predicted_ensembles[i : i + 1],
self.n_class,
sizes[i : i + 1],
coordinates[i : i + 1],
self.size_p,
)
),
k=angle,
axes=(3, 2),
) # merge softmax scores from patches (overlaps)
if self.mode == 3:
# eval global with help from patches
# go through local patches to collect feature maps
# collect predictions from patches
for i in range(len(images)):
j = 0
# while j < self.n**2:
while j < len(coordinates[i]):
patches_var = images_transform(
patches[i][j : j + self.sub_batch_size]
) # b, c, h, w
# fm_patches, output_patches = model.module.collect_local_fm(images_glb[i:i+1], patches_var, self.ratio, self.coordinates, [j, j+self.sub_batch_size], len(images), global_model=global_fixed, template=self.template, n_patch_all=self.n**2) # include cordinates
fm_patches, output_patches = model.module.collect_local_fm(
images_glb[i : i + 1],
patches_var,
ratios[i],
coordinates[i],
[j, j + self.sub_batch_size],
len(images),
global_model=global_fixed,
template=templates[0],
n_patch_all=len(coordinates[i]),
)
predicted_patches[i][
j : j + output_patches.size()[0]
] += (
F.interpolate(
output_patches, size=self.size_p, mode="nearest"
)
.data.cpu()
.numpy()
)
j += self.sub_batch_size
# go through global image
# tmp, fm_global = model.forward(images_glb, None, self.coordinates, self.ratio, mode=self.mode, global_model=None, n_patch=self.n**2) # include cordinates
tmp, fm_global = model.forward(
images_glb, None, None, None, mode=self.mode
)
if flip:
outputs_global += np.flip(
np.rot90(tmp.data.cpu().numpy(), k=angle, axes=(3, 2)),
axis=3,
)
else:
outputs_global += np.rot90(
tmp.data.cpu().numpy(), k=angle, axes=(3, 2)
)
# generate ensembles
for i in range(len(images)):
j = 0
# while j < self.n ** 2:
while j < len(coordinates[i]):
fl = fm_patches[i][j : j + self.sub_batch_size].cuda()
# fg = model.module._crop_global(fm_global[i:i+1], self.coordinates[j:j+self.sub_batch_size], self.ratio)[0]
fg = model.module._crop_global(
fm_global[i : i + 1],
coordinates[i][j : j + self.sub_batch_size],
ratios[i],
)[0]
fg = F.interpolate(
fg, size=fl.size()[2:], mode="bilinear"
)
output_ensembles = model.module.ensemble(
fl, fg
) # include cordinates
# output_ensembles = F.interpolate(model.module.ensemble(fl, fg), self.size_p, **model.module._up_kwargs)
# ensemble predictions
predicted_ensembles[i][
j : j + output_ensembles.size()[0]
] += (
F.interpolate(
output_ensembles,
size=self.size_p,
mode="nearest",
)
.data.cpu()
.numpy()
)
j += self.sub_batch_size
if flip:
# scores_local[i] += np.flip(np.rot90(np.array(patch2global(predicted_patches[i:i+1], self.n_class, self.n, self.step, self.size0, self.size_p, len(images))), k=angle, axes=(3, 2)), axis=3) # merge softmax scores from patches (overlaps)
# scores[i] += np.flip(np.rot90(np.array(patch2global(predicted_ensembles[i:i+1], self.n_class, self.n, self.step, self.size0, self.size_p, len(images))), k=angle, axes=(3, 2)), axis=3) # merge softmax scores from patches (overlaps)
scores_local[i] += np.flip(
np.rot90(
np.array(
patch2global(
predicted_patches[i : i + 1],
self.n_class,
sizes[i : i + 1],
coordinates[i : i + 1],
self.size_p,
)
),
k=angle,
axes=(3, 2),
),
axis=3,
)[
0
] # merge softmax scores from patches (overlaps)
scores[i] += np.flip(
np.rot90(
np.array(
patch2global(
predicted_ensembles[i : i + 1],
self.n_class,
sizes[i : i + 1],
coordinates[i : i + 1],
self.size_p,
)
),
k=angle,
axes=(3, 2),
),
axis=3,
)[
0
] # merge softmax scores from patches (overlaps)
else:
# scores_local[i] += np.rot90(np.array(patch2global(predicted_patches[i:i+1], self.n_class, self.n, self.step, self.size0, self.size_p, len(images))), k=angle, axes=(3, 2)) # merge softmax scores from patches (overlaps)
# scores[i] += np.rot90(np.array(patch2global(predicted_ensembles[i:i+1], self.n_class, self.n, self.step, self.size0, self.size_p, len(images))), k=angle, axes=(3, 2)) # merge softmax scores from patches (overlaps)
scores_local[i] += np.rot90(
np.array(
patch2global(
predicted_patches[i : i + 1],
self.n_class,
sizes[i : i + 1],
coordinates[i : i + 1],
self.size_p,
)
),
k=angle,
axes=(3, 2),
) # merge softmax scores from patches (overlaps)
scores[i] += np.rot90(
np.array(
patch2global(
predicted_ensembles[i : i + 1],
self.n_class,
sizes[i : i + 1],
coordinates[i : i + 1],
self.size_p,
)
),
k=angle,
axes=(3, 2),
) # merge softmax scores from patches (overlaps)
# global predictions
# predictions_global = F.interpolate(torch.Tensor(outputs_global), self.size0, mode='nearest').argmax(1).detach().numpy()
outputs_global = torch.Tensor(outputs_global)
predictions_global = [
F.interpolate(
outputs_global[i : i + 1], images[i].size[::-1], mode="nearest"
)
.argmax(1)
.detach()
.numpy()[0]
for i in range(len(images))
]
if not self.test:
self.metrics_global.update(labels_npy, predictions_global)
if self.mode == 2 or self.mode == 3:
# patch predictions
# predictions_local = scores_local.argmax(1) # b, h, w
predictions_local = [score.argmax(1)[0] for score in scores_local]
if not self.test:
self.metrics_local.update(labels_npy, predictions_local)
# combined/ensemble predictions
# predictions = scores.argmax(1) # b, h, w
predictions = [score.argmax(1)[0] for score in scores]
if not self.test:
self.metrics.update(labels_npy, predictions)
return predictions, predictions_global, predictions_local
else:
return None, predictions_global, None
def validate(self,
do_mirroring=True,
use_train_mode=False,
tiled=True,
step=2,
save_softmax=True,
use_gaussian=True,
compute_global_dice=True,
override=True,
validation_folder_name='validation'):
assert self.was_initialized, "must initialize, ideally with checkpoint (or train first)"
if self.dataset_val is None:
self.load_dataset()
self.do_split()
output_folder = join(self.output_folder, validation_folder_name)
maybe_mkdir_p(output_folder)
if do_mirroring:
mirror_axes = self.data_aug_params['mirror_axes']
else:
mirror_axes = ()
pred_gt_tuples = []
export_pool = Pool(4)
results = []
global_tp = OrderedDict()
global_fp = OrderedDict()
global_fn = OrderedDict()
for k in self.dataset_val.keys():
print(k)
properties = self.dataset[k]['properties']
fname = properties['list_of_data_files'][0].split("/")[-1][:-12]
if override or (not isfile(join(output_folder,
fname + ".nii.gz"))):
data = np.load(self.dataset[k]['data_file'])['data']
transpose_forward = self.plans.get('transpose_forward')
if transpose_forward is not None:
data = data.transpose([0] +
[i + 1 for i in transpose_forward])
print(k, data.shape)
data[-1][data[-1] == -1] = 0
softmax_pred = self.predict_preprocessing_return_softmax(
data[:-1],
do_mirroring,
1,
use_train_mode,
1,
mirror_axes,
tiled,
True,
step,
self.patch_size,
use_gaussian=use_gaussian)
if transpose_forward is not None:
transpose_backward = self.plans.get('transpose_backward')
softmax_pred = softmax_pred.transpose(
[0] + [i + 1 for i in transpose_backward])
if compute_global_dice:
predicted_segmentation = softmax_pred.argmax(0)
gt_segmentation = data[-1]
labels = properties['classes']
labels = [int(i) for i in labels if i > 0]
for l in labels:
if l not in global_fn.keys():
global_fn[l] = 0
if l not in global_fp.keys():
global_fp[l] = 0
if l not in global_tp.keys():
global_tp[l] = 0
conf = ConfusionMatrix(
(predicted_segmentation == l).astype(int),
(gt_segmentation == l).astype(int))
conf.compute()
global_fn[l] += conf.fn
global_fp[l] += conf.fp
global_tp[l] += conf.tp
if save_softmax:
softmax_fname = join(output_folder, fname + ".npz")
else:
softmax_fname = None
if np.prod(softmax_pred.shape) > (2e9 / 4 *
0.9): # *0.9 just to be save
np.save(join(output_folder, fname + ".npy"), softmax_pred)
softmax_pred = join(output_folder, fname + ".npy")
results.append(
export_pool.starmap_async(
store_seg_from_softmax,
((softmax_pred, join(output_folder,
fname + ".nii.gz"), properties, 3,
None, None, None, softmax_fname, None), )))
pred_gt_tuples.append([
join(output_folder, fname + ".nii.gz"),
join(self.gt_niftis_folder, fname + ".nii.gz")
])
_ = [i.get() for i in results]
print("finished prediction, now evaluating...")
task = self.dataset_directory.split("/")[-1]
job_name = self.experiment_name
_ = aggregate_scores(
pred_gt_tuples,
labels=list(range(self.num_classes)),
json_output_file=join(output_folder, "summary.json"),
json_name=job_name + " val tiled %s" % (str(tiled)),
json_author="Fabian",
json_task=task,
num_threads=3)
if compute_global_dice:
global_dice = OrderedDict()
all_labels = list(global_fn.keys())
for l in all_labels:
global_dice[int(l)] = float(
2 * global_tp[l] /
(2 * global_tp[l] + global_fn[l] + global_fp[l]))
write_json(global_dice, join(output_folder, "global_dice.json"))
def run(
data,
weights=None, # model.pt path(s)
batch_size=32, # batch size
imgsz=640, # inference size (pixels)
conf_thres=0.001, # confidence threshold
iou_thres=0.6, # NMS IoU threshold
task='val', # train, val, test, speed or study
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
single_cls=False, # treat as single-class dataset
augment=False, # augmented inference
verbose=False, # verbose output
save_txt=False, # save results to *.txt
save_hybrid=False, # save label+prediction hybrid results to *.txt
save_conf=False, # save confidences in --save-txt labels
save_json=False, # save a COCO-JSON results file
project=ROOT / 'runs/val', # save to project/name
name='exp', # save to project/name
exist_ok=False, # existing project/name ok, do not increment
half=True, # use FP16 half-precision inference
dnn=False, # use OpenCV DNN for ONNX inference
model=None,
dataloader=None,
save_dir=Path(''),
plots=True,
callbacks=Callbacks(),
compute_loss=None,
):
# Initialize/load model and set device
training = model is not None
if training: # called by train.py
device, pt = next(
model.parameters()).device, True # get model device, PyTorch model
half &= device.type != 'cpu' # half precision only supported on CUDA
model.half() if half else model.float()
else: # called directly
device = select_device(device, batch_size=batch_size)
# Directories
save_dir = increment_path(Path(project) / name,
exist_ok=exist_ok) # increment run
(save_dir / 'labels' if save_txt else save_dir).mkdir(
parents=True, exist_ok=True) # make dir
# Load model
model = DetectMultiBackend(weights, device=device, dnn=dnn)
stride, pt = model.stride, model.pt
imgsz = check_img_size(imgsz, s=stride) # check image size
half &= pt and device.type != 'cpu' # half precision only supported by PyTorch on CUDA
if pt:
model.model.half() if half else model.model.float()
else:
half = False
batch_size = 1 # export.py models default to batch-size 1
device = torch.device('cpu')
LOGGER.info(
f'Forcing --batch-size 1 square inference shape(1,3,{imgsz},{imgsz}) for non-PyTorch backends'
)
# Data
data = check_dataset(data) # check
# Configure
model.eval()
is_coco = isinstance(data.get('val'), str) and data['val'].endswith(
'coco/val2017.txt') # COCO dataset
nc = 1 if single_cls else int(data['nc']) # number of classes
iouv = torch.linspace(0.5, 0.95,
10).to(device) # iou vector for [email protected]:0.95
niou = iouv.numel()
# Dataloader
if not training:
if pt and device.type != 'cpu':
model(
torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(
next(model.model.parameters()))) # warmup
pad = 0.0 if task == 'speed' else 0.5
task = task if task in (
'train', 'val', 'test') else 'val' # path to train/val/test images
dataloader = create_dataloader(data[task],
imgsz,
batch_size,
stride,
single_cls,
pad=pad,
rect=pt,
prefix=colorstr(f'{task}: '))[0]
seen = 0
confusion_matrix = ConfusionMatrix(nc=nc)
names = {
k: v
for k, v in enumerate(
model.names if hasattr(model, 'names') else model.module.names)
}
class_map = coco80_to_coco91_class() if is_coco else list(range(1000))
s = ('%20s' + '%11s' * 6) % ('Class', 'Images', 'Labels', 'P', 'R',
'[email protected]', '[email protected]:.95')
dt, p, r, f1, mp, mr, map50, map = [0.0, 0.0,
0.0], 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
loss = torch.zeros(3, device=device)
jdict, stats, ap, ap_class = [], [], [], []
pbar = tqdm(dataloader,
desc=s,
ncols=NCOLS,
bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}') # progress bar
for batch_i, (im, targets, paths, shapes) in enumerate(pbar):
t1 = time_sync()
if pt:
im = im.to(device, non_blocking=True)
targets = targets.to(device)
im = im.half() if half else im.float() # uint8 to fp16/32
im /= 255 # 0 - 255 to 0.0 - 1.0
nb, _, height, width = im.shape # batch size, channels, height, width
t2 = time_sync()
dt[0] += t2 - t1
# Inference
out, train_out = model(im) if training else model(
im, augment=augment, val=True) # inference, loss outputs
dt[1] += time_sync() - t2
# Loss
if compute_loss:
loss += compute_loss([x.float() for x in train_out],
targets)[1] # box, obj, cls
# NMS
targets[:, 2:] *= torch.Tensor([width, height, width,
height]).to(device) # to pixels
lb = [targets[targets[:, 0] == i, 1:]
for i in range(nb)] if save_hybrid else [] # for autolabelling
t3 = time_sync()
out = non_max_suppression(out,
conf_thres,
iou_thres,
labels=lb,
multi_label=True,
agnostic=single_cls)
dt[2] += time_sync() - t3
# Metrics
for si, pred in enumerate(out):
labels = targets[targets[:, 0] == si, 1:]
nl = len(labels)
tcls = labels[:, 0].tolist() if nl else [] # target class
path, shape = Path(paths[si]), shapes[si][0]
seen += 1
if len(pred) == 0:
if nl:
stats.append((torch.zeros(0, niou, dtype=torch.bool),
torch.Tensor(), torch.Tensor(), tcls))
continue
# Predictions
if single_cls:
pred[:, 5] = 0
predn = pred.clone()
scale_coords(im[si].shape[1:], predn[:, :4], shape,
shapes[si][1]) # native-space pred
# Evaluate
if nl:
tbox = xywh2xyxy(labels[:, 1:5]) # target boxes
scale_coords(im[si].shape[1:], tbox, shape,
shapes[si][1]) # native-space labels
labelsn = torch.cat((labels[:, 0:1], tbox),
1) # native-space labels
correct = process_batch(predn, labelsn, iouv)
if plots:
confusion_matrix.process_batch(predn, labelsn)
else:
correct = torch.zeros(pred.shape[0], niou, dtype=torch.bool)
stats.append((correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(),
tcls)) # (correct, conf, pcls, tcls)
# Save/log
if save_txt:
save_one_txt(predn,
save_conf,
shape,
file=save_dir / 'labels' / (path.stem + '.txt'))
if save_json:
save_one_json(predn, jdict, path,
class_map) # append to COCO-JSON dictionary
callbacks.run('on_val_image_end', pred, predn, path, names, im[si])
# Plot images
if plots and batch_i < 3:
f = save_dir / f'val_batch{batch_i}_labels.jpg' # labels
Thread(target=plot_images,
args=(im, targets, paths, f, names),
daemon=True).start()
f = save_dir / f'val_batch{batch_i}_pred.jpg' # predictions
Thread(target=plot_images,
args=(im, output_to_target(out), paths, f, names),
daemon=True).start()
# Compute metrics
stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy
if len(stats) and stats[0].any():
p, r, ap, f1, ap_class = ap_per_class(*stats,
plot=plots,
save_dir=save_dir,
names=names)
ap50, ap = ap[:, 0], ap.mean(1) # [email protected], [email protected]:0.95
mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
nt = np.bincount(stats[3].astype(np.int64),
minlength=nc) # number of targets per class
else:
nt = torch.zeros(1)
# Print results
pf = '%20s' + '%11i' * 2 + '%11.3g' * 4 # print format
LOGGER.info(pf % ('all', seen, nt.sum(), mp, mr, map50, map))
# Print results per class
if (verbose or (nc < 50 and not training)) and nc > 1 and len(stats):
for i, c in enumerate(ap_class):
LOGGER.info(pf %
(names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i]))
# Print speeds
t = tuple(x / seen * 1E3 for x in dt) # speeds per image
if not training:
shape = (batch_size, 3, imgsz, imgsz)
LOGGER.info(
f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {shape}'
% t)
# Plots
if plots:
confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))
callbacks.run('on_val_end')
# Save JSON
if save_json and len(jdict):
w = Path(weights[0] if isinstance(weights, list) else weights
).stem if weights is not None else '' # weights
anno_json = str(
Path(data.get('path', '../coco')) /
'annotations/instances_val2017.json') # annotations json
pred_json = str(save_dir / f"{w}_predictions.json") # predictions json
LOGGER.info(f'\nEvaluating pycocotools mAP... saving {pred_json}...')
with open(pred_json, 'w') as f:
json.dump(jdict, f)
try: # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
check_requirements(['pycocotools'])
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
anno = COCO(anno_json) # init annotations api
pred = anno.loadRes(pred_json) # init predictions api
eval = COCOeval(anno, pred, 'bbox')
if is_coco:
eval.params.imgIds = [
int(Path(x).stem) for x in dataloader.dataset.img_files
] # image IDs to evaluate
eval.evaluate()
eval.accumulate()
eval.summarize()
map, map50 = eval.stats[:
2] # update results ([email protected]:0.95, [email protected])
except Exception as e:
LOGGER.info(f'pycocotools unable to run: {e}')
# Return results
model.float() # for training
if not training:
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
maps = np.zeros(nc) + map
for i, c in enumerate(ap_class):
maps[c] = ap[i]
return (mp, mr, map50, map,
*(loss.cpu() / len(dataloader)).tolist()), maps, t
def run(
data,
weights=None, # model.pt path(s)
batch_size=32, # batch size
imgsz=640, # inference size (pixels)
conf_thres=0.001, # confidence threshold
iou_thres=0.6, # NMS IoU threshold
task='val', # train, val, test, speed or study
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
single_cls=False, # treat as single-class dataset
augment=False, # augmented inference
verbose=False, # verbose output
save_txt=False, # save results to *.txt
save_hybrid=False, # save label+prediction hybrid results to *.txt
save_conf=False, # save confidences in --save-txt labels
save_json=False, # save a COCO-JSON results file
project='runs/val', # save to project/name
name='exp', # save to project/name
exist_ok=False, # existing project/name ok, do not increment
half=True, # use FP16 half-precision inference
model=None,
dataloader=None,
save_dir=Path(''),
plots=True,
wandb_logger=None,
compute_loss=None,
):
# Initialize/load model and set device
training = model is not None
if training: # called by train.py
device = next(model.parameters()).device # get model device
else: # called directly
device = select_device(device, batch_size=batch_size)
# Directories
save_dir = increment_path(Path(project) / name,
exist_ok=exist_ok) # increment run
(save_dir / 'labels' if save_txt else save_dir).mkdir(
parents=True, exist_ok=True) # make dir
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
gs = max(int(model.stride.max()), 32) # grid size (max stride)
imgsz = check_img_size(imgsz, s=gs) # check image size
# Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99
# if device.type != 'cpu' and torch.cuda.device_count() > 1:
# model = nn.DataParallel(model)
# Data
with open(data) as f:
data = yaml.safe_load(f)
check_dataset(data) # check
# Half
half &= device.type != 'cpu' # half precision only supported on CUDA
if half:
model.half()
# Configure
model.eval()
is_coco = type(data['val']) is str and data['val'].endswith(
'coco/val2017.txt') # COCO dataset
nc = 1 if single_cls else int(data['nc']) # number of classes
iouv = torch.linspace(0.5, 0.95,
10).to(device) # iou vector for [email protected]:0.95
niou = iouv.numel()
# Dataloader
if not training:
if device.type != 'cpu':
model(
torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(
next(model.parameters()))) # run once
task = task if task in (
'train', 'val', 'test') else 'val' # path to train/val/test images
dataloader = create_dataloader(data[task],
imgsz,
batch_size,
gs,
single_cls,
pad=0.5,
rect=True,
prefix=colorstr(f'{task}: '))[0]
seen = 0
confusion_matrix = ConfusionMatrix(nc=nc)
names = {
k: v
for k, v in enumerate(
model.names if hasattr(model, 'names') else model.module.names)
}
class_map = coco80_to_coco91_class() if is_coco else list(range(1000))
s = ('%20s' + '%11s' * 6) % ('Class', 'Images', 'Labels', 'P', 'R',
'[email protected]', '[email protected]:.95')
p, r, f1, mp, mr, map50, map, t0, t1, t2 = 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.
loss = torch.zeros(3, device=device)
jdict, stats, ap, ap_class = [], [], [], []
for batch_i, (img, targets, paths,
shapes) in enumerate(tqdm(dataloader, desc=s)):
t_ = time_sync()
img = img.to(device, non_blocking=True)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
targets = targets.to(device)
nb, _, height, width = img.shape # batch size, channels, height, width
t = time_sync()
t0 += t - t_
# Run model
out, train_out = model(
img, augment=augment) # inference and training outputs
t1 += time_sync() - t
# Compute loss
if compute_loss:
loss += compute_loss([x.float() for x in train_out],
targets)[1][:3] # box, obj, cls
# Run NMS
targets[:, 2:] *= torch.Tensor([width, height, width,
height]).to(device) # to pixels
lb = [targets[targets[:, 0] == i, 1:]
for i in range(nb)] if save_hybrid else [] # for autolabelling
t = time_sync()
out = non_max_suppression(out,
conf_thres,
iou_thres,
labels=lb,
multi_label=True,
agnostic=single_cls)
t2 += time_sync() - t
# Statistics per image
for si, pred in enumerate(out):
labels = targets[targets[:, 0] == si, 1:]
nl = len(labels)
tcls = labels[:, 0].tolist() if nl else [] # target class
path, shape = Path(paths[si]), shapes[si][0]
seen += 1
if len(pred) == 0:
if nl:
stats.append((torch.zeros(0, niou, dtype=torch.bool),
torch.Tensor(), torch.Tensor(), tcls))
continue
# Predictions
if single_cls:
pred[:, 5] = 0
predn = pred.clone()
scale_coords(img[si].shape[1:], predn[:, :4], shape,
shapes[si][1]) # native-space pred
# Evaluate
if nl:
tbox = xywh2xyxy(labels[:, 1:5]) # target boxes
scale_coords(img[si].shape[1:], tbox, shape,
shapes[si][1]) # native-space labels
labelsn = torch.cat((labels[:, 0:1], tbox),
1) # native-space labels
correct = process_batch(predn, labelsn, iouv)
if plots:
confusion_matrix.process_batch(predn, labelsn)
else:
correct = torch.zeros(pred.shape[0], niou, dtype=torch.bool)
stats.append((correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(),
tcls)) # (correct, conf, pcls, tcls)
# Save/log
if save_txt:
save_one_txt(predn,
save_conf,
shape,
file=save_dir / 'labels' / (path.stem + '.txt'))
if save_json:
save_one_json(predn, jdict, path,
class_map) # append to COCO-JSON dictionary
if wandb_logger and wandb_logger.wandb_run:
wandb_logger.val_one_image(pred, predn, path, names, img[si])
# Plot images
if plots and batch_i < 3:
f = save_dir / f'val_batch{batch_i}_labels.jpg' # labels
Thread(target=plot_images,
args=(img, targets, paths, f, names),
daemon=True).start()
f = save_dir / f'val_batch{batch_i}_pred.jpg' # predictions
Thread(target=plot_images,
args=(img, output_to_target(out), paths, f, names),
daemon=True).start()
# Compute statistics
stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy
if len(stats) and stats[0].any():
p, r, ap, f1, ap_class = ap_per_class(*stats,
plot=plots,
save_dir=save_dir,
names=names)
ap50, ap = ap[:, 0], ap.mean(1) # [email protected], [email protected]:0.95
mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
nt = np.bincount(stats[3].astype(np.int64),
minlength=nc) # number of targets per class
else:
nt = torch.zeros(1)
# Print results
pf = '%20s' + '%11i' * 2 + '%11.3g' * 4 # print format
print(pf % ('all', seen, nt.sum(), mp, mr, map50, map))
# Print results per class
if (verbose or (nc < 50 and not training)) and nc > 1 and len(stats):
for i, c in enumerate(ap_class):
print(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i]))
# Print speeds
t = tuple(x / seen * 1E3 for x in (t0, t1, t2)) # speeds per image
if not training:
shape = (batch_size, 3, imgsz, imgsz)
print(
f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {shape}'
% t)
# Plots
if plots:
confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))
if wandb_logger and wandb_logger.wandb:
val_batches = [
wandb_logger.wandb.Image(str(f), caption=f.name)
for f in sorted(save_dir.glob('val*.jpg'))
]
wandb_logger.log({"Validation": val_batches})
# Save JSON
if save_json and len(jdict):
w = Path(weights[0] if isinstance(weights, list) else weights
).stem if weights is not None else '' # weights
anno_json = str(
Path(data.get('path', '../coco')) /
'annotations/instances_val2017.json') # annotations json
pred_json = str(save_dir / f"{w}_predictions.json") # predictions json
print(f'\nEvaluating pycocotools mAP... saving {pred_json}...')
with open(pred_json, 'w') as f:
json.dump(jdict, f)
try: # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
check_requirements(['pycocotools'])
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
anno = COCO(anno_json) # init annotations api
pred = anno.loadRes(pred_json) # init predictions api
eval = COCOeval(anno, pred, 'bbox')
if is_coco:
eval.params.imgIds = [
int(Path(x).stem) for x in dataloader.dataset.img_files
] # image IDs to evaluate
eval.evaluate()
eval.accumulate()
eval.summarize()
map, map50 = eval.stats[:
2] # update results ([email protected]:0.95, [email protected])
except Exception as e:
print(f'pycocotools unable to run: {e}')
# Return results
model.float() # for training
if not training:
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
print(f"Results saved to {save_dir}{s}")
maps = np.zeros(nc) + map
for i, c in enumerate(ap_class):
maps[c] = ap[i]
return (mp, mr, map50, map,
*(loss.cpu() / len(dataloader)).tolist()), maps, t
def tst(
data,
weights=None,
batch_size=32,
imgsz=640,
conf_thres=0.001,
iou_thres=0.6, # for NMS
save_json=False,
single_cls=False,
augment=False,
verbose=False,
model=None,
dataloader=None,
save_dir=Path(''), # for saving images
save_txt=False, # for auto-labelling
save_hybrid=False, # for hybrid auto-labelling
save_conf=False, # save auto-label confidences
plots=True,
log_imgs=0): # number of logged images
# Initialize/load model and set device
training = model is not None
if training: # called by train.py
device = next(model.parameters()).device # get model device
else: # called directly
set_logging()
device = select_device(opt.device, batch_size=batch_size)
# Directories
save_dir = Path(
increment_path(Path(opt.project) / opt.name,
exist_ok=opt.exist_ok)) # increment run
(save_dir / 'labels' if save_txt else save_dir).mkdir(
parents=True, exist_ok=True) # make dir
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
# Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99
# if device.type != 'cpu' and torch.cuda.device_count() > 1:
# model = nn.DataParallel(model)
# Half
half = device.type != 'cpu' # half precision only supported on CUDA
if half:
model.half()
# Configure
model.eval()
is_coco = data.endswith('coco.yaml') # is COCO dataset
with open(data) as f:
data = yaml.load(f, Loader=yaml.FullLoader) # model dict
check_dataset(data) # check
nc = 1 if single_cls else int(data['nc']) # number of classes
iouv = torch.linspace(0.5, 0.95,
10).to(device) # iou vector for [email protected]:0.95
niou = iouv.numel()
# Logging
log_imgs, wandb = min(log_imgs, 100), None # ceil
try:
import wandb # Weights & Biases
except ImportError:
log_imgs = 0
# Dataloader
if not training:
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
_ = model(img.half() if half else img
) if device.type != 'cpu' else None # run once
path = data['test'] if opt.task == 'test' else data[
'val'] # path to val/test images
dataloader = create_dataloader(
path,
imgsz,
batch_size,
model.stride.max(),
opt,
pad=0.5,
rect=True,
prefix=colorstr('test: ' if opt.task == 'test' else 'val: '))[0]
seen = 0
confusion_matrix = ConfusionMatrix(nc=nc)
names = {
k: v
for k, v in enumerate(
model.names if hasattr(model, 'names') else model.module.names)
}
coco91class = coco80_to_coco91_class()
s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R',
'[email protected]', '[email protected]:.95')
p, r, f1, mp, mr, map50, map, t0, t1 = 0., 0., 0., 0., 0., 0., 0., 0., 0.
loss = torch.zeros(3, device=device)
jdict, stats, ap, ap_class, wandb_images = [], [], [], [], []
for batch_i, (img, targets, paths,
shapes) in enumerate(tqdm(dataloader, desc=s)):
img = img.to(device, non_blocking=True)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
targets = targets.to(device)
nb, _, height, width = img.shape # batch size, channels, height, width
with torch.no_grad():
# Run model
t = time_synchronized()
inf_out, train_out = model(
img, augment=augment) # inference and training outputs
t0 += time_synchronized() - t
# Compute loss
if training:
loss += compute_loss([x.float() for x in train_out], targets,
model)[1][:3] # box, obj, cls
# Run NMS
targets[:, 2:] *= torch.Tensor([width, height, width,
height]).to(device) # to pixels
lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)
] if save_hybrid else [] # for autolabelling
t = time_synchronized()
output = non_max_suppression(inf_out,
conf_thres=conf_thres,
iou_thres=iou_thres,
labels=lb)
t1 += time_synchronized() - t
# Statistics per image
for si, pred in enumerate(output):
labels = targets[targets[:, 0] == si, 1:]
nl = len(labels)
tcls = labels[:, 0].tolist() if nl else [] # target class
path = Path(paths[si])
seen += 1
if len(pred) == 0:
if nl:
stats.append((torch.zeros(0, niou, dtype=torch.bool),
torch.Tensor(), torch.Tensor(), tcls))
continue
# Predictions
predn = pred.clone()
scale_coords(img[si].shape[1:], predn[:, :4], shapes[si][0],
shapes[si][1]) # native-space pred
# Append to text file
if save_txt:
gn = torch.tensor(shapes[si][0])[[1, 0, 1, 0
]] # normalization gain whwh
for *xyxy, conf, cls in predn.tolist():
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) /
gn).view(-1).tolist() # normalized xywh
line = (cls, *xywh,
conf) if save_conf else (cls,
*xywh) # label format
with open(save_dir / 'labels' / (path.stem + '.txt'),
'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n')
# W&B logging
if plots and len(wandb_images) < log_imgs:
box_data = [{
"position": {
"minX": xyxy[0],
"minY": xyxy[1],
"maxX": xyxy[2],
"maxY": xyxy[3]
},
"class_id": int(cls),
"box_caption": "%s %.3f" % (names[cls], conf),
"scores": {
"class_score": conf
},
"domain": "pixel"
} for *xyxy, conf, cls in pred.tolist()]
boxes = {
"predictions": {
"box_data": box_data,
"class_labels": names
}
} # inference-space
wandb_images.append(
wandb.Image(img[si], boxes=boxes, caption=path.name))
# Append to pycocotools JSON dictionary
if save_json:
# [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ...
image_id = int(
path.stem) if path.stem.isnumeric() else path.stem
box = xyxy2xywh(predn[:, :4]) # xywh
box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner
for p, b in zip(pred.tolist(), box.tolist()):
jdict.append({
'image_id':
image_id,
'category_id':
coco91class[int(p[5])] if is_coco else int(p[5]),
'bbox': [round(x, 3) for x in b],
'score':
round(p[4], 5)
})
# Assign all predictions as incorrect
correct = torch.zeros(pred.shape[0],
niou,
dtype=torch.bool,
device=device)
if nl:
detected = [] # target indices
tcls_tensor = labels[:, 0]
# target boxes
tbox = xywh2xyxy(labels[:, 1:5])
scale_coords(img[si].shape[1:], tbox, shapes[si][0],
shapes[si][1]) # native-space labels
if plots:
confusion_matrix.process_batch(
pred, torch.cat((labels[:, 0:1], tbox), 1))
# Per target class
for cls in torch.unique(tcls_tensor):
ti = (cls == tcls_tensor).nonzero(as_tuple=False).view(
-1) # prediction indices
pi = (cls == pred[:, 5]).nonzero(as_tuple=False).view(
-1) # target indices
# Search for detections
if pi.shape[0]:
# Prediction to target ious
ious, i = box_iou(predn[pi, :4], tbox[ti]).max(
1) # best ious, indices
# Append detections
detected_set = set()
for j in (ious > iouv[0]).nonzero(as_tuple=False):
d = ti[i[j]] # detected target
if d.item() not in detected_set:
detected_set.add(d.item())
detected.append(d)
correct[
pi[j]] = ious[j] > iouv # iou_thres is 1xn
if len(
detected
) == nl: # all targets already located in image
break
# Append statistics (correct, conf, pcls, tcls)
stats.append(
(correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls))
# Plot images
if plots and batch_i < 3:
f = save_dir / f'test_batch{batch_i}_labels.jpg' # labels
Thread(target=plot_images,
args=(img, targets, paths, f, names),
daemon=True).start()
f = save_dir / f'test_batch{batch_i}_pred.jpg' # predictions
Thread(target=plot_images,
args=(img, output_to_target(output), paths, f, names),
daemon=True).start()
# Compute statistics
stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy
if len(stats) and stats[0].any():
p, r, ap, f1, ap_class = ap_per_class(*stats,
plot=plots,
save_dir=save_dir,
names=names)
p, r, ap50, ap = p[:, 0], r[:, 0], ap[:, 0], ap.mean(
1) # [P, R, [email protected], [email protected]:0.95]
mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
nt = np.bincount(stats[3].astype(np.int64),
minlength=nc) # number of targets per class
else:
nt = torch.zeros(1)
# Print results
pf = '%20s' + '%12.3g' * 6 # print format
print(pf % ('all', seen, nt.sum(), mp, mr, map50, map))
# Print results per class
if (verbose or (nc <= 20 and not training)) and nc > 1 and len(stats):
for i, c in enumerate(ap_class):
print(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i]))
# Print speeds
t = tuple(x / seen * 1E3
for x in (t0, t1, t0 + t1)) + (imgsz, imgsz, batch_size) # tuple
if not training:
print(
'Speed: %.1f/%.1f/%.1f ms inference/NMS/total per %gx%g image at batch-size %g'
% t)
# Plots
if plots:
confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))
if wandb and wandb.run:
wandb.log({"Images": wandb_images})
wandb.log({
"Validation": [
wandb.Image(str(f), caption=f.name)
for f in sorted(save_dir.glob('test*.jpg'))
]
})
# Save JSON
if save_json and len(jdict):
w = Path(weights[0] if isinstance(weights, list) else weights
).stem if weights is not None else '' # weights
anno_json = '../coco/annotations/instances_val2017.json' # annotations json
pred_json = str(save_dir / f"{w}_predictions.json") # predictions json
print('\nEvaluating pycocotools mAP... saving %s...' % pred_json)
with open(pred_json, 'w') as f:
json.dump(jdict, f)
try: # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
anno = COCO(anno_json) # init annotations api
pred = anno.loadRes(pred_json) # init predictions api
eval = COCOeval(anno, pred, 'bbox')
if is_coco:
eval.params.imgIds = [
int(Path(x).stem) for x in dataloader.dataset.img_files
] # image IDs to evaluate
eval.evaluate()
eval.accumulate()
eval.summarize()
map, map50 = eval.stats[:
2] # update results ([email protected]:0.95, [email protected])
except Exception as e:
print(f'pycocotools unable to run: {e}')
# Return results
if not training:
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
print(f"Results saved to {save_dir}{s}")
model.float() # for training
maps = np.zeros(nc) + map
for i, c in enumerate(ap_class):
maps[c] = ap[i]
return (mp, mr, map50, map,
*(loss.cpu() / len(dataloader)).tolist()), maps, t
class Evaluator(object):
def __init__(self, n_class, sub_batchsize, mode=1, test=False):
self.metrics_global = ConfusionMatrix(n_class)
self.metrics_local = ConfusionMatrix(n_class)
self.metrics = ConfusionMatrix(n_class)
self.n_class = n_class
self.sub_batchsize = sub_batchsize
self.mode = mode
self.test = test
def get_scores(self):
score_train = self.metrics.get_scores()
score_train_local = self.metrics_local.get_scores()
score_train_global = self.metrics_global.get_scores()
return score_train, score_train_global, score_train_local
def reset_metrics(self):
self.metrics.reset()
self.metrics_local.reset()
self.metrics_global.reset()
def eval_test(self, sample, model):
with torch.no_grad():
ids = sample['id']
h, w = sample['output_size'][0]
if not self.test:
labels = sample['label'].squeeze(1).long()
labels_npy = np.array(labels)
if self.mode == 1: # global
img_g = sample['image_g'].cuda()
outputs_g = model.forward(img_g)
outputs_g = F.interpolate(outputs_g, size=(h, w), mode='bilinear')
if self.mode == 2: # local
img_l = sample['image_l'].cuda()
batch_size = img_l.size(0)
idx = 0
outputs_l = []
while idx+self.sub_batchsize <= batch_size:
output_l = model.forward(img_l[idx:idx+self.sub_batchsize])
output_l = F.interpolate(output_l, size=(h, w), mode='bilinear')
outputs_l.append(output_l)
idx += self.sub_batchsize
outputs_l = torch.cat(outputs_l, dim=0)
if self.mode == 3: # global&local
img_g = sample['image_g'].cuda()
img_l = sample['image_l'].cuda()
batch_size = img_l.size(0)
idx = 0
outputs = []; outputs_g = []; outputs_l = []
while idx+self.sub_batchsize <= batch_size:
output, output_g, output_l, mse = model.forward(img_g[idx:idx+self.sub_batchsize], img_l[idx:idx+self.sub_batchsize])
outputs.append(output); outputs_g.append(output_g); outputs_l.append(output_l)
idx += self.sub_batchsize
outputs = torch.cat(outputs, dim=0); outputs_g = torch.cat(outputs_g, dim=0); outputs_l = torch.cat(outputs_l, dim=0)
# no target
# outputs, outputs_g, outputs_l, mse = model.forward(img_g, img_l)
# predictions
if self.mode == 1:
outputs_g = outputs_g.cpu()
predictions_global = [outputs_g[i:i+1].argmax(1).detach().numpy() for i in range(len(labels))]
if not self.test:
self.metrics_global.update(labels_npy, predictions_global)
return None, predictions_global, None
if self.mode == 2:
outputs_l = outputs_l.cpu()
predictions_local = [outputs_l[i:i+1].argmax(1).detach().numpy() for i in range(len(labels))]
if not self.test:
self.metrics_local.update(labels_npy, predictions_local)
return None, None, predictions_local
if self.mode == 3:
outputs_g = outputs_g.cpu(); outputs_l = outputs_l.cpu(); outputs = outputs.cpu()
predictions_global = [outputs_g[i:i+1].argmax(1).detach().numpy() for i in range(len(labels))]
predictions_local = [outputs_l[i:i+1].argmax(1).detach().numpy() for i in range(len(labels))]
predictions = [outputs[i:i+1].argmax(1).detach().numpy() for i in range(len(labels))]
if not self.test:
self.metrics_global.update(labels_npy, predictions_global)
self.metrics_local.update(labels_npy, predictions_local)
self.metrics.update(labels_npy, predictions)
return predictions, predictions_global, predictions_local
