def eval_coco(self, val_dataset):
index2category = json.load(open("coco_index2category.json"))
logging.info('Start Evaling')
coco_result = []
coco_img_ids = set([])
for step, samples in enumerate(val_dataset):
images, labels = samples['image'], samples['label']
image_size = images.size(2)
image_paths, origin_sizes = samples['image_path'], samples['origin_size']
with torch.no_grad():
outputs = self.net(images)
#output = self.yolo_loss(outputs)
output_list = []
for i in range(3):
output_list.append(self.yolo_loss[i](outputs[i]))
output = torch.cat(output_list, 1)
batch_detections = non_max_suppression(output, self.config.num_classes, conf_thres=0.001, nms_thres=0.45)
for idx, detections in enumerate(batch_detections):
image_id = int(os.path.basename(image_paths[idx])[-16:-4])
coco_img_ids.add(image_id)
if detections is not None:
origin_size = eval(origin_sizes[idx])
detections = detections.cpu().numpy()
dim_diff = np.abs(origin_size[0] - origin_size[1])
pad1, pad2 = dim_diff // 2, dim_diff - dim_diff // 2
pad = ((pad1, pad2), (0, 0), (0, 0)) if origin_size[1] <= origin_size[0] else ((0, 0), (pad1, pad2), (0, 0))
scale = origin_size[0] if origin_size[1] <= origin_size[0] else origin_size[1]
for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
x1 = x1 / self.config.image_size * scale
x2 = x2 / self.config.image_size * scale
y1 = y1 / self.config.image_size * scale
y2 = y2 / self.config.image_size * scale
x1 -= pad[1][0]
y1 -= pad[0][0]
x2 -= pad[1][0]
y2 -= pad[0][0]
w = x2 - x1
h = y2 - y1
coco_result.append({
"image_id":image_id,
"category_id":index2category[str(int(cls_pred.item()))],
"bbox":(float(x1), float(y1), float(w), float(h)),
"score":float(conf),
})
logging.info("Now have finished [%.3d/%.3d]"%(step, len(val_dataset)))
save_path = "coco_results.json"
with open(save_path, "w") as f:
json.dump(coco_result, f, sort_keys=True, indent=4, separators=(',', ':'))
logging.info('Save result in {}'.format(save_path))
logging.info('Using COCO APi to evaluate')
cocoGt = COCO(self.config.annotation)
cocoDt = cocoGt.loadRes(save_path)
cocoEval = COCOeval(cocoGt, cocoDt, "bbox")
cocoEval.params.imgIds = list(coco_img_ids)
cocoEval.evaluate()
cocoEval.accumulate()
cocoEval.summarize()
def validate(net):
n_gt = 0
correct = 0
for step, samples in enumerate(dataloader):
images, labels, image_paths = samples["image"], samples[
"label"], samples["img_path"]
labels = labels.cuda()
with torch.no_grad():
time1 = datetime.datetime.now()
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(yolo_losses[i](outputs[i]))
output = torch.cat(output_list, 1)
output = non_max_suppression(output, 1, conf_thres=0.5)
if ((datetime.datetime.now() - time1).seconds > 5):
logging.info('Batch %d time is too long ' % (step))
n_gt = 1
break
# print("time2", (datetime.datetime.now() - time1).seconds*1000+(datetime.datetime.now() - time1).microseconds//1000)
# calculate
for sample_i in range(labels.size(0)):
# Get labels for sample where width is not zero (dummies)
target_sample = labels[sample_i, labels[sample_i, :, 3] != 0]
for obj_cls, tx, ty, tw, th in target_sample:
# Get rescaled gt coordinates
tx1, tx2 = config["img_w"] * (
tx - tw / 2), config["img_w"] * (tx + tw / 2)
ty1, ty2 = config["img_h"] * (
ty - th / 2), config["img_h"] * (ty + th / 2)
n_gt += 1
box_gt = torch.cat([
coord.unsqueeze(0) for coord in [tx1, ty1, tx2, ty2]
]).view(1, -1)
sample_pred = output[sample_i]
if sample_pred is not None:
# Iterate through predictions where the class predicted is same as gt
for x1, y1, x2, y2, conf, obj_conf, obj_pred in sample_pred[
sample_pred[:, 6] == obj_cls]:
box_pred = torch.cat([
coord.unsqueeze(0)
for coord in [x1, y1, x2, y2]
]).view(1, -1)
iou = bbox_iou(box_pred, box_gt)
if iou >= config["iou_thres"]:
correct += 1
break
if n_gt:
logging.info('Batch [%d/%d] mAP: %.5f' %
(step, len(dataloader), float(correct / n_gt)))
logging.info('Mean Average Precision: %.5f' % float(correct / n_gt))
def inference(self, image, classes, colors):
image_origin = image
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = cv2.resize(image,
(self.config.image_size, self.config.image_size),
interpolation=cv2.INTER_LINEAR)
image = np.expand_dims(image, 0)
image = image.astype(np.float32)
image /= 255
image = np.transpose(image, (0, 3, 1, 2))
image = image.astype(np.float32)
image = torch.from_numpy(image)
start_time = time.time()
if torch.cuda.is_available():
image = image.cuda()
with torch.no_grad():
outputs = self.net(image)
output_list = []
for i in range(3):
output_list.append(self.yolo_loss[i](outputs[i]))
output = torch.cat(output_list, 1)
batch_detections = non_max_suppression(output,
self.config.num_classes,
conf_thres=0.5,
nms_thres=0.4)
spand_time = float(time.time() - start_time)
detection = batch_detections[0]
if detection is not None:
origin_size = image_origin.shape[:2]
detection = detection.cpu().numpy()
for x1, y1, x2, y2, conf, cls_conf, cls_pred in detection:
x1 = int(x1 / self.config.image_size * origin_size[1])
x2 = int(x2 / self.config.image_size * origin_size[1])
y1 = int(y1 / self.config.image_size * origin_size[0])
y2 = int(y2 / self.config.image_size * origin_size[0])
color = colors[int(cls_pred)]
image_origin = cv2.rectangle(image_origin, (x1, y1), (x2, y2),
color, 3)
image_origin = cv2.rectangle(image_origin, (x1, y1),
(x2, y1 + 20),
color,
thickness=-1)
caption = "{}:{:.2f}".format(classes[int(cls_pred)], cls_conf)
image_origin = cv2.putText(image_origin, caption,
(x1, y1 + 15),
cv2.FONT_HERSHEY_SIMPLEX, 0.6,
(255, 255, 255), 2)
return image_origin, spand_time
def detect(config):
is_training = False
# Load and initialize network
net = ModelMain(config, is_training=is_training)
net.train(is_training)
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
# Restore pretrain model
if config["pretrain_snapshot"]:
state_dict = torch.load(config["pretrain_snapshot"])
net.load_state_dict(state_dict)
else:
logging.warning("missing pretrain_snapshot!!!")
# YOLO loss with 3 scales
yolo_losses = []
for i in range(3):
yolo_losses.append(
YOLOLoss(config["yolo"]["anchors"][i], config["yolo"]["classes"],
(config["img_w"], config["img_h"])))
# Load tested img
imgfile = config["img_path"]
img = Image.open(imgfile).convert('RGB')
resized = img.resize((config["img_w"], config["img_h"]))
input = image2torch(resized)
input = input.to(torch.device("cuda"))
start = time.time()
outputs = net(input)
output_list = []
for i in range(3):
output_list.append(yolo_losses[i](outputs[i]))
output = torch.cat(output_list, 1)
output = non_max_suppression(output,
config["yolo"]["classes"],
conf_thres=0.5,
nms_thres=0.4)
finish = time.time()
print('%s: Predicted in %f seconds.' % (imgfile, (finish - start)))
namefile = config["classname_path"]
class_names = load_class_names(namefile)
plot_boxes(img, output, 'predictions.jpg', class_names)
def ref(config):
"""
:param config:
:return:
"""
# File name
img_fn = '/home/yz/cde/ProposalYOLO/feature/ref/img-ref-attn.txt'
ftr_fn = '/home/yz/cde/ProposalYOLO/feature/ref/ftr-ref-attn.txt'
img_strm = open(img_fn, 'a')
ftr_strm = open(ftr_fn, 'a')
# Load and initialize network
is_training = False
# net = ProposalModel(config, is_training=is_training)
net = ProposalAttention(config, is_training=is_training)
net.train(is_training)
ROIdelegator = RoIPooling(pooled_height=1,
pooled_width=1,
spatial_scale=1.0 / 32)
# Forward hook
layer = net.targeted_layer()
features = list()
def hook_feature(module, input, output):
features.append(output[0])
features.append(output[1])
features.append(output[2])
layer.register_forward_hook(hook_feature)
ftr_strm.write('%d\n' % (net.backbone.layers_out_filters[-1]))
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
# Restore pretrain model
if config["pretrain_snapshot"]:
logging.info("load checkpoint from {}".format(
config["pretrain_snapshot"]))
state_dict = torch.load(config["pretrain_snapshot"])
net.load_state_dict(state_dict)
else:
raise Exception("missing pretrain_snapshot!!!")
# YOLO loss with 3 scales
yolo_losses = []
for i in range(3):
yolo_losses.append(
ProposalLoss(config["yolo"]["anchors"][i],
(config["img_w"], config["img_h"])))
# Prepare images path
images_path = [
y for x in os.walk(config["ref_path"])
for y in glob.glob(os.path.join(x[0], "*.jpg"))
]
images_path.sort()
print 'Num. of images: ', len(images_path)
if len(images_path) == 0:
raise Exception("no image found in {}".format(config["ref_path"]))
# Start inference
if config["pca"]:
pca = pickle.load(open(config["mapping_fn"], 'rb'))
batch_size = config["batch_size"]
for step in range(0, len(images_path), batch_size):
# preprocess
images = []
images_origin = []
features = []
name = ""
for path in images_path[step * batch_size:(step + 1) * batch_size]:
logging.info("processing: {}".format(path))
image = cv2.imread(path, cv2.IMREAD_COLOR)
name = path.split('/')[-1][:-4]
if image is None:
logging.error("read path error: {}. skip it.".format(path))
continue
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
images_origin.append(image) # keep for save result
image = cv2.resize(image, (config["img_w"], config["img_h"]),
interpolation=cv2.INTER_LINEAR)
image = image.astype(np.float32)
image /= 255.0
image = np.transpose(image, (2, 0, 1))
image = image.astype(np.float32)
images.append(image)
images = np.asarray(images)
images = torch.from_numpy(images).cuda()
# inference
with torch.no_grad():
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(yolo_losses[i](outputs[i]))
output = torch.cat(output_list, 1)
batch_detections = non_max_suppression(output)
# ROI Pooling
if batch_detections[0] is None:
continue
rois = batch_detections[0].data.cpu().numpy()[:, :-1]
image_index = np.zeros((rois.shape[0], 1))
rois = np.rint(np.concatenate((image_index, rois), axis=1))
rois = torch.Tensor(rois).cuda()
feature_map = features[2]
output = ROIdelegator(feature_map, rois)
output = output.view(output.size(0),
output.size(1) * output.size(2) * output.size(3))
output = output.data.cpu().numpy()
# Output
cnt = output.shape[0]
assert cnt == rois.shape[0]
dim = output.shape[1]
if config["pca"]:
output = pca.transform(output)
output = normalize(output, norm='l2', axis=1)
ori_h, ori_w = images_origin[0].shape[:2]
pre_h, pre_w = config["img_h"], config["img_w"]
for i in xrange(cnt):
x1 = torch.clamp(rois[i][1].data, min=0, max=pre_w) / pre_w * ori_w
y1 = torch.clamp(rois[i][2].data, min=0, max=pre_h) / pre_h * ori_h
x2 = torch.clamp(rois[i][3].data, min=0, max=pre_w) / pre_w * ori_w
y2 = torch.clamp(rois[i][4].data, min=0, max=pre_h) / pre_h * ori_h
img_strm.write('%s-%04d %d %d %d %d\n' %
(name, i + 1, x1, y1, x2, y2))
for j in xrange(dim):
if j < dim - 1:
ftr_strm.write('%f ' % output[i][j])
else:
ftr_strm.write('%f\n' % output[i][j])
"""
def pca(config):
"""
:param config:
:return:
"""
# Load and initialize network
is_training = False
# net = ProposalModel(config, is_training=is_training)
net = ProposalAttention(config, is_training=is_training)
net.train(is_training)
ROIdelegator = RoIPooling(pooled_height=1,
pooled_width=1,
spatial_scale=1.0 / 32)
# Forward hook
layer = net.targeted_layer()
features = list()
def hook_feature(module, input, output):
features.append(output[0])
features.append(output[1])
features.append(output[2])
layer.register_forward_hook(hook_feature)
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
# Restore pretrain model
if config["pretrain_snapshot"]:
logging.info("load checkpoint from {}".format(
config["pretrain_snapshot"]))
state_dict = torch.load(config["pretrain_snapshot"])
net.load_state_dict(state_dict)
else:
raise Exception("missing pretrain_snapshot!!!")
# YOLO loss with 3 scales
yolo_losses = []
for i in range(3):
yolo_losses.append(
ProposalLoss(config["yolo"]["anchors"][i],
(config["img_w"], config["img_h"])))
# Prepare images path
images_path = [
y for x in os.walk(config["pca_path"])
for y in glob.glob(os.path.join(x[0], "*.jpg"))
]
images_path.sort()
print 'Num. of images: ', len(images_path)
if len(images_path) == 0:
raise Exception("no image found in {}".format(config["ref_path"]))
# Start inference
batch_size = config["batch_size"]
feats = list()
for step in range(0, len(images_path), batch_size):
# preprocess
images = []
images_origin = []
features = []
name = ""
for path in images_path[step * batch_size:(step + 1) * batch_size]:
logging.info("processing: {}".format(path))
image = cv2.imread(path, cv2.IMREAD_COLOR)
name = path.split('/')[-1][:-4]
if image is None:
logging.error("read path error: {}. skip it.".format(path))
continue
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
images_origin.append(image) # keep for save result
image = cv2.resize(image, (config["img_w"], config["img_h"]),
interpolation=cv2.INTER_LINEAR)
image = image.astype(np.float32)
image /= 255.0
image = np.transpose(image, (2, 0, 1))
image = image.astype(np.float32)
images.append(image)
images = np.asarray(images)
images = torch.from_numpy(images).cuda()
# inference
with torch.no_grad():
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(yolo_losses[i](outputs[i]))
output = torch.cat(output_list, 1)
batch_detections = non_max_suppression(output)
# ROI Pooling
if batch_detections[0] is None:
continue
rois = batch_detections[0].data.cpu().numpy()[:, :-1]
image_index = np.zeros((rois.shape[0], 1))
rois = np.rint(np.concatenate((image_index, rois), axis=1))
rois = torch.Tensor(rois).cuda()
feature_map = features[2]
output = ROIdelegator(feature_map, rois)
output = output.view(output.size(0),
output.size(1) * output.size(2) * output.size(3))
output = output.data.cpu().numpy()
# output = normalize(output, norm='l2', axis=1)
# print output.shape
# Output
feats.append(output)
feats = np.concatenate(tuple(feats), axis=0)
print 'Learning pca with dim:', feats.shape
# PCA
feats = normalize(feats, norm='l2', axis=1)
pca = PCA(feats.shape[1], whiten=True)
pca.fit(feats)
pickle.dump(pca, open(config["mapping_fn"], 'wb'))
def test(config):
is_training = False
anchors = [int(x) for x in config["yolo"]["anchors"].split(",")]
anchors = [[[anchors[i], anchors[i + 1]], [anchors[i + 2], anchors[i + 3]], [anchors[i + 4], anchors[i + 5]]] for i
in range(0, len(anchors), 6)]
anchors.reverse()
config["yolo"]["anchors"] = []
for i in range(3):
config["yolo"]["anchors"].append(anchors[i])
# Load and initialize network
net = ModelMain(config, is_training=is_training)
net.train(is_training)
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
# Restore pretrain model
if config["pretrain_snapshot"]:
logging.info("load checkpoint from {}".format(config["pretrain_snapshot"]))
state_dict = torch.load(config["pretrain_snapshot"])
net.load_state_dict(state_dict)
else:
raise Exception("missing pretrain_snapshot!!!")
# YOLO loss with 3 scales
yolo_losses = []
for i in range(3):
yolo_losses.append(YOLOLayer(config["batch_size"],i,config["yolo"]["anchors"][i],
config["yolo"]["classes"], (config["img_w"], config["img_h"])))
# prepare images path
images_name = os.listdir(config["images_path"])
images_path = [os.path.join(config["images_path"], name) for name in images_name]
if len(images_path) == 0:
raise Exception("no image found in {}".format(config["images_path"]))
cap = cv2.VideoCapture(0)
# cap = cv2.VideoCapture("./007.avi")
img_i = 0
start = time.time()
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
img_i += 1
# preprocess
images = []
images_origin = []
image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
images_origin.append(image) # keep for save result
image = cv2.resize(image, (config["img_w"], config["img_h"]),
interpolation=cv2.INTER_LINEAR)
image = image.astype(np.float32)
image /= 255.0
image = np.transpose(image, (2, 0, 1))
image = image.astype(np.float32)
images.append(image)
images = np.asarray(images)
images = torch.from_numpy(images).cuda()
# inference
with torch.no_grad():
time1=datetime.datetime.now()
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(yolo_losses[i](outputs[i]))
output = torch.cat(output_list, 1)
print("time1",(datetime.datetime.now()-time1).microseconds)
batch_detections = non_max_suppression(output, config["yolo"]["classes"],
conf_thres=config["confidence_threshold"])
print("time2", (datetime.datetime.now() - time1).microseconds)
# write result images. Draw bounding boxes and labels of detections
classes = open(config["classes_names_path"], "r").read().split("\n")[:-1]
if not os.path.isdir("./output/"):
os.makedirs("./output/")
for idx, detections in enumerate(batch_detections):
img_show = images_origin[idx]
img_show = cv2.cvtColor(img_show, cv2.COLOR_RGB2BGR)
if detections is not None:
unique_labels = detections[:, -1].cpu().unique()
n_cls_preds = len(unique_labels)
boxes=[]
for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
# Rescale coordinates to original dimensions
ori_h, ori_w = images_origin[idx].shape[:2]
pre_h, pre_w = config["img_h"], config["img_w"]
box_h = ((y2 - y1) / pre_h) * ori_h
box_w = ((x2 - x1) / pre_w) * ori_w
y1 = (y1 / pre_h) * ori_h
x1 = (x1 / pre_w) * ori_w
# Create a Rectangle patch
box = BoundBox(x1, y1, x1 + box_w, y1 + box_h, cls_conf.item(), int(cls_pred))
boxes.append(box)
img_show = draw_boxes(img_show, boxes, labels)
# image_show = cv2.rectangle(images_origin[idx], (x1, y1), (x1 + box_w, y1 + box_h), (0, 255, 0), 1)
cv2.imshow('1', img_show)
cv2.waitKey(1)
logging.info("Save all results to ./output/")
def detect_image(name, response, config, net, yolo_losses, classes, complex_yolo_416, pid):
start_time = time.time()
images_path = [os.path.join(config["images_path"], name)]
if len(images_path) == 0:
raise Exception("no image with name {} found in {}".format(name, config["images_path"]))
# Start inference
batch_size = config["batch_size"]
for step in range(0, len(images_path), batch_size):
images = []
images_origin = []
for path in images_path[step * batch_size: (step + 1) * batch_size]:
logging.info("processing: {}".format(path))
image = cv2.imread(path, cv2.IMREAD_COLOR)
if image is None:
logging.error("read path error: {}. skip it.".format(path))
continue
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
images_origin.append(image) # keep for save result
image = cv2.resize(image, (config["img_w"], config["img_h"]),
interpolation=cv2.INTER_LINEAR)
image = image.astype(np.float32)
image /= 255.0
image = np.transpose(image, (2, 0, 1))
image = image.astype(np.float32)
images.append(image)
images = np.asarray(images)
images = torch.from_numpy(images)
lock.acquire()
# inference
with torch.no_grad():
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(yolo_losses[i](outputs[i]))
output = torch.cat(output_list, 1)
batch_detections = non_max_suppression(output, config["yolo"]["classes"],
conf_thres=config["confidence_threshold"],
nms_thres=0.45)
lock.release()
# write result images. Draw bounding boxes and labels of detections
if not os.path.isdir("./output/"):
os.makedirs("./output/")
for idx, detections in enumerate(batch_detections):
plt.figure()
fig, ax = plt.subplots(1)
ax.imshow(images_origin[idx])
if detections is not None:
unique_labels = detections[:, -1].cpu().unique()
n_cls_preds = len(unique_labels)
bbox_colors = random.sample(colors, n_cls_preds)
for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
color = bbox_colors[int(np.where(unique_labels == int(cls_pred))[0])]
# Rescale coordinates to original dimensions
ori_h, ori_w = images_origin[idx].shape[:2]
pre_h, pre_w = config["img_h"], config["img_w"]
box_h = ((y2 - y1) / pre_h) * ori_h
box_w = ((x2 - x1) / pre_w) * ori_w
y1 = (y1 / pre_h) * ori_h
x1 = (x1 / pre_w) * ori_w
# Create a Rectangle patch
bbox = patches.Rectangle((x1, y1), box_w, box_h, linewidth=2,
edgecolor=color,
facecolor='none')
# Add the bbox to the plot
ax.add_patch(bbox)
# Add label
plt.text(x1, y1, s=classes[int(cls_pred)], color='white',
verticalalignment='top',
bbox={'color': color, 'pad': 0})
# Save generated image with detections
plt.axis('off')
plt.gca().xaxis.set_major_locator(NullLocator())
plt.gca().yaxis.set_major_locator(NullLocator())
plt.savefig('output/{}'.format(name), bbox_inches='tight', pad_inches=0.0)
plt.close()
computation_time = time.time() - start_time
cpu = psutil.cpu_percent() / 100
complex_yolo_416.append(computation_time * cpu * 2.8)
print("\tyolo3 + {} finished in {}s, system response in {} s, cpu in {} cycles (10^9)"
.format(round(cpu, 3), round(computation_time, 4)
, round(np.average(response), 4)
, round(np.average(complex_yolo_416), 3)))
# logging.info("Save all results to ./output/")
return computation_time, cpu * 100, complex_yolo_416
def eval_voc(self, val_dataset):
logging.info('Start Evaling')
def voc_ap(rec, prec, use_07_metric=False):
""" ap = voc_ap(rec, prec, [use_07_metric])
Compute VOC AP given precision and recall.
If use_07_metric is true, uses the
VOC 07 11 point method (default:False).
"""
_rec = np.arange(0., 1.1, 0.1)
_prec = []
if use_07_metric:
# 11 point metric
ap = 0.
for t in np.arange(0., 1.1, 0.1):
if np.sum(rec >= t) == 0:
p = 0
else:
p = np.max(prec[rec >= t])
_prec.append(p)
ap = ap + p / 11.
else:
# correct AP calculation
# first append sentinel values at the end
mrec = np.concatenate(([0.], rec, [1.]))
mpre = np.concatenate(([0.], prec, [0.]))
# compute the precision envelope
for i in range(mpre.size - 1, 0, -1):
mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])
# to calculate area under PR curve, look for points
# where X axis (recall) changes value
i = np.where(mrec[1:] != mrec[:-1])[0]
# and sum (\Delta recall) * prec
ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])
return ap
def parse_rec(imagename):
filename = imagename.replace('jpg', 'xml')
tree = ET.parse(filename)
objects = []
for obj in tree.findall('object'):
obj_struct = {}
obj_struct['pose'] = obj.fine('pose').text()
obj_struct['truncated'] = int(obj.find('truncated').text)
obj_struct['difficult'] = int(obj.find('difficult').text)
bbox = obj.find('bndbox')
obj_struct['bbox'] = [
int(bbox.find('xmin').text),
int(bbox.find('ymin').text),
int(bbox.find('xmax').text),
int(bbox.find('ymax').text)
]
objects.append(obj_struct)
return objects
for step, samples in enumerate(val_dataset):
images, labels = samples['image'], samples['label']
image_paths, origin_sizes = samples['image_path'], samples[
'origin_size']
with torch.no_grad():
outputs = self.net(images)
output_list = []
for i in range(3):
output_list.append(self.yolo_loss[i](outputs[i], labels))
output = torch.cat(output_list, 1)
batch_detections = non_max_suppression(output,
self.config.num_classes,
conf_thres=0.01,
nms_thres=0.4)
for idx, detections in enumerate(batch_detections):
image_id = int(os.path.basename(image_paths[idx])[:6])
if detections is not None:
origin_size = eval(origin_sizes[idx])
detections = detections.cpu().numpy()
for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
x1 = x1 / self.config.image_size * origin_size[0]
x2 = x2 / self.config.image_size * origin_size[0]
y1 = y1 / self.config.image_size * origin_size[1]
y2 = y2 / self.config.image_size * origin_size[1]
def test(config):
"""
:param config:
:return:
"""
is_training = False
# Load and initialize network
# net = ProposalModel(config, is_training=is_training)
net = ProposalAttention(config, is_training=is_training)
net.train(is_training)
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
# Restore pretrain model
if config["pretrain_snapshot"]:
logging.info("load checkpoint from {}".format(
config["pretrain_snapshot"]))
state_dict = torch.load(config["pretrain_snapshot"])
net.load_state_dict(state_dict)
else:
raise Exception("missing pretrain_snapshot!!!")
# YOLO loss with 3 scales
yolo_losses = []
for i in range(3):
yolo_losses.append(
ProposalLoss(config["yolo"]["anchors"][i],
(config["img_w"], config["img_h"])))
# prepare images path
images_name = os.listdir(config["images_path"])
# shuffle(images_name)
# images_name = images_name[:40]
images_name.sort()
images_path = [
os.path.join(config["images_path"], name) for name in images_name
]
if len(images_path) == 0:
raise Exception("no image found in {}".format(config["images_path"]))
# Start inference
batch_size = config["batch_size"]
for step in range(0, len(images_path), batch_size):
# preprocess
images = []
images_origin = []
for path in images_path[step * batch_size:(step + 1) * batch_size]:
logging.info("processing: {}".format(path))
name = path.split('/')[-1][:-4]
image = cv2.imread(path, cv2.IMREAD_COLOR)
if image is None:
logging.error("read path error: {}. skip it.".format(path))
continue
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
images_origin.append(image) # keep for save result
image = cv2.resize(image, (config["img_w"], config["img_h"]),
interpolation=cv2.INTER_LINEAR)
image = image.astype(np.float32)
image /= 255.0
image = np.transpose(image, (2, 0, 1))
image = image.astype(np.float32)
images.append(image)
images = np.asarray(images)
images = torch.from_numpy(images).cuda()
# inference
with torch.no_grad():
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(yolo_losses[i](outputs[i]))
output = torch.cat(output_list, 1)
# print output.size()
batch_detections = non_max_suppression(output)
# write result images. Draw bounding boxes and labels of detections
if not os.path.isdir("./output/"):
os.makedirs("./output/")
for idx, detections in enumerate(batch_detections):
plt.figure()
fig, ax = plt.subplots(1)
ax.imshow(images_origin[idx])
if detections is not None:
for x1, y1, x2, y2, conf in detections:
bbox_colors = random.sample(colors, len(batch_detections))
color = bbox_colors[idx]
# Rescale coordinates to original dimensions
ori_h, ori_w = images_origin[idx].shape[:2]
pre_h, pre_w = config["img_h"], config["img_w"]
if x1 < 0:
x1 = 1
if x2 > pre_w:
x2 = pre_w - 1
if y1 < 0:
y1 = 1
if y2 > pre_h:
y2 = pre_h - 1
box_h = ((y2 - y1) / pre_h) * ori_h
box_w = ((x2 - x1) / pre_w) * ori_w
y1 = (y1 / pre_h) * ori_h
x1 = (x1 / pre_w) * ori_w
# Create a Rectangle patch
bbox = patches.Rectangle((x1, y1),
box_w,
box_h,
linewidth=2,
edgecolor=color,
facecolor='none')
# Add the bbox to the plot
ax.add_patch(bbox)
ax.text(x1,
y1,
'{:.2f}'.format(conf),
bbox=dict(facecolor=color, alpha=0.9),
fontsize=8,
color='white')
else:
print 'Nothing detected.'
# Save generated image with detections
plt.axis('off')
plt.gca().xaxis.set_major_locator(NullLocator())
plt.gca().yaxis.set_major_locator(NullLocator())
plt.savefig('output/{}-attn.png'.format(name),
bbox_inches='tight',
pad_inches=0.0)
plt.close()
logging.info("Save all results to ./output/")
def evaluate(config):
is_training = False
# Load and initialize network
net = ModelMain(config, is_training=is_training)
net.train(is_training)
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
# Restore pretrain model
if config["pretrain_snapshot"]:
state_dict = torch.load(config["pretrain_snapshot"])
net.load_state_dict(state_dict)
else:
logging.warning("missing pretrain_snapshot!!!")
# YOLO loss with 3 scales
yolo_losses = []
for i in range(3):
yolo_losses.append(
YOLOLoss(config["yolo"]["anchors"][i], config["yolo"]["classes"],
(config["img_w"], config["img_h"])))
# DataLoader
dataloader = torch.utils.data.DataLoader(COCODataset(
config["val_path"], (config["img_w"], config["img_h"]),
is_training=False),
batch_size=config["batch_size"],
shuffle=False,
num_workers=16,
pin_memory=False)
# Start the eval loop
logging.info("Start eval.")
n_gt = 0
correct = 0
for step, samples in enumerate(dataloader):
images, labels = samples["image"], samples["label"]
labels = labels.cuda()
with torch.no_grad():
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(yolo_losses[i](outputs[i]))
output = torch.cat(output_list, 1)
output = non_max_suppression(output, 80, conf_thres=0.2)
# calculate
for sample_i in range(labels.size(0)):
# Get labels for sample where width is not zero (dummies)
target_sample = labels[sample_i, labels[sample_i, :, 3] != 0]
for obj_cls, tx, ty, tw, th in target_sample:
# Get rescaled gt coordinates
tx1, tx2 = config["img_w"] * (
tx - tw / 2), config["img_w"] * (tx + tw / 2)
ty1, ty2 = config["img_h"] * (
ty - th / 2), config["img_h"] * (ty + th / 2)
n_gt += 1
box_gt = torch.cat([
coord.unsqueeze(0) for coord in [tx1, ty1, tx2, ty2]
]).view(1, -1)
sample_pred = output[sample_i]
if sample_pred is not None:
# Iterate through predictions where the class predicted is same as gt
for x1, y1, x2, y2, conf, obj_conf, obj_pred in sample_pred[
sample_pred[:, 6] == obj_cls]:
box_pred = torch.cat([
coord.unsqueeze(0)
for coord in [x1, y1, x2, y2]
]).view(1, -1)
iou = bbox_iou(box_pred, box_gt)
if iou >= config["iou_thres"]:
correct += 1
break
if n_gt:
logging.info('Batch [%d/%d] mAP: %.5f' %
(step, len(dataloader), float(correct / n_gt)))
logging.info('Mean Average Precision: %.5f' % float(correct / n_gt))
def evaluate(config):
is_training = False
# Load and initialize network
net = ModelMain(config, is_training=is_training)
net.train(is_training)
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
# Restore pretrain model
if config["pretrain_snapshot"]:
logging.info("Load checkpoint: {}".format(config["pretrain_snapshot"]))
state_dict = torch.load(config["pretrain_snapshot"])
net.load_state_dict(state_dict)
else:
logging.warning("missing pretrain_snapshot!!!")
# YOLO loss with 3 scales
yolo_losses = []
for i in range(3):
yolo_losses.append(
YOLOLoss(config["yolo"]["anchors"][i], config["yolo"]["classes"],
(config["img_w"], config["img_h"])))
# DataLoader.
dataloader = torch.utils.data.DataLoader(COCODataset(
config["val_path"], (config["img_w"], config["img_h"]),
is_training=False),
batch_size=config["batch_size"],
shuffle=False,
num_workers=8,
pin_memory=False)
# Coco Prepare.
index2category = json.load(open("coco_index2category.json"))
# Start the eval loop
logging.info("Start eval.")
coco_results = []
coco_img_ids = set([])
APs = []
for step, samples in enumerate(dataloader):
images, labels = samples["image"], samples["label"]
image_paths, origin_sizes = samples["image_path"], samples[
"origin_size"]
with torch.no_grad():
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(yolo_losses[i](outputs[i]))
output = torch.cat(output_list, 1)
batch_detections = non_max_suppression(output,
config["yolo"]["classes"],
conf_thres=0.0001,
nms_thres=0.45)
for idx, detections in enumerate(batch_detections):
correct = []
annotations = labels[idx, labels[idx, :, 3] != 0]
image_id = int(os.path.basename(image_paths[idx])[-16:-4])
coco_img_ids.add(image_id)
if detections is None:
if annotations.size(0) != 0:
APs.append(0)
continue
detections = detections[np.argsort(-detections[:, 4])]
origin_size = eval(origin_sizes[idx])
detections = detections.cpu().numpy()
# ===========================================================================================================================
# The amount of padding that was added
pad_x = max(origin_size[1] - origin_size[0],
0) * (config["img_w"] / max(origin_size))
pad_y = max(origin_size[0] - origin_size[1],
0) * (config["img_w"] / max(origin_size))
# Image height and width after padding is removed
unpad_h = config["img_w"] - pad_y
unpad_w = config["img_w"] - pad_x
# ===========================================================================================================================
if annotations.size(0) == 0:
correct.extend([0 for _ in range(len(detections))])
else:
target_boxes = torch.FloatTensor(annotations[:, 1:].shape)
target_boxes[:,
0] = (annotations[:, 1] - annotations[:, 3] / 2)
target_boxes[:,
1] = (annotations[:, 2] - annotations[:, 4] / 2)
target_boxes[:,
2] = (annotations[:, 1] + annotations[:, 3] / 2)
target_boxes[:,
3] = (annotations[:, 2] + annotations[:, 4] / 2)
target_boxes *= config["img_w"]
detected = []
for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
pred_bbox = (x1, y1, x2, y2)
#x1 = x1 / config["img_w"] * origin_size[0]
#x2 = x2 / config["img_w"] * origin_size[0]
#y1 = y1 / config["img_h"] * origin_size[1]
#y2 = y2 / config["img_h"] * origin_size[1]
#w = x2 - x1
#h = y2 - y1
h = ((y2 - y1) / unpad_h) * origin_size[1]
w = ((x2 - x1) / unpad_w) * origin_size[0]
y1 = ((y1 - pad_y // 2) / unpad_h) * origin_size[1]
x1 = ((x1 - pad_x // 2) / unpad_w) * origin_size[0]
coco_results.append({
"image_id":
image_id,
"category_id":
index2category[str(int(cls_pred.item()))],
"bbox": (float(x1), float(y1), float(w), float(h)),
"score":
float(conf),
})
pred_bbox = torch.FloatTensor(pred_bbox).view(1, -1)
# Compute iou with target boxes
iou = bbox_iou(pred_bbox, target_boxes)
# Extract index of largest overlap
best_i = np.argmax(iou)
# If overlap exceeds threshold and classification is correct mark as correct
if iou[best_i] > config[
'iou_thres'] and cls_pred == annotations[
best_i, 0] and best_i not in detected:
correct.append(1)
detected.append(best_i)
else:
correct.append(0)
true_positives = np.array(correct)
false_positives = 1 - true_positives
# Compute cumulative false positives and true positives
false_positives = np.cumsum(false_positives)
true_positives = np.cumsum(true_positives)
# Compute recall and precision at all ranks
recall = true_positives / annotations.size(0) if annotations.size(
0) else true_positives
precision = true_positives / np.maximum(
true_positives + false_positives,
np.finfo(np.float64).eps)
# Compute average precision
AP = compute_ap(recall, precision)
APs.append(AP)
print("+ Sample [%d/%d] AP: %.4f (%.4f)" %
(len(APs), 5000, AP, np.mean(APs)))
logging.info("Now {}/{}".format(step, len(dataloader)))
print("Mean Average Precision: %.4f" % np.mean(APs))
save_results_path = "coco_results.json"
with open(save_results_path, "w") as f:
json.dump(coco_results,
f,
sort_keys=True,
indent=4,
separators=(',', ':'))
logging.info("Save coco format results to {}".format(save_results_path))
# COCO api
logging.info("Using coco-evaluate tools to evaluate.")
cocoGt = COCO(config["annotation_path"])
cocoDt = cocoGt.loadRes(save_results_path)
cocoEval = COCOeval(cocoGt, cocoDt, "bbox")
cocoEval.params.imgIds = list(coco_img_ids) # real imgIds
cocoEval.evaluate()
cocoEval.accumulate()
cocoEval.summarize()
def evaluate(config):
# checkpoint_paths = {'58': r'\\192.168.25.58\Team-CV\checkpoints\torch_yolov3'}
checkpoint_paths = {'39': r'F:\Team-CV\checkpoints\shuffle_v2/'}
# checkpoint_paths = {'68': r'E:\github\YOLOv3_PyTorch\evaluate\weights'}
post_weights = {k: 0 for k in checkpoint_paths.keys()}
weight_index = {k: 0 for k in checkpoint_paths.keys()}
time_inter = 10
dataloader = torch.utils.data.DataLoader(COCODataset(
config["train_path"], (config["img_w"], config["img_h"]),
is_training=False,
is_scene=True),
batch_size=config["batch_size"],
shuffle=False,
num_workers=0,
pin_memory=False,
drop_last=True) # DataLoader
net, yolo_losses = build_yolov3(config)
while 1:
for key, checkpoint_path in checkpoint_paths.items():
os.makedirs(checkpoint_path + '/result', exist_ok=True)
checkpoint_weights = os.listdir(checkpoint_path)
checkpoint_result = os.listdir(checkpoint_path + '/result')
checkpoint_result = [
cweight.split("_")[2][:-4] for cweight in checkpoint_result
if cweight.endswith('ini')
]
checkpoint_weights = [
cweight for cweight in checkpoint_weights
if cweight.endswith('weights')
]
if weight_index[key] >= len(checkpoint_weights):
print('weight_index[key]', weight_index[key],
len(checkpoint_weights))
time.sleep(time_inter)
continue
if post_weights[key] == checkpoint_weights[weight_index[key]]:
print('post_weights[key]', post_weights[key])
time.sleep(time_inter)
continue
post_weights[key] = checkpoint_weights[weight_index[key]]
if post_weights[key].endswith("_.weights"): #检查权重是否保存完
print("post_weights[key].split('_')",
post_weights[key].split('_'))
time.sleep(time_inter)
continue
if checkpoint_weights[weight_index[key]].split(
"_")[1][:-8] in checkpoint_result:
print('weight_index[key] +', weight_index[key])
weight_index[key] += 1
time.sleep(time_inter // 20)
continue
weight_index[key] += 1
try:
if config["pretrain_snapshot"]: # Restore pretrain model
state_dict = torch.load(config["pretrain_snapshot"])
logging.info("loading model from %s" %
config["pretrain_snapshot"])
net.load_state_dict(state_dict)
else:
state_dict = torch.load(
os.path.join(checkpoint_path, post_weights[key]))
logging.info(
"loading model from %s" %
os.path.join(checkpoint_path, post_weights[key]))
net.load_state_dict(state_dict)
except Exception as E:
print(E)
time.sleep(time_inter)
continue
logging.info("Start eval.") # Start the eval loop
n_gt = 0
correct = 0
imagepath_list = []
for step, samples in enumerate(dataloader):
images, labels, image_paths = samples["image"], samples[
"label"], samples["img_path"]
labels = labels.cuda()
with torch.no_grad():
time1 = datetime.datetime.now()
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(yolo_losses[i](outputs[i]))
output = torch.cat(output_list, 1)
output = non_max_suppression(output, 1, conf_thres=0.5)
if ((datetime.datetime.now() - time1).seconds > 5):
logging.info('Batch %d time is too long ' % (step))
n_gt = 1
break
print(
"time2",
(datetime.datetime.now() - time1).seconds * 1000 +
(datetime.datetime.now() - time1).microseconds // 1000)
# calculate
for sample_i in range(labels.size(0)):
# Get labels for sample where width is not zero (dummies)
target_sample = labels[sample_i,
labels[sample_i, :, 3] != 0]
for obj_cls, tx, ty, tw, th in target_sample:
# Get rescaled gt coordinates
tx1, tx2 = config["img_w"] * (
tx - tw / 2), config["img_w"] * (tx + tw / 2)
ty1, ty2 = config["img_h"] * (
ty - th / 2), config["img_h"] * (ty + th / 2)
n_gt += 1
box_gt = torch.cat([
coord.unsqueeze(0)
for coord in [tx1, ty1, tx2, ty2]
]).view(1, -1)
sample_pred = output[sample_i]
if sample_pred is not None:
# Iterate through predictions where the class predicted is same as gt
for x1, y1, x2, y2, conf, obj_conf, obj_pred in sample_pred[
sample_pred[:, 6] == obj_cls.cuda()]:
box_pred = torch.cat([
coord.unsqueeze(0)
for coord in [x1, y1, x2, y2]
]).view(1, -1)
iou = bbox_iou(box_pred, box_gt)
if iou >= config["iou_thres"]:
correct += 1
break
else:
if image_paths[
sample_i] not in imagepath_list:
imagepath_list.append(
image_paths[sample_i])
else:
if image_paths[sample_i] not in imagepath_list:
imagepath_list.append(
image_paths[sample_i])
if n_gt:
logging.info('Batch [%d/%d] err_count:%d mAP: %.5f' %
(step, len(dataloader), len(imagepath_list),
float(correct / n_gt)))
logging.info('Mean Average Precision: %.5f' %
float(correct / n_gt))
Mean_Average = float(correct / n_gt)
ini_name = os.path.join(
checkpoint_path + '/result/',
'%.4f_%s.ini' % ((float(post_weights[key].split("_")[0]) +
float(correct / n_gt)) / 2,
post_weights[key].replace(".weights", "")))
write_ini(ini_name, Mean_Average, imagepath_list)
break
def test(config):
is_training = False
anchors = [int(x) for x in config["yolo"]["anchors"].split(",")]
anchors = [[[anchors[i], anchors[i + 1]], [anchors[i + 2], anchors[i + 3]],
[anchors[i + 4], anchors[i + 5]]]
for i in range(0, len(anchors), 6)]
anchors.reverse()
config["yolo"]["anchors"] = []
for i in range(3):
config["yolo"]["anchors"].append(anchors[i])
# Load and initialize network
net = ModelMain(config, is_training=is_training)
net.train(is_training)
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
# Restore pretrain model
if config["pretrain_snapshot"]:
logging.info("load checkpoint from {}".format(
config["pretrain_snapshot"]))
state_dict = torch.load(config["pretrain_snapshot"])
net.load_state_dict(state_dict)
else:
raise Exception("missing pretrain_snapshot!!!")
# YOLO loss with 3 scales
yolo_losses = []
for i in range(3):
yolo_losses.append(
YOLOLayer(config["batch_size"], i, config["yolo"]["anchors"][i],
config["yolo"]["classes"],
(config["img_w"], config["img_h"])))
# prepare images path
images_name = os.listdir(config["images_path"])
images_path = [
os.path.join(config["images_path"], name) for name in images_name
]
if len(images_path) == 0:
raise Exception("no image found in {}".format(config["images_path"]))
# Start inference
batch_size = config["batch_size"]
for step in range(0, len(images_path), batch_size):
# preprocess
images = []
images_origin = []
for path in images_path[step * batch_size:(step + 1) * batch_size]:
logging.info("processing: {}".format(path))
image = cv2.imread(path, cv2.IMREAD_COLOR)
if image is None:
logging.error("read path error: {}. skip it.".format(path))
continue
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
images_origin.append(image) # keep for save result
image = cv2.resize(image, (config["img_w"], config["img_h"]),
interpolation=cv2.INTER_LINEAR)
image = image.astype(np.float32)
image /= 255.0
image = np.transpose(image, (2, 0, 1))
image = image.astype(np.float32)
images.append(image)
images = np.asarray(images)
images = torch.from_numpy(images).cuda()
# inference
with torch.no_grad():
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(yolo_losses[i](outputs[i]))
output = torch.cat(output_list, 1)
batch_detections = non_max_suppression(
output,
config["yolo"]["classes"],
conf_thres=config["confidence_threshold"])
# write result images. Draw bounding boxes and labels of detections
classes = open(config["classes_names_path"],
"r").read().split("\n")[:-1]
for idx, detections in enumerate(batch_detections):
if detections is not None:
unique_labels = detections[:, -1].cpu().unique()
n_cls_preds = len(unique_labels)
boxes = []
for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
# Rescale coordinates to original dimensions
ori_h, ori_w = images_origin[idx].shape[:2]
pre_h, pre_w = config["img_h"], config["img_w"]
box_h = ((y2 - y1) / pre_h) * ori_h
box_w = ((x2 - x1) / pre_w) * ori_w
y1 = (y1 / pre_h) * ori_h
x1 = (x1 / pre_w) * ori_w
# Create a Rectangle patch
box = BoundBox(x1, y1, x1 + box_w, y1 + box_h,
cls_conf.item(), int(cls_pred),
classes[int(cls_pred)])
boxes.append(box)
# Save generated image with detections
img_show = draw_boxes(images_origin[idx], boxes, labels, 0.5)
img_show = cv2.resize(img_show,
(img_show.shape[1], img_show.shape[0]),
interpolation=cv2.INTER_CUBIC)
# outVideo.write(img_show)
cv2.imshow("ai", img_show)
cv2.waitKey()
logging.info("Save all results to ./output/")
def test(model, conf_thres=0.001, nms_thres=0.5):
device = "cuda:0"
classes = ['person']
if model is None:
pass
else:
device = model.device # get model device
data_parameters = read_yaml(model.data_yaml)
classes = data_parameters['classes']
batch_size = model.batch_size
# Dataset
data_parameters[
"data_path"] = '/home/lingc1/data/sports-training-data/player_detection/validate_dataset_5k_half_size'
data_set = Yolov3Data(data_parameters, None, index_file='val_test')
dataloader = DataLoader(data_set,
batch_size,
shuffle=False,
num_workers=0,
collate_fn=data_set.collate_fn)
print(('%20s' + '%10s' * 6) %
('Class', 'Images', 'Targets', 'P', 'R', 'mAP', 'F1'))
loss, p_80, r_80, f1_80, mp_80, mr_80, map_80, mf1_80 = 0., 0., 0., 0., 0., 0., 0., 0.
seen = 0
images_num = 0
output_results = ""
class_recs = {}
nc = len(classes)
names = classes
det_lines = []
imagenames = []
npos_cls = {}
for i in range(nc):
npos_cls[i] = 0
class_recs[i] = {}
for i, (imgs, targets) in enumerate(dataloader):
imgs = imgs.to(device)
targets = targets.to(device)
target_number = len(targets)
preds = model.inference(imgs, None)
output = non_max_suppression(preds,
conf_thres=conf_thres,
nms_thres=nms_thres)
true_targets = targets[torch.sum(targets[:, 1:6], 1) != 0]
# Statistics per image
# remove the targets that fills 0 for data distribution.
true_targets = targets[torch.sum(targets[:, 1:6], 1) != 0]
# npos += len(true_targets)
for si, pred in enumerate(output):
images_num += 1
if pred is not None and len(pred) > 0:
# Rescale boxes from 416 to true image size
# pred[:, :4] = scale_coords(imgs.shape[2:], pred[:, :4], im0_shape).round()
for *xyxy, conf, cls_conf, cls in pred:
x = int(xyxy[0])
y = int(xyxy[1])
w = int((xyxy[2] - xyxy[0]).round())
h = int((xyxy[3] - xyxy[1]).round())
output_line = "{:s},{:d},{:d},{:d},{:d},{:f}\n".format(
Path(str(si)).name, x, y, w, h, conf)
det_lines.append(output_line)
output_results = output_results + output_line
labels = true_targets[true_targets[:, 0] == si, 1:]
nl = len(labels)
tcls = labels[:, 4].tolist() if nl else [] # target class
seen += 1
imagename = os.path.splitext(Path(str(si)).name)[0]
imagenames.append(imagename)
if nl:
unique_classes = np.unique(tcls).astype('int32')
# detected = []
for cls in unique_classes:
cls_idx = np.where(tcls == cls)[0]
tcls_tensor = labels[:, 4]
tcls_tensor = tcls_tensor[cls_idx]
npos_cls[cls] += len(tcls_tensor)
# target boxes
tbox = xywh2xyxy(labels[:, 0:4])
tbox = tbox[cls_idx]
tbox[:, [0, 2]] *= imgs.shape[3]
tbox[:, [1, 3]] *= imgs.shape[2]
bbox = np.array(tbox.cpu().numpy().round(), dtype=int)
det = [False] * len(tcls_tensor)
difficult = np.array(det)
class_recs[cls][imagename] = {
'bbox': bbox,
'difficult': difficult,
'det': det
}
p_80, r_80, ap_80, f1_80 = [], [], [], []
p_50, r_50, ap_50, f1_50 = [], [], [], []
ap_80_iou = 0.8
ap_50_iou = 0.5
class_recs_80 = copy.deepcopy(class_recs)
class_recs_50 = copy.deepcopy(class_recs)
for cls in range(nc):
rec_cls, prec_cls, ap_cls = voc_eval(det_lines,
npos_cls[cls],
imagenames,
class_recs_80[cls],
ovthresh=ap_80_iou,
use_07_metric=True)
f1_cls = 2 * prec_cls[-1] * rec_cls[-1] / (prec_cls[-1] + rec_cls[-1] +
1e-16)
p_80.append(prec_cls[-1])
r_80.append(rec_cls[-1])
ap_80.append(ap_cls)
f1_80.append(f1_cls)
print("AP 80")
print("person ap is: %.6f" % (ap_cls * 100))
print("recall is %.6f" % (rec_cls[-1] * 100))
print("precision is %.6f" % (prec_cls[-1] * 100))
rec_cls, prec_cls, ap_cls = voc_eval(det_lines,
npos_cls[cls],
imagenames,
class_recs_50[cls],
ovthresh=ap_50_iou,
use_07_metric=True)
f1_cls = 2 * prec_cls[-1] * rec_cls[-1] / (prec_cls[-1] + rec_cls[-1] +
1e-16)
p_50.append(prec_cls[-1])
r_50.append(rec_cls[-1])
ap_50.append(ap_cls)
f1_50.append(f1_cls)
print("AP 50")
print("person ap is: %.6f" % (ap_cls * 100))
print("recall is %.6f" % (rec_cls[-1] * 100))
print("precision is %.6f" % (prec_cls[-1] * 100))
mp_80, mr_80, map_80, mf1_80 = np.mean(p_80) * 100, np.mean(
r_80) * 100, np.mean(ap_80) * 100, np.mean(f1_80) * 100
mp_50, mr_50, map_50, mf1_50 = np.mean(p_50) * 100, np.mean(
r_50) * 100, np.mean(ap_50) * 100, np.mean(f1_50) * 100
# Print results
all_target_sum = 0
for _, cls_npos in npos_cls.items():
all_target_sum += cls_npos
pf = '%20s' + '%10.6g' * 6 # print format
# print(pf % ('all', seen, nt.sum(), mp, mr, map, mf1), end='\n\n')
print(pf % ('all', seen, all_target_sum, mp_80, mr_80, map_80, mf1_80),
end='\n\n')
print(pf % ('all', seen, all_target_sum, mp_50, mr_50, map_50, mf1_50),
end='\n\n')
# Print results per class
# if nc > 1:
# for i, c in enumerate(ap_class):
# print(pf % (names[c], seen, npos_cls[c], p[i], r[i], ap[i], f1[i]))
if nc > 1:
for i in range(nc):
print(pf % (names[i], seen, npos_cls[i], p_80[i], r_80[i],
ap_80[i], f1_80[i]))
print(pf % (names[i], seen, npos_cls[i], p_50[i], r_50[i],
ap_50[i], f1_50[i]))
# Return results
maps = np.zeros(nc)
# for i, c in enumerate(ap_class):
# maps[c] = ap[i]
for i in range(nc):
maps[i] = ap_80[i]
return (mp_80, mr_80, map_80, mf1_80, loss / len(dataloader), mp_50, mr_50,
map_50, mf1_50), maps
def eval(config):
"""
:param config:
:return:
"""
is_training = False
# Load and initialize network
# net = ProposalModel(config, is_training=is_training)
net = ProposalAttention(config, is_training=is_training)
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
# Restore pretrain model
if config["pretrain_snapshot"]:
state_dict = torch.load(config["pretrain_snapshot"])
net.load_state_dict(state_dict)
# YOLO loss with 3 scales
val_losses = []
for i in range(3):
val_losses.append(
ProposalLoss(config["yolo"]["anchors"][i],
(config["img_w"], config["img_h"])))
# DataLoader
val_loader = torch.utils.data.DataLoader(
COCOvalDataset(config["val_path"], (config["img_w"], config["img_h"])),
batch_size=16, # set batch size by 1
shuffle=False,
num_workers=2,
pin_memory=False)
""" VALIDATION """
total = 0.0
proposal = 0.0
correct = 0.0
net.eval()
img_cnt = 0
recall_cnt = 0.0
for step, samples in enumerate(val_loader):
images, labels = samples["image"], samples["label"]
with torch.no_grad():
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(val_losses[i](outputs[i]))
output = torch.cat(output_list, 1)
batch_detections = non_max_suppression(output)
# one image at a time !!!
for label_i in range(labels.size(0)):
total_avg = 0
correct_avg = 0
# calculate total
targets = labels[label_i]
for tx, ty, tw, th in targets:
if tw > 0:
total += 1
total_avg += 1
else:
continue
# calculate proposal
if batch_detections[label_i] is None:
continue
img_cnt += 1
predictions = batch_detections[label_i]
proposal += predictions.size(0)
# calculate correct
for tx, ty, tw, th in targets:
x1, x2 = config["img_w"] * (tx - tw / 2.0), config["img_w"] * (
tx + tw / 2.0)
y1, y2 = config["img_h"] * (ty - th / 2.0), config["img_h"] * (
ty + th / 2.0)
box_gt = [x1, y1, x2, y2, 1.0]
box_gt = torch.from_numpy(np.array(box_gt)).float().cuda()
best_iou = 0.0
for pred_i in range(predictions.size(0)):
iou = bbox_iou(predictions[pred_i].unsqueeze(0),
box_gt.unsqueeze(0))
iou = iou.item()
best_iou = max(iou, best_iou)
if best_iou >= 0.5:
correct += 1
correct_avg += 1
recall_cnt += float(correct_avg / float(total_avg))
if (step + 1) % 100 == 0:
print 'Total: %d\tProposal: %d\tCorrect: %d\tPrecision: %.4f\tRecall: %.4f' % (
total, proposal, correct, correct /
(proposal + 1e-6), correct / (total + 1e-6))
precision = correct / (proposal + 1e-6)
recall = correct / (total + 1e-6)
fscore = (2.0 * precision * recall) / (precision + recall + 1e-6)
print("Precision: %.4f\tRecall: %.4f\tFscore: %.4f" %
(precision, recall, fscore))
print("Avg Recall: %.4f" % (recall_cnt / float(img_cnt + 1e-6)))
def test(config):
is_training = False
# Load and initialize network
net = ModelMain(config, is_training=is_training)
net.train(is_training)
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
# Restore pretrain model
if config["pretrain_snapshot"]:
logging.info("load checkpoint from {}".format(
config["pretrain_snapshot"]))
state_dict = torch.load(config["pretrain_snapshot"])
net.load_state_dict(state_dict)
else:
raise Exception("missing pretrain_snapshot!!!")
# YOLO loss with 3 scales
yolo_losses = []
for i in range(3):
yolo_losses.append(
YOLOLayer(config["yolo"]["anchors"][i], config["yolo"]["classes"],
(config["img_w"], config["img_h"])))
# prepare images path
images_name = os.listdir(config["images_path"])
images_path = [
os.path.join(config["images_path"], name) for name in images_name
]
if len(images_path) == 0:
raise Exception("no image found in {}".format(config["images_path"]))
# Start testing FPS of different batch size
for batch_size in range(1, 10):
# preprocess
images = []
for path in images_path[:batch_size]:
image = cv2.imread(path, cv2.IMREAD_COLOR)
if image is None:
logging.error("read path error: {}. skip it.".format(path))
continue
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = cv2.resize(image, (config["img_w"], config["img_h"]),
interpolation=cv2.INTER_LINEAR)
image = image.astype(np.float32)
image /= 255.0
image = np.transpose(image, (2, 0, 1))
image = image.astype(np.float32)
images.append(image)
for i in range(batch_size - len(images)):
images.append(images[0]) # fill len to batch_sze
images = np.asarray(images)
images = torch.from_numpy(images).cuda()
# inference in 30 times and calculate average
inference_times = []
for i in range(30):
start_time = time.time()
with torch.no_grad():
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(yolo_losses[i](outputs[i]))
output = torch.cat(output_list, 1)
batch_detections = non_max_suppression(
output,
config["yolo"]["classes"],
conf_thres=config["confidence_threshold"])
torch.cuda.synchronize() # wait all done.
end_time = time.time()
inference_times.append(end_time - start_time)
inference_time = sum(inference_times) / len(
inference_times) / batch_size
fps = 1.0 / inference_time
logging.info(
"Batch_Size: {}, Inference_Time: {:.5f} s/image, FPS: {}".format(
batch_size, inference_time, fps))
def test(config):
is_training = False
anchors = [int(x) for x in config["yolo"]["anchors"].split(",")]
anchors = [[[anchors[i], anchors[i + 1]], [anchors[i + 2], anchors[i + 3]],
[anchors[i + 4], anchors[i + 5]]]
for i in range(0, len(anchors), 6)]
anchors.reverse()
config["yolo"]["anchors"] = []
for i in range(3):
config["yolo"]["anchors"].append(anchors[i])
net = ModelMain(config, is_training=is_training)
net.train(is_training)
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
ini_files = [
inifile for inifile in os.listdir(
os.path.join(config['test_weights'], 'result'))
if inifile.endswith('.ini')
]
accuracy_s = [(inifile[:-4]).split('_')[-1] for inifile in ini_files]
accuracy_ints = list(map(float, accuracy_s))
max_index = accuracy_ints.index(max(accuracy_ints))
# for kkk,ini_file in enumerate(ini_files):
ini_list_config = configparser.ConfigParser()
config_file_path = os.path.join(config['test_weights'], 'result',
ini_files[max_index])
Bi_picpath = os.path.join(config['test_weights'], 'result',
ini_files[max_index]).replace('.ini', '')
os.makedirs(Bi_picpath, exist_ok=True)
ini_list_config.read(config_file_path)
ini_session = ini_list_config.sections()
accuracy = ini_list_config.items(ini_session[0])
err_jpgfiles = ini_list_config.items(ini_session[1])
weight_file = os.path.join(
config['test_weights'],
'%s.weights' % ini_files[max_index].split('_')[0])
if weight_file: # Restore pretrain model
logging.info("load checkpoint from {}".format(weight_file))
state_dict = torch.load(weight_file)
net.load_state_dict(state_dict)
else:
raise Exception("missing pretrain_snapshot!!!")
yolo_losses = []
for i in range(3):
yolo_losses.append(
YOLOLayer(config["batch_size"], i, config["yolo"]["anchors"][i],
config["yolo"]["classes"],
(config["img_w"], config["img_h"])))
# images_name = os.listdir(config["images_path"]) # prepare images path
# images_path = [os.path.join(config["images_path"], name) for name in images_name]
# if len(images_path) == 0:
# raise Exception("no image found in {}".format(config["images_path"]))
# batch_size = config["batch_size"]# Start inference
# for step in range(0, len(images_path), batch_size):
for _jpg_images in err_jpgfiles:
images = [] # preprocess
images_origin = []
jpg_path = str(_jpg_images[1])
logging.info("processing: {}".format(jpg_path))
bbox_list = read_gt_boxes(jpg_path)
image = cv2.imread(jpg_path, cv2.IMREAD_COLOR)
if image is None:
logging.error("read path error: {}. skip it.".format(jpg_path))
continue
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
images_origin.append(image) # keep for save result
image = cv2.resize(image, (config["img_w"], config["img_h"]),
interpolation=cv2.INTER_LINEAR)
image = image.astype(np.float32)
image /= 255.0
image = np.transpose(image, (2, 0, 1))
image = image.astype(np.float32)
images.append(image)
images = np.asarray(images)
images = torch.from_numpy(images).cuda()
with torch.no_grad(): # inference
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(yolo_losses[i](outputs[i]))
output = torch.cat(output_list, 1)
batch_detections = non_max_suppression(
output,
config["yolo"]["classes"],
conf_thres=config["confidence_threshold"])
classes = open(config["classes_names_path"],
"r").read().split("\n")[:-1]
for idx, detections in enumerate(batch_detections):
image_show = images_origin[idx]
if detections is not None:
unique_labels = detections[:, -1].cpu().unique()
n_cls_preds = len(unique_labels)
for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
ori_h, ori_w = images_origin[
idx].shape[:
2] # Rescale coordinates to original dimensions
pre_h, pre_w = config["img_h"], config["img_w"]
box_h = ((y2 - y1) / pre_h) * ori_h
box_w = ((x2 - x1) / pre_w) * ori_w
y1 = (y1 / pre_h) * ori_h
x1 = (x1 / pre_w) * ori_w
image_show = cv2.rectangle(images_origin[idx], (x1, y1),
(x1 + box_w, y1 + box_h),
(0, 255, 0), 2)
for (x1, x2, y1, y2) in bbox_list:
[x1, x2, y1, y2] = map(int, [x1, x2, y1, y2])
cv2.rectangle(image_show, (x1, y1), (x2, y2), (255, 0, 0),
2)
pic_name = (jpg_path.split('/')[-1]).split('.')[0]
image_show = cv2.cvtColor(image_show, cv2.COLOR_RGB2BGR)
cv2.imwrite(
os.path.join(Bi_picpath,
'%s.jpg' % os.path.basename(pic_name)),
image_show)
def evaluate(config):
is_training = False
# Load and initialize network
net = ModelMain(config, is_training=is_training)
net.train(is_training)
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
# Restore pretrain model
if config["pretrain_snapshot"]:
state_dict = torch.load(config["pretrain_snapshot"])
net.load_state_dict(state_dict)
else:
logging.warning("missing pretrain_snapshot!!!")
# YOLO loss with 3 scales
yolo_losses = []
for i in range(3):
yolo_losses.append(YOLOLoss(config["yolo"]["anchors"][i],
config["yolo"]["classes"], (config["img_w"], config["img_h"])))
# DataLoader
dataloader = torch.utils.data.DataLoader(dataset=COCODataset(config["val_path"], config["img_w"]),
batch_size=config["batch_size"],
shuffle=True, num_workers=1, pin_memory=False)
# Start the eval loop
logging.info("Start eval.")
n_gt = 0
correct = 0
logging.info('%s' % str(dataloader))
gt_histro={}
pred_histro = {}
correct_histro = {}
for i in range(config["yolo"]["classes"]):
gt_histro[i] = 1
pred_histro[i] = 1
correct_histro[i] = 0
# images 是一个batch里的全部图片，labels是一个batch里面的全部标签
for step, (images, labels) in enumerate(dataloader):
labels = labels.cuda()
with torch.no_grad():
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(yolo_losses[i](outputs[i]))
# 把三个尺度上的预测结果在第1维度(第0维度是batch里的照片，第1维度是一张照片里面的各个预测框，第2维度是各个预测数值)上拼接起来
batch_output = torch.cat(output_list, dim=1)
logging.info('%s' % str(batch_output.shape))
# 进行NMS抑制
batch_output = non_max_suppression(prediction=batch_output, num_classes=config["yolo"]["classes"], conf_thres=config["conf_thresh"], nms_thres=config["nms_thresh"])
# calculate
for sample_index_in_batch in range(labels.size(0)):
# fetched img sample in tensor( C(RxGxB) x H x W ), transform to cv2 format in H x W x C(BxGxR)
sample_image = images[sample_index_in_batch].numpy()
sample_image = np.transpose(sample_image, (1, 2, 0))
sample_image = cv2.cvtColor(sample_image, cv2.COLOR_RGB2BGR)
logging.debug("fetched img %d size %s" % (sample_index_in_batch, sample_image.shape))
# Get labels for sample where width is not zero (dummies)(init all labels to zeros in array)
target_sample = labels[sample_index_in_batch, labels[sample_index_in_batch, :, 3] != 0]
# get prediction for this sample
sample_pred = batch_output[sample_index_in_batch]
if sample_pred is not None:
for x1, y1, x2, y2, conf, obj_conf, obj_pred in sample_pred: # for each prediction box
# logging.info("%d" % obj_cls)
box_pred = torch.cat([coord.unsqueeze(0) for coord in [x1, y1, x2, y2]]).view(1, -1)
sample_image = draw_prediction(sample_image,conf, obj_conf, int(obj_pred), (x1, y1, x2, y2), config)
# 每一个ground truth的 分类编号obj_cls、相对中心x、相对中心y、相对宽w、相对高h
for obj_cls, tx, ty, tw, th in target_sample:
# Get rescaled gt coordinates
# 转化为输入像素尺寸的 左上角像素tx1 ty1，右下角像素tx2 ty2
tx1, tx2 = config["img_w"] * (tx - tw / 2), config["img_w"] * (tx + tw / 2)
ty1, ty2 = config["img_h"] * (ty - th / 2), config["img_h"] * (ty + th / 2)
# 计算ground truth数量，用于统计信息
n_gt += 1
gt_histro[int(obj_cls)] += 1
# 转化为 shape(1,4)的tensor，用来计算IoU
box_gt = torch.cat([coord.unsqueeze(0) for coord in [tx1, ty1, tx2, ty2]]).view(1, -1)
# logging.info('%s' % str(box_gt.shape))
sample_pred = batch_output[sample_index_in_batch]
if sample_pred is not None:
# Iterate through predictions where the class predicted is same as gt
# 对于每一个ground truth，遍历预测结果
for x1, y1, x2, y2, conf, obj_conf, obj_pred in sample_pred[sample_pred[:, 6] == obj_cls]: # 如果当前预测分类 == 当前真实分类
#logging.info("%d" % obj_cls)
box_pred = torch.cat([coord.unsqueeze(0) for coord in [x1, y1, x2, y2]]).view(1, -1)
pred_histro[int(obj_pred)] += 1
iou = bbox_iou(box_pred, box_gt)
if iou >= config["iou_thresh"]:
correct += 1
correct_histro[int(obj_pred)] += 1
break
if n_gt:
types = config["types"]
reverse_types = {} # 建立一个反向的types
for key in types.keys():
reverse_types[types[key]] = key
logging.info('Batch [%d/%d] mAP: %.5f' % (step, len(dataloader), float(correct / n_gt)))
logging.info('mAP Histro:%s' % str([ reverse_types[i] +':'+ str(int(100 * correct_histro[i] / gt_histro[i])) for i in range(config["yolo"]["classes"] ) ]))
logging.info('Recall His:%s' % str([ reverse_types[i] +':'+ str(int(100 * correct_histro[i] / pred_histro[i])) for i in range(config["yolo"]["classes"]) ]))
logging.info('Mean Average Precision: %.5f' % float(correct / n_gt))
def test(config,int_dir='result'):
is_training = False
anchors = [int(x) for x in config["yolo"]["anchors"].split(",")]
anchors = [[[anchors[i], anchors[i + 1]], [anchors[i + 2], anchors[i + 3]], [anchors[i + 4], anchors[i + 5]]] for i
in range(0, len(anchors), 6)]
anchors.reverse()
config["yolo"]["anchors"] = []
for i in range(3):
config["yolo"]["anchors"].append(anchors[i])
net = ModelMain(config, is_training=is_training)
net.train(is_training)
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
ini_files = os.listdir(os.path.join(config['test_weights'], int_dir))
for kkk,ini_file in enumerate(ini_files):
ini_list_config = configparser.ConfigParser()
config_file_path = os.path.join(config['test_weights'], int_dir,ini_files[-kkk-1])
ini_list_config.read(config_file_path)
ini_session = ini_list_config.sections()
# accuracy = ini_list_config.items(ini_session[0])
err_jpgfiles = ini_list_config.items(ini_session[1])
aaa = glob.glob(os.path.join(config['test_weights'],'*_%s.weights'%ini_files[-kkk-1].split('_')[-1].split('.')[0]))
weight_file = aaa[0]#os.path.join(config['test_weights'],'%s.weights'%ini_files[-kkk-1].split('_')[0])
if weight_file: # Restore pretrain model
logging.info("load checkpoint from {}".format(weight_file))
state_dict = torch.load(weight_file)
net.load_state_dict(state_dict)
else:
raise Exception("missing pretrain_snapshot!!!")
yolo_losses = []
for i in range(3):
yolo_losses.append(YOLOLayer(1, i, config["yolo"]["anchors"][i],
config["yolo"]["classes"], (config["img_w"], config["img_h"])))
for index, _jpg_images in enumerate(err_jpgfiles):
images = []# preprocess
images_origin = []
jpg_path = str(_jpg_images[1])
print(str(index+1),jpg_path)
bbox_list = read_gt_boxes(jpg_path)
image = cv2.imread(jpg_path, cv2.IMREAD_COLOR)
if image is None:
logging.error("read path error: {}. skip it.".format(jpg_path))
continue
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
images_origin.append(image) # keep for save result
image = cv2.resize(image, (config["img_w"], config["img_h"]),interpolation=cv2.INTER_LINEAR)
image = image.astype(np.float32)
image /= 255.0
image = np.transpose(image, (2, 0, 1))
image = image.astype(np.float32)
images.append(image)
images = np.asarray(images)
images = torch.from_numpy(images).cuda()
with torch.no_grad():# inference
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(yolo_losses[i](outputs[i]))
output = torch.cat(output_list, 1)
batch_detections = non_max_suppression(output, config["yolo"]["classes"],
conf_thres=config["confidence_threshold"])
classes = open(config["classes_names_path"], "r").read().split("\n")[:-1]
if not os.path.isdir("./output/"):
os.makedirs("./output/")
for idx, detections in enumerate(batch_detections):
image_show=images_origin[idx]
if detections is not None:
unique_labels = detections[:, -1].cpu().unique()
n_cls_preds = len(unique_labels)
for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
ori_h, ori_w = images_origin[idx].shape[:2]# Rescale coordinates to original dimensions
pre_h, pre_w = config["img_h"], config["img_w"]
box_h = ((y2 - y1) / pre_h) * ori_h
box_w = ((x2 - x1) / pre_w) * ori_w
y1 = (y1 / pre_h) * ori_h
x1 = (x1 / pre_w) * ori_w
#绿色代表预测，红色代表标注
image_show = cv2.rectangle(images_origin[idx], (x1, y1), (x1 + box_w, y1 + box_h), (0, 255, 0),2)
for (x1, x2, y1, y2) in bbox_list:
[x1, x2, y1, y2] = map(int, [x1, x2, y1, y2])
cv2.rectangle(image_show, (x1, y1), (x2, y2), (0, 0, 255), 2)
cv2.imshow('1', image_show)
cv2.waitKey()
def test(config):
is_training = False
anchors = [int(x) for x in config["yolo"]["anchors"].split(",")]
anchors = [[[anchors[i], anchors[i + 1]], [anchors[i + 2], anchors[i + 3]], [anchors[i + 4], anchors[i + 5]]] for i
in range(0, len(anchors), 6)]
anchors.reverse()
config["yolo"]["anchors"] = []
for i in range(3):
config["yolo"]["anchors"].append(anchors[i])
# Load and initialize network
net = ModelMain(config, is_training=is_training)
net.train(is_training)
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
# Restore pretrain model
if config["pretrain_snapshot"]:
logging.info("load checkpoint from {}".format(config["pretrain_snapshot"]))
state_dict = torch.load(config["pretrain_snapshot"])
net.load_state_dict(state_dict)
else:
raise Exception("missing pretrain_snapshot!!!")
# YOLO loss with 3 scales
yolo_losses = []
for i in range(3):
yolo_losses.append(YOLOLayer(config["batch_size"],i,config["yolo"]["anchors"][i],
config["yolo"]["classes"], (config["img_w"], config["img_h"])))
# prepare images path
images_path = os.listdir(config["images_path"])
images_path = [file for file in images_path if file.endswith('.jpg')]
# images_path = [os.path.join(config["images_path"], name) for name in images_name]
if len(images_path) == 0:
raise Exception("no image found in {}".format(config["images_path"]))
# Start inference
batch_size = config["batch_size"]
bgimage = cv2.imread(os.path.join(config["images_path"], images_path[0]), cv2.IMREAD_COLOR)
bgimage = cv2.cvtColor(bgimage, cv2.COLOR_BGR2GRAY)
for step in range(0, len(images_path)-1, batch_size):
# preprocess
images = []
images_origin = []
for path in images_path[step*batch_size: (step+1)*batch_size]:
if not path.endswith(".jpg") and (not path.endswith(".png")) and not path.endswith(".JPEG"):
continue
image = cv2.imread(os.path.join(config["images_path"], path), cv2.IMREAD_COLOR)
if image is None:
logging.error("read path error: {}. skip it.".format(path))
continue
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
images_origin.append(image) # keep for save result
image = cv2.resize(image, (config["img_w"], config["img_h"]),
interpolation=cv2.INTER_LINEAR)
image = image.astype(np.float32)
image /= 255.0
image = np.transpose(image, (2, 0, 1))
image = image.astype(np.float32)
images.append(image)
images = np.asarray(images)
images = torch.from_numpy(images).cuda()
# inference
with torch.no_grad():
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(yolo_losses[i](outputs[i]))
output = torch.cat(output_list, 1)
batch_detections = non_max_suppression(output, config["yolo"]["classes"],
conf_thres=config["confidence_threshold"])
# write result images. Draw bounding boxes and labels of detections
classes = open(config["classes_names_path"], "r").read().split("\n")[:-1]
for idx, detections in enumerate(batch_detections):
image_show =images_origin[idx]
if detections is not None:
anno = savexml.GEN_Annotations(path + '.jpg')
anno.set_size(1280, 720, 3)
unique_labels = detections[:, -1].cpu().unique()
n_cls_preds = len(unique_labels)
bbox_list = []
for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
# Rescale coordinates to original dimensions
ori_h, ori_w = images_origin[idx].shape[:2]
pre_h, pre_w = config["img_h"], config["img_w"]
box_h = ((y2 - y1) / pre_h) * ori_h
box_w = ((x2 - x1) / pre_w) * ori_w
y1 = (y1 / pre_h) * ori_h
x1 = (x1 / pre_w) * ori_w
# Create a Rectangle patch
bbox_list.append((x1, y1,box_w,box_h))
image_show = cv2.rectangle(images_origin[idx], (x1, y1), (x1 + box_w, y1 + box_h), (0, 0, 255), 1)
boundbox = bg_judge(images_origin[idx],bgimage,bbox_list)
print('boundbox',boundbox,bbox_list)
for (x,y,w,h) in boundbox:
image_show = cv2.rectangle(image_show, (x, y), (x + w, y + h), (0, 255, 0), 1)
# anno.add_pic_attr("mouse", int(x1.cpu().data), int(y1.cpu().data), int(box_w.cpu().data) , int(box_h.cpu().data) ,"0")
#
# xml_path = os.path.join(config["images_path"], path).replace('rec_pic',r'detect_pic1\Annotations').replace('jpg','xml')
# anno.savefile(xml_path)
# cv2.imwrite(os.path.join(config["images_path"], path).replace('rec_pic',r'detect_pic1\rec_pic'),images_origin[idx])
cv2.imshow('1', image_show)
cv2.waitKey(1)
logging.info("Save all results to ./output/")
def vis(config):
"""
:param config:
:return:
"""
is_training = False
# Load and initialize network
# net = ProposalModel(config, is_training=is_training)
net = ProposalAttention(config, is_training=is_training)
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
# Restore pretrain model
if config["pretrain_snapshot"]:
state_dict = torch.load(config["pretrain_snapshot"])
net.load_state_dict(state_dict)
# YOLO loss with 3 scales
val_losses = []
for i in range(3):
val_losses.append(
ProposalLoss(config["yolo"]["anchors"][i],
(config["img_w"], config["img_h"])))
# DataLoader
val_loader = torch.utils.data.DataLoader(COCOvalDataset(
config["val_path"], (config["img_w"], config["img_h"])),
batch_size=4,
shuffle=False,
num_workers=2,
pin_memory=False)
sample = None
step_d = 1
step_i = 1
cnt = 3
for i, sample in enumerate(val_loader):
if i == cnt:
break
# Detection
images, labels = sample["image"], sample["label"]
with torch.no_grad():
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(val_losses[i](outputs[i]))
output = torch.cat(output_list, 1)
batch_detections = non_max_suppression(output)
for idx, detections in enumerate(batch_detections):
plt.figure()
fig, ax = plt.subplots(1)
ax.imshow(np.transpose(images[idx].numpy(), (1, 2, 0)))
if detections is not None:
for x1, y1, x2, y2, conf in detections:
bbox_colors = random.sample(colors,
len(batch_detections))
color = bbox_colors[idx]
pre_h, pre_w = config["img_h"], config["img_w"]
if x1 < 0:
x1 = 0
if x2 > pre_w:
x2 = pre_w
if y1 < 0:
y1 = 0
if y2 > pre_h:
y2 = pre_h
box_h = (y2 - y1)
box_w = (x2 - x1)
# Create a Rectangle patch
bbox = patches.Rectangle((x1, y1),
box_w,
box_h,
linewidth=2,
edgecolor=color,
facecolor='none')
# Add the bbox to the plot
ax.add_patch(bbox)
# ax.text(x1, y1, '{:.2f}'.format(conf), bbox=dict(facecolor=color, alpha=0.9), fontsize=8, color='white')
else:
print 'Nothing detected.'
# Save generated image with detections
plt.axis('off')
plt.gca().xaxis.set_major_locator(NullLocator())
plt.gca().yaxis.set_major_locator(NullLocator())
plt.savefig('vis10/img{}_detect.jpg'.format(step_d),
bbox_inches='tight',
pad_inches=0.0)
plt.close()
step_d += 1
# Ground-truth
for i, (image,
label) in enumerate(zip(sample['image'], sample['label'])):
plt.figure()
fig, ax = plt.subplots(1)
ax.imshow(np.transpose(image.numpy(), (1, 2, 0)))
for l in label:
if l.sum() == 0:
continue
x1 = int((l[0] - l[2] / 2) * config["img_w"])
y1 = int((l[1] - l[3] / 2) * config["img_h"])
x2 = int((l[0] + l[2] / 2) * config["img_w"])
y2 = int((l[1] + l[3] / 2) * config["img_h"])
box_h = (y2 - y1)
box_w = (x2 - x1)
bbox = patches.Rectangle((x1, y1),
box_w,
box_h,
linewidth=2,
edgecolor='blue',
facecolor='none')
ax.add_patch(bbox)
# cv2.rectangle(image, (x1, y1), (x2, y2), (0, 0, 255))
plt.axis('off')
plt.gca().xaxis.set_major_locator(NullLocator())
plt.gca().yaxis.set_major_locator(NullLocator())
plt.savefig('vis10/img{}_input.jpg'.format(step_i),
bbox_inches='tight',
pad_inches=0.0)
plt.close()
step_i += 1
def eval_voc(self, val_dataset, classes, iou_thresh=0.5):
logging.info('Start Evaling')
results = {}
def voc_ap(rec, prec, use_07_metric=False):
""" ap = voc_ap(rec, prec, [use_07_metric])
Compute VOC AP given precision and recall.
If use_07_metric is true, uses the
VOC 07 11 point method (default:False).
"""
_rec = np.arange(0., 1.1, 0.1)
_prec = []
if use_07_metric:
# 11 point metric
ap = 0.
for t in np.arange(0., 1.1, 0.1):
if np.sum(rec >= t) == 0:
p = 0
else:
p = np.max(prec[rec >= t])
_prec.append(p)
ap = ap + p / 11.
else:
# correct AP calculation
# first append sentinel values at the end
mrec = np.concatenate(([0.], rec, [1.]))
mpre = np.concatenate(([0.], prec, [0.]))
# compute the precision envelope
for i in range(mpre.size - 1, 0, -1):
mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])
# to calculate area under PR curve, look for points
# where X axis (recall) changes value
i = np.where(mrec[1:] != mrec[:-1])[0]
# and sum (\Delta recall) * prec
ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])
return ap
def caculate_ap(correct, conf, pred_cls, total, classes):
correct, conf, pred_cls = np.array(correct), np.array(
conf), np.array(pred_cls)
index = np.argsort(-conf)
correct, conf, pred_cls = correct[index], conf[index], pred_cls[
index]
ap = []
AP = {}
for i, c in enumerate(classes):
k = pred_cls == i
n_gt = total[c]
n_p = sum(k)
if n_gt == 0 and n_p == 0:
continue
elif n_p == 0 or n_gt == 0:
ap.append(0)
AP[c] = 0
else:
fpc = np.cumsum(1 - correct[k])
tpc = np.cumsum(correct[k])
rec = tpc / n_gt
prec = tpc / (tpc + fpc)
_ap = voc_ap(rec, prec)
ap.append(_ap)
AP[c] = _ap
mAP = np.array(ap).mean()
return mAP, AP
def parse_rec(imagename, classes):
filename = imagename.replace('jpg', 'xml')
tree = ET.parse(filename)
objects = []
for obj in tree.findall('object'):
difficult = obj.find('difficult').text
cls = obj.find('name').text
if cls not in classes or int(difficult) == 1:
continue
cls_id = classes.index(cls)
xmlbox = obj.find('bndbox')
obj = [
float(xmlbox.find('xmin').text),
float(xmlbox.find('xmax').text),
float(xmlbox.find('ymin').text),
float(xmlbox.find('ymax').text), cls_id
]
objects.append(obj)
return np.asarray(objects)
total = {}
for cls in classes:
total[cls] = 0
correct = []
conf_list = []
pred_list = []
for step, samples in enumerate(val_dataset):
images, labels = samples['image'], samples['label']
image_paths, origin_sizes = samples['image_path'], samples[
'origin_size']
logging.info("Now have finished [%.3d/%.3d]" %
(step, len(val_dataset)))
with torch.no_grad():
outputs = self.net(images)
output_list = []
for i in range(3):
output_list.append(self.yolo_loss[i](outputs[i]))
output = torch.cat(output_list, 1)
batch_detections = non_max_suppression(output,
self.config.num_classes,
conf_thres=0.001,
nms_thres=0.4)
for idx, detections in enumerate(batch_detections):
image_path = image_paths[idx]
label = labels[idx]
for t in range(label.size(0)):
if label[t, :].sum() == 0:
label = label[:t, :]
break
label_cls = np.array(label[:, 0])
for cls_id in label_cls:
total[classes[int(cls_id)]] += 1
if detections is None:
if label.size(0) != 0:
label_cls = np.unique(label_cls)
for cls_id in label_cls:
correct.append(0)
conf_list.append(1)
pred_list.append(int(cls_id))
continue
if label.size(0) == 0:
for *pred_box, conf, cls_conf, cls_pred in detections:
correct.append(0)
conf_list.append(conf)
pred_list.append(int(cls_pred))
else:
detections = detections[np.argsort(-detections[:, 4])]
detected = []
for *pred_box, conf, cls_conf, cls_pred in detections:
pred_box = torch.FloatTensor(pred_box).view(1, -1)
pred_box[:, 2:] = pred_box[:, 2:] - pred_box[:, :2]
pred_box[:, :2] = pred_box[:, :2] + pred_box[:, 2:] / 2
pred_box = pred_box / self.config.image_size
ious = bbox_iou(pred_box, label[:, 1:])
best_i = np.argmax(ious)
if ious[best_i] > iou_thresh and int(cls_pred) == int(
label[best_i, 0]) and best_i not in detected:
correct.append(1)
detected.append(best_i)
else:
correct.append(0)
pred_list.append(int(cls_pred))
conf_list.append(float(conf))
results['correct'] = correct
results['conf'] = conf_list
results['pred_cls'] = pred_list
results['total'] = total
with open('results.json', 'w') as f:
json.dump(results, f)
logging.info('Having saved to results.json')
logging.info('Begin calculating....')
with open('results.json', 'r') as result_file:
results = json.load(result_file)
mAP, AP_class = caculate_ap(correct=results['correct'],
conf=results['conf'],
pred_cls=results['pred_cls'],
total=results['total'],
classes=classes)
logging.info('mAP(IoU=0.5):{:.1f}'.format(mAP * 100))
def evaluate(config):
is_training = False
# Load and initialize network
net = ModelMain(config, is_training=is_training)
net.train(is_training)
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
# Restore pretrain model
if config["pretrain_snapshot"]:
logging.info("Load checkpoint: {}".format(config["pretrain_snapshot"]))
state_dict = torch.load(config["pretrain_snapshot"])
net.load_state_dict(state_dict)
else:
logging.warning("missing pretrain_snapshot!!!")
# YOLO loss with 3 scales
yolo_losses = []
for i in range(3):
yolo_losses.append(
YOLOLoss(config["yolo"]["anchors"][i], config["yolo"]["classes"],
(config["img_w"], config["img_h"])))
# DataLoader.
dataloader = torch.utils.data.DataLoader(COCODataset(
config["val_path"], (config["img_w"], config["img_h"]),
is_training=False),
batch_size=config["batch_size"],
shuffle=False,
num_workers=8,
pin_memory=False)
# Coco Prepare.
index2category = json.load(open("coco_index2category.json"))
# Start the eval loop
logging.info("Start eval.")
coco_results = []
coco_img_ids = set([])
for step, samples in enumerate(dataloader):
images, labels = samples["image"], samples["label"]
image_paths, origin_sizes = samples["image_path"], samples[
"origin_size"]
with torch.no_grad():
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(yolo_losses[i](outputs[i]))
output = torch.cat(output_list, 1)
batch_detections = non_max_suppression(output,
config["yolo"]["classes"],
conf_thres=0.01,
nms_thres=0.45)
for idx, detections in enumerate(batch_detections):
image_id = int(os.path.basename(image_paths[idx])[-16:-4])
coco_img_ids.add(image_id)
if detections is not None:
origin_size = eval(origin_sizes[idx])
detections = detections.cpu().numpy()
for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
x1 = x1 / config["img_w"] * origin_size[0]
x2 = x2 / config["img_w"] * origin_size[0]
y1 = y1 / config["img_h"] * origin_size[1]
y2 = y2 / config["img_h"] * origin_size[1]
w = x2 - x1
h = y2 - y1
coco_results.append({
"image_id":
image_id,
"category_id":
index2category[str(int(cls_pred.item()))],
"bbox": (float(x1), float(y1), float(w), float(h)),
"score":
float(conf),
})
logging.info("Now {}/{}".format(step, len(dataloader)))
save_results_path = "coco_results.json"
with open(save_results_path, "w") as f:
json.dump(coco_results,
f,
sort_keys=True,
indent=4,
separators=(',', ':'))
logging.info("Save coco format results to {}".format(save_results_path))
# COCO api
logging.info("Using coco-evaluate tools to evaluate.")
cocoGt = COCO(config["annotation_path"])
cocoDt = cocoGt.loadRes(save_results_path)
cocoEval = COCOeval(cocoGt, cocoDt, "bbox")
cocoEval.params.imgIds = list(coco_img_ids) # real imgIds
cocoEval.evaluate()
cocoEval.accumulate()
cocoEval.summarize()
def test(config):
is_training = False
# Load and initialize network
net = ModelMain(config, is_training=is_training)
net.train(is_training)
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
# Restore pretrain model
if config["pretrain_snapshot"]:
logging.info("load checkpoint from {}".format(
config["pretrain_snapshot"]))
state_dict = torch.load(config["pretrain_snapshot"])
net.load_state_dict(state_dict)
else:
raise Exception("missing pretrain_snapshot!!!")
# YOLO loss with 3 scales
yolo_losses = []
for i in range(3):
yolo_losses.append(
YOLOLoss(config["yolo"]["anchors"][i], config["yolo"]["classes"],
(config["img_w"], config["img_h"])))
# prepare images path
images_name = os.listdir(config["images_path"])
images_path = [
os.path.join(config["images_path"], name) for name in images_name
]
if len(images_path) == 0:
raise Exception("no image found in {}".format(config["images_path"]))
# Start inference
batch_size = config["batch_size"]
for step in range(0, len(images_path), batch_size):
# preprocess
images = []
images_origin = []
for path in images_path[step * batch_size:(step + 1) * batch_size]:
logging.info("processing: {}".format(path))
image = cv2.imread(path, cv2.IMREAD_COLOR)
if image is None:
logging.error("read path error: {}. skip it.".format(path))
continue
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
images_origin.append(image) # keep for save result
image = cv2.resize(image, (config["img_w"], config["img_h"]),
interpolation=cv2.INTER_LINEAR)
image = image.astype(np.float32)
image /= 255.0
image = np.transpose(image, (2, 0, 1))
image = image.astype(np.float32)
images.append(image)
images = np.asarray(images)
images = torch.from_numpy(images).cuda()
# inference
with torch.no_grad():
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(yolo_losses[i](outputs[i]))
output = torch.cat(output_list, 1)
batch_detections = non_max_suppression(
output,
config["yolo"]["classes"],
conf_thres=config["confidence_threshold"])
# write result images. Draw bounding boxes and labels of detections
classes = open(config["classes_names_path"],
"r").read().split("\n")[:-1]
if not os.path.isdir("./output/"):
os.makedirs("./output/")
for idx, detections in enumerate(batch_detections):
plt.figure()
fig, ax = plt.subplots(1)
ax.imshow(images_origin[idx])
if detections is not None:
unique_labels = detections[:, -1].cpu().unique()
n_cls_preds = len(unique_labels)
bbox_colors = random.sample(colors, n_cls_preds)
for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
color = bbox_colors[int(
np.where(unique_labels == int(cls_pred))[0])]
# Rescale coordinates to original dimensions
ori_h, ori_w = images_origin[idx].shape[:2]
pre_h, pre_w = config["img_h"], config["img_w"]
box_h = ((y2 - y1) / pre_h) * ori_h
box_w = ((x2 - x1) / pre_w) * ori_w
y1 = (y1 / pre_h) * ori_h
x1 = (x1 / pre_w) * ori_w
# Create a Rectangle patch
bbox = patches.Rectangle((x1, y1),
box_w,
box_h,
linewidth=2,
edgecolor=color,
facecolor='none')
# Add the bbox to the plot
ax.add_patch(bbox)
# Add label
plt.text(x1,
y1,
s=classes[int(cls_pred)],
color='white',
verticalalignment='top',
bbox={
'color': color,
'pad': 0
})
# Save generated image with detections
plt.axis('off')
plt.gca().xaxis.set_major_locator(NullLocator())
plt.gca().yaxis.set_major_locator(NullLocator())
plt.savefig('output/{}_{}.jpg'.format(step, idx),
bbox_inches='tight',
pad_inches=0.0)
plt.close()
logging.info("Save all results to ./output/")
def test():
is_traning = False # 不训练，测试
# Load and initialize network
net = ModelMain(is_training=is_traning)
net.train(is_traning)
# Set data parallel
net = nn.DataParallel(net)
net = net.cuda()
# Restore pretrain model
if trained_model_dir:
logging.info("load checkpoint from {}".format(trained_model_dir))
state_dict = torch.load(trained_model_dir)
net.load_state_dict(state_dict)
else:
raise Exception("missing pretrain_snapshot!!!")
# YOLO loss with 3 scales
yolo_losses = []
for i in range(3):
yolo_losses.append(YOLOLoss(anchors[i], yolo_class_num,
(img_h, img_w)))
# prepare the images path
images_name = os.listdir("./images/")
images_path = [os.path.join("./images/", name) for name in images_name]
print('images_name：', images_name)
print('images_path：', len(images_path), images_path)
if len(images_path) == 0:
raise Exception("no image found in {}".format("./images/"))
# Start inference
batch_size = 16
for step in range(0, len(images_path),
batch_size): # range(0, 4, 16) step = 0, 4, 8, 12
logging.info('Batch_size:{}'.format(batch_size))
# preprocess
images = [] # 输入网络图片组
images_origin = [] # 原始图片组
for path in images_path[step * batch_size:(step + 1) * batch_size]:
logging.info("processing: {}".format(path))
image = cv2.imread(path, cv2.IMREAD_COLOR)
# cv2.imshow('Image', image)
# cv2.waitKey(0)
logging.info(" √ Successfully Processed! √")
if image is None:
logging.error("read path error: {}. skip it.".format(path))
continue
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
images_origin.append(image) # 预处理完毕，加入原始图片组
# 进一步将图片处理为网络可以接受的数据类型（resize、归一化等）
image = cv2.resize(image, (img_h, img_w),
interpolation=cv2.INTER_LINEAR)
image = image.astype(np.float32)
image /= 255.0
image = np.transpose(image, (2, 0, 1))
image = image.astype(np.float32)
images.append(image) # 归一化完毕，加入输入网络图片组
images = np.asarray(images)
images = torch.from_numpy(images).cuda()
logging.info("\nImages Convert to Tensor of CUDA Done!")
# inference
with torch.no_grad():
outputs = net(images)
output_list = []
for i in range(3):
output_list.append(yolo_losses[i](outputs[i]))
output = torch.cat(output_list, 1)
batch_detections = non_max_suppression(prediction=output,
num_classes=yolo_class_num,
conf_thres=0.5)
logging.info("\nNet Detection Done!\n")
# write result images: Draw BBox
classes = open(classes_name_path,
'r').read().split("\n")[:-1] # 读取coco.names
if not os.path.isdir("./output/"):
os.makedirs("./output/")
for idx, detections in enumerate(batch_detections):
plt.figure()
fig, ax = plt.subplots(1)
ax.imshow(images_origin[idx])
if detections is not None:
unique_labels = detections[:, -1].cpu().unique()
n_cls_preds = len(unique_labels)
bbox_colors = random.sample(colors, n_cls_preds)
# print('Final Detections: ', detections)
for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
color = bbox_colors[int(
np.where(unique_labels == int(cls_pred))[0])]
# Rescale coordinates to original dimensions
ori_h, ori_w = images_origin[idx].shape[:2]
pre_h, pre_w = img_h, img_w # 416, 416
box_h = ((y2 - y1) / pre_h) * ori_h
box_w = ((x2 - x1) / pre_w) * ori_w
y1 = (y1 / pre_h) * ori_h
x1 = (x1 / pre_w) * ori_w
# Create a Rectangle patch
bbox = patches.Rectangle((x1, y1),
box_w,
box_h,
linewidth=2,
edgecolor=color,
facecolor='none')
# Add the bbox to the plot
ax.add_patch(bbox)
# Add label
plt.text(x1,
y1,
s=classes[int(cls_pred)],
color='white',
verticalalignment='top',
bbox={
'color': color,
'pad': 0
})
# Save generated image with detections
plt.axis('off')
plt.gca().xaxis.set_major_locator(NullLocator())
plt.gca().yaxis.set_major_locator(NullLocator())
plt.savefig('output/{}_{}.jpg'.format(step, idx),
bbox_inches='tight',
pad_inches=0.0)
plt.close()
logging.info("All the Test Process Succeed! Enjoy it!")
