net = net.to(DEVICE)
print(f'It took {timer.end("Load Model")} seconds to load the model.')
predictor = Predictor(net, config.image_size, config.image_mean,
config.image_std,
nms_method=args.nms_method,
iou_threshold=config.iou_threshold,
candidate_size=200,
sigma=0.5,
device=device)
results = []
for i in range(len(dataset)):
print("process image", i)
timer.start("Load Image")
image = dataset.get_image(i)
print("Load Image: {:4f} seconds.".format(timer.end("Load Image")))
timer.start("Predict")
boxes, labels, probs = predictor.predict(image)
if args.net == 'basenet':
continue
else:
net.detach_hidden()
print("Prediction: {:4f} seconds.".format(timer.end("Predict")))
indexes = torch.ones(labels.size(0), 1, dtype=torch.float32) * i
results.append(torch.cat([
indexes.reshape(-1, 1),
labels.reshape(-1, 1).float(),
probs.reshape(-1, 1),
boxes + 1.0 # matlab's indexes start from 1
], dim=1))
results = torch.cat(results)
def main():
parser = argparse.ArgumentParser(
description="SSD Evaluation on VOC Dataset.")
parser.add_argument("--trained_model", type=str)
parser.add_argument(
"--dataset_type",
default="voc",
type=str,
help='Specify dataset type. Currently support voc and open_images.')
parser.add_argument(
"--dataset",
type=str,
help="The root directory of the VOC dataset or Open Images dataset.")
parser.add_argument("--label_file", type=str, help="The label file path.")
parser.add_argument("--use_cuda", type=str2bool, default=True)
parser.add_argument("--use_2007_metric", type=str2bool, default=True)
parser.add_argument("--nms_method", type=str, default="hard")
parser.add_argument("--iou_threshold",
type=float,
default=0.5,
help="The threshold of Intersection over Union.")
parser.add_argument("--eval_dir",
default="eval_results",
type=str,
help="The directory to store evaluation results.")
parser.add_argument('--mb2_width_mult',
default=1.0,
type=float,
help='Width Multiplifier for MobilenetV2')
args = parser.parse_args()
DEVICE = torch.device(
"cuda:0" if torch.cuda.is_available() and args.use_cuda else "cpu")
eval_path = pathlib.Path(args.eval_dir)
eval_path.mkdir(exist_ok=True)
timer = Timer()
class_names = [name.strip() for name in open(args.label_file).readlines()]
if args.dataset_type == "voc":
dataset = VOCDataset(args.dataset, is_test=True)
true_case_stat, all_gb_boxes, all_difficult_cases = group_annotation_by_class(
dataset)
net = create_mobilenetv2_ssd_lite(len(class_names),
width_mult=args.mb2_width_mult,
is_test=True)
timer.start("Load Model")
net.load(args.trained_model)
net = net.to(DEVICE)
print(f'It took {timer.end("Load Model")} seconds to load the model.')
predictor = create_mobilenetv2_ssd_lite_predictor(
net, nms_method=args.nms_method, device=DEVICE)
results = []
for i in range(len(dataset)):
print("process image", i)
timer.start("Load Image")
image = dataset.get_image(i)
print("Load Image: {:4f} seconds.".format(timer.end("Load Image")))
timer.start("Predict")
boxes, labels, probs = predictor.predict(image)
print("Prediction: {:4f} seconds.".format(timer.end("Predict")))
indexes = torch.ones(labels.size(0), 1, dtype=torch.float32) * i
results.append(
torch.cat(
[
indexes.reshape(-1, 1),
labels.reshape(-1, 1).float(),
probs.reshape(-1, 1),
boxes + 1.0 # matlab's indexes start from 1
],
dim=1))
results = torch.cat(results)
for class_index, class_name in enumerate(class_names):
if class_index == 0: continue # ignore background
prediction_path = eval_path / f"det_test_{class_name}.txt"
with open(prediction_path, "w") as f:
sub = results[results[:, 1] == class_index, :]
for i in range(sub.size(0)):
prob_box = sub[i, 2:].numpy()
image_id = dataset.ids[int(sub[i, 0])]
print(image_id + " " + " ".join([str(v) for v in prob_box]),
file=f)
aps = []
print("\n\nAverage Precision Per-class:")
for class_index, class_name in enumerate(class_names):
if class_index == 0:
continue
prediction_path = eval_path / f"det_test_{class_name}.txt"
ap = compute_average_precision_per_class(
true_case_stat[class_index], all_gb_boxes[class_index],
all_difficult_cases[class_index], prediction_path,
args.iou_threshold, args.use_2007_metric)
aps.append(ap)
print(f"{class_name}: {ap}")
print(f"\nAverage Precision Across All Classes:{sum(aps)/len(aps)}")
class Predictor:
def __init__(self,
net,
size,
mean=0.0,
std=1.0,
nms_method=None,
iou_threshold=0.45,
filter_threshold=0.01,
candidate_size=200,
sigma=0.5,
device=None):
"""Implement Predictor while testing of the model
Arguments:
net: an object of net to used for prediction
size: variable containing size of image as input
mean: an array variable containing mean value of all the channels of input
std: variable containing standard deviation of image as input
nms_method : string variable contaning type of negative mining
iou_threshold : a float variable containing threshold value of IOU
filter_threshold : a float variable containing threshold value of filter
candidate_size: only consider the candidates with the highest scores
sigma: the parameter in score re-computation. (for soft negative mining)
scores[i] = scores[i] * exp(-(iou_i)^2 / simga)
device : variable containing device on which net needs to do computation
"""
self.net = net
self.transform = PredictionTransform(size, mean, std)
self.iou_threshold = iou_threshold
self.filter_threshold = filter_threshold
self.candidate_size = candidate_size
self.nms_method = nms_method
self.sigma = sigma
print("Predictor arg device", device)
if device:
self.device = device
else:
self.device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
self.net.to(self.device)
self.net.eval()
self.timer = Timer()
def predict(self, image, top_k=-1, prob_threshold=None):
"""Implement Predictor while testing of the model
Arguments:
image: image input for predictor
prob_threshold: threshold for probability
top_k: keep top_k results. If k <= 0, keep all the results.
Returns:
predicted boxes, labels and their probability
"""
cpu_device = torch.device("cpu")
height, width, _ = image.shape
image = self.transform(image)
images = image.unsqueeze(0)
images = images.to(self.device)
with torch.no_grad():
self.timer.start()
scores, boxes = self.net.forward(images)
print("Inference time: ", self.timer.end())
boxes = boxes[0]
scores = scores[0]
if not prob_threshold:
prob_threshold = self.filter_threshold
# this version of nms is slower on GPU, so we move data to CPU.
boxes = boxes.to(cpu_device)
scores = scores.to(cpu_device)
picked_box_probs = []
picked_labels = []
for class_index in range(1, scores.size(1)):
probs = scores[:, class_index]
mask = probs > prob_threshold
probs = probs[mask]
if probs.size(0) == 0:
continue
subset_boxes = boxes[mask, :]
box_probs = torch.cat([subset_boxes, probs.reshape(-1, 1)], dim=1)
box_probs = box_utils.nms(box_probs,
self.nms_method,
score_threshold=prob_threshold,
iou_threshold=self.iou_threshold,
sigma=self.sigma,
top_k=top_k,
candidate_size=self.candidate_size)
picked_box_probs.append(box_probs)
picked_labels.extend([class_index] * box_probs.size(0))
if not picked_box_probs:
return torch.tensor([]), torch.tensor([]), torch.tensor([])
picked_box_probs = torch.cat(picked_box_probs)
picked_box_probs[:, 0] *= width
picked_box_probs[:, 1] *= height
picked_box_probs[:, 2] *= width
picked_box_probs[:, 3] *= height
return picked_box_probs[:, :4], torch.tensor(
picked_labels), picked_box_probs[:, 4]
class Predictor:
def __init__(self, net, size, mean=0.0, std=1.0, nms_method=None,
iou_threshold=0.45, filter_threshold=0.01, candidate_size=200, sigma=0.5, device=None):
self.net = net
self.transform = PredictionTransform(size, mean, std)
self.iou_threshold = iou_threshold
self.filter_threshold = filter_threshold
self.candidate_size = candidate_size
self.nms_method = nms_method
self.sigma = sigma
if device:
self.device = device
else:
self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
self.net.to(self.device)
self.net.eval()
self.timer = Timer()
def predict(self, image, top_k=-1, prob_threshold=None):
cpu_device = torch.device("cpu")
height, width, _ = image.shape
image = self.transform(image)
images = image.unsqueeze(0)
images = images.to(self.device)
with torch.no_grad():
self.timer.start()
scores, boxes = self.net.forward(images)
print("Inference time: ", self.timer.end())
boxes = boxes[0]
scores = scores[0]
if not prob_threshold:
prob_threshold = self.filter_threshold
# this version of nms is slower on GPU, so we move data to CPU.
boxes = boxes.to(cpu_device)
scores = scores.to(cpu_device)
picked_box_probs = []
picked_labels = []
for class_index in range(1, scores.size(1)):
probs = scores[:, class_index]
mask = probs > prob_threshold
probs = probs[mask]
if probs.size(0) == 0:
continue
subset_boxes = boxes[mask, :]
box_probs = torch.cat([subset_boxes, probs.reshape(-1, 1)], dim=1)
box_probs = box_utils.nms(box_probs, self.nms_method,
score_threshold=prob_threshold,
iou_threshold=self.iou_threshold,
sigma=self.sigma,
top_k=top_k,
candidate_size=self.candidate_size)
picked_box_probs.append(box_probs)
picked_labels.extend([class_index] * box_probs.size(0))
if not picked_box_probs:
return torch.tensor([]), torch.tensor([]), torch.tensor([])
picked_box_probs = torch.cat(picked_box_probs)
picked_box_probs[:, 0] *= width
picked_box_probs[:, 1] *= height
picked_box_probs[:, 2] *= width
picked_box_probs[:, 3] *= height
return picked_box_probs[:, :4], torch.tensor(picked_labels), picked_box_probs[:, 4]
