def append_visulize(self,
imgpath,
boxesPre,
labelsPre,
scoresPre,
boxGT,
labelGT,
savepath=None):
imPre = np.array(Image.open(imgpath).convert('RGB'))
imGT = imPre.copy()
# scoreGT = np.ones(shape=(labelGT.shape[0],))
visualize_boxes(image=imPre,
boxes=boxesPre,
labels=labelsPre,
probs=scoresPre,
class_labels=self.cateNames)
# visualize_boxes(image=imGT, boxes=boxGT, labels=labelGT, probs=scoreGT,
# class_labels=self.cateNames)
whitepad = np.zeros(shape=(imPre.shape[0], 10, 3), dtype=np.uint8)
imshow = np.concatenate((imGT, whitepad, imPre), axis=1)
self.visual_imgs.append(imshow)
# plt.imshow(imPre)
# plt.show()
if savepath:
import os
plt.imsave(
os.path.join(savepath, '{}.png'.format(len(self.visual_imgs))),
imshow)
def evaluate(cfg, args):
device = select_device(args.device)
# Initialize/load model
if cfg.MODEL.META_ARCHITECTURE:
# Initialize model
model = YOLOv3(cfg).to(device)
# Load weights
if cfg.TEST.WEIGHTS.endswith(".pth"):
state = torch.load(cfg.TEST.WEIGHTS, map_location=device)
model.load_state_dict(state["state_dict"])
else:
load_darknet_weights(model, cfg.TEST.WEIGHTS)
if device.type != "cpu" and torch.cuda.device_count() > 1:
model = nn.DataParallel(model)
else:
warnings.warn(
"WARNNING: Backbone network cannot be empty! "
f"Default load Darknet53 meta architecture for `{cfg.CONFIG_FILE}`!"
)
model = YOLOv3(cfg).to(device)
if cfg.TEST.TASK == "visual":
images = os.listdir(os.path.join(os.getcwd(), "data", "test"))
for filename in images:
path = os.path.join(
os.path.join(os.getcwd(), "data", "test", filename))
images = cv2.imread(path)
assert images is not None
bboxes_prd = Evaluator(model, cfg=cfg).get_bbox(images)
if bboxes_prd.shape[0] != 0:
boxes = bboxes_prd[..., :4]
class_inds = bboxes_prd[..., 5].astype(np.int32)
scores = bboxes_prd[..., 4]
visualize_boxes(image=images,
boxes=boxes,
labels=class_inds,
probs=scores,
class_labels=cfg.CLASSES)
path = os.path.join(f"./outputs/{filename}")
cv2.imwrite(path, images)
elif cfg.TEST.TASK == "eval":
maps = 0.
with torch.no_grad():
aps = Evaluator(model, visiual=True,
cfg=cfg).calculate_aps(cfg.TEST.MULTI_SCALE,
cfg.TEST.FLIP)
for i in aps:
print(f"{i:25s} --> mAP : {aps[i]:.4f}")
maps += aps[i]
maps = maps / len(cfg.CLASSES)
print(f'mAP:{maps:.6f}')
return maps
def _inference_epoch(self,
imgdir,
outdir=None,
multiscale=True,
flip=True):
from dataset import get_imgdir
from utils.visualize import visualize_boxes
self.model.eval()
dataloader = get_imgdir(imgdir, batch_size=8, net_size=self.net_size)
for i, (imgpath, imgs, ori_shapes) in enumerate(dataloader):
if i == 50:
break
ori_shapes = ori_shapes.cuda()
if not multiscale:
INPUT_SIZES = [self.net_size]
else:
INPUT_SIZES = [
self.net_size - 32, self.net_size, self.net_size + 32
]
pyramids = makeImgPyramids(imgs.numpy().transpose(0, 2, 3, 1),
scales=INPUT_SIZES,
flip=flip)
# produce outputFeatures for each scale
img2multi = defaultdict(list)
for idx, pyramid in enumerate(pyramids):
pyramid = torch.from_numpy(pyramid.transpose(0, 3, 1,
2)).cuda()
with torch.no_grad():
grids = self.model(pyramid)
for imgidx in range(imgs.shape[0]):
img2multi[imgidx].append([grid[imgidx] for grid in grids])
# append prediction for each image per scale/flip
for imgidx, scalegrids in img2multi.items():
allboxes = []
allscores = []
for _grids, _scale in zip(scalegrids[:len(INPUT_SIZES)],
INPUT_SIZES):
_boxes, _scores = predict_yolo(
_grids,
self.anchors,
_scale,
ori_shapes[imgidx],
num_classes=self.num_classes)
allboxes.append(_boxes)
allscores.append(_scores)
if flip:
for _grids, _scale in zip(scalegrids[len(INPUT_SIZES):],
INPUT_SIZES):
_boxes, _scores = predict_yolo(
_grids,
self.anchors,
_scale,
ori_shapes[imgidx],
num_classes=self.num_classes)
_boxes = bbox_flip(_boxes.squeeze(0),
flip_x=True,
size=ori_shapes[imgidx])
_boxes = _boxes[np.newaxis, :]
allboxes.append(_boxes)
allscores.append(_scores)
nms_boxes, nms_scores, nms_labels = torch_nms(
torch.cat(allboxes, dim=1),
torch.cat(allscores, dim=1),
num_classes=self.num_classes)
if nms_boxes is not None:
detected_img = visualize_boxes(
np.array(Image.open(imgpath[imgidx]).convert('RGB')),
boxes=nms_boxes.cpu().numpy(),
labels=nms_labels.cpu().numpy(),
probs=nms_scores.cpu().numpy(),
class_labels=self.labels)
if outdir is not None:
plt.imsave(
os.path.join(outdir,
imgpath[imgidx].split('/')[-1]),
detected_img)