def net_initilize(self, net_type, model_path, class_names):
if net_type == 'vgg16-ssd':
net = create_vgg_ssd(len(class_names), is_test=True)
elif net_type == 'mb1-ssd':
net = create_mobilenetv1_ssd(len(class_names), is_test=True)
elif net_type == 'mb1-ssd-lite':
net = create_mobilenetv1_ssd_lite(len(class_names), is_test=True)
else:
print(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
)
sys.exit(1)
if net_type == 'vgg16-ssd':
predictor = create_vgg_ssd_predictor(net, candidate_size=200)
elif net_type == 'mb1-ssd':
predictor = create_mobilenetv1_ssd_predictor(net,
candidate_size=200)
elif net_type == 'mb1-ssd-lite':
predictor = create_mobilenetv1_ssd_lite_predictor(
net, candidate_size=200)
else:
print(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
)
sys.exit(1)
self.net = net
self.net.load(model_path)
self.pdc = predictor
def vgg16_ssd_eval_func(model, data_loader, device):
eval_path = os.path.join(os.getcwd(), ".neutrino-torch-zoo/voc/eval_results")
Path(eval_path).mkdir(parents=True, exist_ok=True)
net = create_vgg_ssd(len(VOC_CLASS_NAMES), is_test=True)
net.load_state_dict(model.state_dict())
predictor = create_vgg_ssd_predictor(net, nms_method='hard', device=device)
from eval_ssd import ssd_eval
return ssd_eval(predictor=predictor, dataset=data_loader, data_path=eval_path, class_names=VOC_CLASS_NAMES)
def select_net(model_path, net_type, num_classes):
'''
选择模型
:param model_path: 模型路径
:param net_type: 模型类型
:param num_classes: label个数，label=0,是背景
:return:
'''
if net_type == 'vgg16-ssd':
net = create_vgg_ssd(num_classes, is_test=True)
elif net_type == 'mb1-ssd':
net = create_mobilenetv1_ssd(num_classes, is_test=True)
elif net_type == 'mb1-ssd-lite':
net = create_mobilenetv1_ssd_lite(num_classes, is_test=True)
elif net_type == 'mb2-ssd-lite':
net = create_mobilenetv2_ssd_lite(num_classes,
is_test=True,
device=device)
elif net_type == 'sq-ssd-lite':
net = create_squeezenet_ssd_lite(num_classes, is_test=True)
else:
print(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
)
sys.exit(1)
net.load(model_path)
if net_type == 'vgg16-ssd':
predictor = create_vgg_ssd_predictor(net, candidate_size=200)
elif net_type == 'mb1-ssd':
predictor = create_mobilenetv1_ssd_predictor(net, candidate_size=200)
elif net_type == 'mb1-ssd-lite':
predictor = create_mobilenetv1_ssd_lite_predictor(net,
candidate_size=200)
elif net_type == 'mb2-ssd-lite':
predictor = create_mobilenetv2_ssd_lite_predictor(net,
candidate_size=200,
device=device)
elif net_type == 'sq-ssd-lite':
predictor = create_squeezenet_ssd_lite_predictor(net,
candidate_size=200)
else:
predictor = create_vgg_ssd_predictor(net, candidate_size=200)
return predictor
def create_predictor(model, model_type):
if model_type == 'vgg16-ssd':
return create_vgg_ssd_predictor(model, candidate_size=200)
elif model_type == 'mb1-ssd':
return create_mobilenetv1_ssd_predictor(model, candidate_size=200)
elif model_type == 'mb1-ssd-lite':
return create_mobilenetv1_ssd_lite_predictor(model, candidate_size=200)
elif model_type == 'mb2-ssd-lite':
return create_mobilenetv2_ssd_lite_predictor(model, candidate_size=200)
elif model_type == 'sq-ssd-lite':
return create_squeezenet_ssd_lite_predictor(model, candidate_size=200)
else:
print(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
)
sys.exit(1)
def _create_predictor(self, net_type):
if net_type == 'vgg16-ssd':
return create_vgg_ssd_predictor(self.net, candidate_size=200)
elif net_type == 'mb1-ssd':
return create_mobilenetv1_ssd_predictor(self.net,
candidate_size=200)
elif net_type == 'mb1-ssd-lite':
return create_mobilenetv1_ssd_lite_predictor(self.net,
candidate_size=200)
elif net_type == 'mb2-ssd-lite':
return create_mobilenetv2_ssd_lite_predictor(self.net,
candidate_size=200)
elif net_type == 'sq-ssd-lite':
return create_squeezenet_ssd_lite_predictor(self.net,
candidate_size=200)
else:
raise RuntimeError(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
)
def prepare_predictor(net_type, model_path, label_path):
class_names = [name.strip() for name in open(label_path).readlines()]
num_classes = len(class_names)
if net_type == 'vgg16-ssd':
net = create_vgg_ssd(num_classes, is_test=True)
elif net_type == 'mb1-ssd':
net = create_mobilenetv1_ssd(num_classes, is_test=True)
elif net_type == 'mb1-ssd-lite':
net = create_mobilenetv1_ssd_lite(num_classes, is_test=True)
elif net_type == 'mb2-ssd-lite':
net = create_mobilenetv2_ssd_lite(num_classes, is_test=True)
elif net_type == 'sq-ssd-lite':
net = create_squeezenet_ssd_lite(num_classes, is_test=True)
else:
raise ValueError(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
)
net.load(model_path)
if net_type == 'vgg16-ssd':
predictor = create_vgg_ssd_predictor(net, candidate_size=200)
elif net_type == 'mb1-ssd':
predictor = create_mobilenetv1_ssd_predictor(net, candidate_size=200)
elif net_type == 'mb1-ssd-lite':
predictor = create_mobilenetv1_ssd_lite_predictor(net,
candidate_size=200)
elif net_type == 'mb2-ssd-lite':
predictor = create_mobilenetv2_ssd_lite_predictor(net,
candidate_size=200)
elif net_type == 'sq-ssd-lite':
predictor = create_squeezenet_ssd_lite_predictor(net,
candidate_size=200)
else:
raise ValueError(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
)
return class_names, predictor
print("# Creating ssd model.")
model = create_vgg_ssd(len(class_names), is_test=True)
model.load(args["model_path"])
if args["split_point"] != 0:
_, model2 = create_vgg_ssd_new(len(class_names), model, split=args["split_point"], is_test=True)
predictor = PredictorM2(
model2,
nms_method=None,
iou_threshold=config.iou_threshold,
candidate_size=200,
device=None,
)
else:
predictor = create_vgg_ssd_predictor(model, candidate_size=200)
print("# Done!")
# If server and client run in same local directory,
# need a separate place to store the uploads.
try:
os.mkdir("uploads")
except FileExistsError:
pass
# start looping over all the frames
times = []
while True:
if args["split_point"] != 0:
(rpiName, jpg_buffer) = imageHub.recv_image()
def main(args):
DEVICE = torch.device(
"cuda:0" if torch.cuda.is_available() and args.use_cuda else "cpu")
if not os.path.exists(args.eval_dir):
os.mkdir(args.eval_dir)
timer = Timer()
class_names = [name.strip() for name in open(args.label_file).readlines()]
# dataset = Folder_image_set()
# true_case_stat, all_gb_boxes, all_difficult_cases = group_annotation_by_class(dataset)
if args.net == 'vgg16-ssd':
net = create_vgg_ssd(len(class_names), is_test=True)
elif args.net == 'mb1-ssd':
net = create_mobilenetv1_ssd(len(class_names), is_test=True)
elif args.net == 'mb1-ssd-lite':
net = create_mobilenetv1_ssd_lite(len(class_names), is_test=True)
elif args.net == 'sq-ssd-lite':
net = create_squeezenet_ssd_lite(len(class_names), is_test=True)
elif args.net == 'mb2-ssd-lite':
net = create_mobilenetv2_ssd_lite(len(class_names),
width_mult=args.mb2_width_mult,
is_test=True)
else:
logging.fatal(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
)
parser.print_help(sys.stderr)
sys.exit(1)
#train_transform = MatchPrior(config.priors, config.center_variance,
# config.size_variance, 0.5)
#test_transform = TestTransform(config.image_size, config.image_mean, config.image_std)
# test_transform = TestTransform(config.image_size, config.image_mean, config.image_std)
# dataset = FolderDataset("/media/wy_disk/wy_file/Detection/dataset/datasets/ECP_Golden_pattern", transform = test_transform)
# dataset = FolderDataset("/media/wy_disk/wy_file/Detection/dataset/datasets/ECP_Golden_pattern")
dataset = FolderDataset("/media/wy_disk/ChenYen/VIRAT/dataset_orgnize/val")
timer.start("Load Model")
net.load(args.trained_model)
net = net.to(DEVICE)
print('It took {} seconds to load the model.'.format(
timer.end("Load Model")))
if args.net == 'vgg16-ssd':
predictor = create_vgg_ssd_predictor(net,
nms_method=args.nms_method,
device=DEVICE)
elif args.net == 'mb1-ssd':
predictor = create_mobilenetv1_ssd_predictor(
net, nms_method=args.nms_method, device=DEVICE)
elif args.net == 'mb1-ssd-lite':
predictor = create_mobilenetv1_ssd_lite_predictor(
net, nms_method=args.nms_method, device=DEVICE)
elif args.net == 'sq-ssd-lite':
predictor = create_squeezenet_ssd_lite_predictor(
net, nms_method=args.nms_method, device=DEVICE)
elif args.net == 'mb2-ssd-lite':
predictor = create_mobilenetv2_ssd_lite_predictor(
net, nms_method=args.nms_method, device=DEVICE)
else:
logging.fatal(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
)
parser.print_help(sys.stderr)
sys.exit(1)
results = []
eval_path = "eval_results"
# eval_whole_image(dataset,5, predictor)
eval_subblock_image(dataset, 5, predictor)
import pdb
pdb.set_trace()
for i in range(len(dataset)):
print("process image", i)
timer.start("Load Image")
import pdb
pdb.set_trace()
image = dataset.get_image(i)
print("Load Image: {:4f} seconds.".format(timer.end("Load Image")))
timer.start("Predict")
boxes, labels, probs = predictor.predict(image, 10, 0.5)
print("Prediction: {:4f} seconds.".format(timer.end("Predict")))
# indexes = torch.ones(labels.size(0), 1, dtype=torch.float32) * i
# print("index:p\t{}".format(sum(probs>0.5)))
# import pdb;pdb.set_trace()
boxes, labels, probs = boxes.data.numpy(), labels.data.numpy(
), probs.data.numpy()
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
for box, _label, _prob in zip(boxes, labels, probs):
if _prob < 0.7: continue
print(box)
box = box.astype(int)
# import pdb;pdb.set_trace()
print(box)
cv2.rectangle(image, (box[0], box[1]), (box[2], box[3]),
(255, 255, 0), 4)
# str(str.split(class_names[_label]," ")[1])
cv2.putText(
image,
dataset.class_names[_label],
(box[0] + 20, box[1] + 40),
cv2.FONT_HERSHEY_SIMPLEX,
1, # font scale
(255, 0, 255),
2) # line type
print(boxes.shape[0])
cv2.imshow('annotated', image)
# key = cv2.waitKey(0)
if cv2.waitKey(0) & 0xFF == ord('q'):
break
elif net_type == 'mb1-ssd-lite':
net = create_mobilenetv1_ssd_lite(len(class_names), is_test=True)
elif net_type == 'mb2-ssd-lite':
net = create_mobilenetv2_ssd_lite(len(class_names), is_test=True)
elif net_type == 'sq-ssd-lite':
net = create_squeezenet_ssd_lite(len(class_names), is_test=True)
else:
print(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
)
sys.exit(1)
net.load(model_path)
if net_type == 'vgg16-ssd':
predictor = create_vgg_ssd_predictor(net,
candidate_size=200,
filter_threshold=0.01)
elif net_type == 'mb1-ssd':
predictor = create_mobilenetv1_ssd_predictor(net, candidate_size=200)
elif net_type == 'mb1-ssd-lite':
predictor = create_mobilenetv1_ssd_lite_predictor(net, candidate_size=200)
elif net_type == 'mb2-ssd-lite':
predictor = create_mobilenetv2_ssd_lite_predictor(net, candidate_size=200)
elif net_type == 'sq-ssd-lite':
predictor = create_squeezenet_ssd_lite_predictor(net, candidate_size=200)
else:
print(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
)
sys.exit(1)
elif args.net == 'mb3-small-ssd-lite':
net = create_mobilenetv3_small_ssd_lite(len(class_names), is_test=True)
else:
logging.fatal(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
)
parser.print_help(sys.stderr)
sys.exit(1)
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.')
if args.net == 'vgg16-ssd':
predictor = create_vgg_ssd_predictor(net,
nms_method=args.nms_method,
iou_threshold=0.5,
device=DEVICE)
elif args.net == 'vgg16-ssd1':
predictor = create_vgg_ssd_predictor1(net,
nms_method=args.nms_method,
iou_threshold=0.5,
device=DEVICE)
elif args.net == 'vgg16-ssd2':
predictor = create_vgg_ssd_predictor2(net,
nms_method=args.nms_method,
iou_threshold=0.2,
device=DEVICE)
elif args.net == 'mb1-ssd':
predictor = create_mobilenetv1_ssd_predictor(
net, nms_method=args.nms_method, iou_threshold=0.5, device=DEVICE)
elif args.net == 'mb1-ssd-lite':
help='Specify dataset type. Currently support voc and open_images.')
parser.add_argument("--label_path",default="./models/voc-model-labels.txt", 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('--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")
class_names = [name.strip() for name in open(args.label_path).readlines()]
num_classes = len(class_names)
if args.net == 'vgg16-ssd':
net = create_vgg_ssd(num_classes, is_test=True)
net.load(args.model_path)
predictor = create_vgg_ssd_predictor(net, candidate_size=200, device=torch.device("cpu"))
elif args.net == "mb2-ssd-lite":
net = create_mobilenetv2_ssd_lite(num_classes, is_test=True)
net.load(args.model_path)
predictor = create_mobilenetv2_ssd_lite_predictor(net, candidate_size=200,
device=torch.device('cuda'))
timer = Timer()
cap = cv2.VideoCapture("./test_video/test_video_traffic.mp4") # capture from file
#fourcc = cv2.VideoWriter_fourcc(*'XVID')
#video = cv2.VideoWriter("./mobilenet_v2.avi", fourcc, 30.0, (534, 300))
#start = time.time()
while True:
def main(args):
net_type = args.net_type
model_path = args.weights_path
label_path = args.label_path
class_names = [name.strip() for name in open(label_path).readlines()]
num_classes = len(class_names)
if args.live:
cap = cv2.VideoCapture(0)
cap.set(3, 640)
cap.set(4, 480)
else:
cap = cv2.VideoCapture(args.video_path)
Fourcc = cv2.VideoWriter_fourcc('M', 'P', '4', 'V')
writer = cv2.VideoWriter('result.mp4',
fourcc=Fourcc,
fps=15,
frameSize=(640, 480))
num_gpus = torch.cuda.device_count()
device = 'cuda' if num_gpus else 'cpu'
if net_type == 'vgg16-ssd':
net = create_vgg_ssd(len(class_names), is_test=True)
elif net_type == 'mb1-ssd':
net = create_mobilenetv1_ssd(len(class_names), is_test=True)
elif net_type == 'mb1-ssd-lite':
net = create_mobilenetv1_ssd_lite(len(class_names), is_test=True)
elif net_type == 'mb2-ssd-lite':
net = create_mobilenetv2_ssd_lite(len(class_names), is_test=True)
elif net_type == 'sq-ssd-lite':
net = create_squeezenet_ssd_lite(len(class_names), is_test=True)
#elif net_type == 'mb3-ssd-lite':
# net = create_mobilenetv3_ssd_lite(len(class_names), is_test=True)
else:
print(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
)
sys.exit(1)
net.load(model_path)
if net_type == 'vgg16-ssd':
predictor = create_vgg_ssd_predictor(net,
candidate_size=20,
device=device)
elif net_type == 'mb1-ssd':
predictor = create_mobilenetv1_ssd_predictor(net,
candidate_size=20,
device=device)
elif net_type == 'mb1-ssd-lite':
predictor = create_mobilenetv1_ssd_lite_predictor(net,
candidate_size=20,
device=device)
elif net_type == 'mb2-ssd-lite':
predictor = create_mobilenetv2_ssd_lite_predictor(net,
candidate_size=20,
device=device)
elif net_type == 'sq-ssd-lite':
predictor = create_squeezenet_ssd_lite_predictor(net,
candidate_size=20,
device=device)
#elif net_type == 'mb3-ssd-lite':
# predictor = create_mobilenetv3_ssd_lite_predictor(net, candidate_size=10)
else:
print(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
)
sys.exit(1)
timer = Timer()
while True:
_, orig_image = cap.read()
if orig_image is None:
print('END')
break
image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)
timer.start()
boxes, labels, probs = predictor.predict(image, 10, 0.4)
interval = timer.end()
print('Time: {:.2f}s, Detect Objects: {:d}.'.format(
interval, labels.size(0)))
for i in range(boxes.size(0)):
box = boxes[i, :]
label = f"{class_names[labels[i]]}: {probs[i]:.2f}"
cv2.rectangle(orig_image, (box[0], box[1]), (box[2], box[3]),
(255, 255, 0), 4)
cv2.putText(
orig_image,
label,
(box[0] + 20, box[1] + 40),
cv2.FONT_HERSHEY_SIMPLEX,
1, # font scale
(255, 0, 255),
2) # line type
writer.write(orig_image)
cv2.imshow('annotated', orig_image)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
writer.release()
cv2.destroyAllWindows()
if args.out_video:
shutil.move('result.mp4', args.out_video)
else:
os.remove('result.mp4')
def main(args):
DEVICE = torch.device("cuda:0" if torch.cuda.is_available()
and args['flow_control']['use_cuda'] else "cpu")
# eval_path = pathlib.Path(args.eval_dir)
# eval_path.mkdir(exist_ok=True)
if not os.path.exists(args['flow_control']['eval_dir']):
os.mkdir(args['flow_control']['eval_dir'])
timer = Timer()
class_names = [
name.strip()
for name in open(args['flow_control']['label_file']).readlines()
]
_net = args['flow_control']['net']
_dataset_type = args['flow_control']['dataset_type']
if _dataset_type == "voc":
raise NotImplementedError("Not implement error")
dataset = VOCDataset(args['flow_control']['dataset'], is_test=True)
elif _dataset_type == 'open_images':
raise NotImplementedError("Not implement error")
dataset = OpenImagesDataset(args['flow_control']['dataset'],
dataset_type="test")
elif _dataset_type == "coco":
# dataset = CocoDetection("/home/wenyen4desh/datasets/coco/test2017","/home/wenyen4desh/datasets/annotations/image_info_test2017.json")
#dataset = CocoDetection("../../dataset/datasets/coco/val2017","../../dataset/datasets/coco/annotations/instances_val2017.json")
# dataset = CocoDetection("/home/wenyen4desh/datasets/coco/train2017","/home/wenyen4desh/datasets/coco/annotations/instances_train2017.json")
dataset = CocoDetection(args['Datasets']['coco']['val_image_path'],
args['Datasets']['coco']['val_anno_path'])
elif _dataset_type == "ecp":
dataset = EuroCity_Dataset(args['Datasets']['ecp']['val_image_path'],
args['Datasets']['ecp']['val_anno_path'])
true_case_stat, all_gb_boxes, all_difficult_cases = group_annotation_by_class(
dataset)
if _net == 'vgg16-ssd':
net = create_vgg_ssd(len(class_names), is_test=True)
elif _net == 'mb1-ssd':
net = create_mobilenetv1_ssd(len(class_names), is_test=True)
elif _net == 'mb1-ssd-lite':
net = create_mobilenetv1_ssd_lite(len(class_names), is_test=True)
elif _net == 'sq-ssd-lite':
net = create_squeezenet_ssd_lite(len(class_names), is_test=True)
elif _net == 'mb2-ssd-lite':
net = create_mobilenetv2_ssd_lite(
len(class_names),
width_mult=args['flow_control']['mb2_width_mult'],
is_test=True)
else:
logging.fatal(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
)
parser.print_help(sys.stderr)
sys.exit(1)
#train_transform = MatchPrior(config.priors, config.center_variance,
# config.size_variance, 0.5)
#test_transform = TestTransform(config.image_size, config.image_mean, config.image_std)
import pdb
pdb.set_trace()
############################## automatically validation ############################################
timer.start("Load Model")
net.load(args['flow_control']['trained_model'])
net = net.to(DEVICE)
print('It took {} seconds to load the model.'.format(
timer.end("Load Model")))
_nms_method = args['flow_control']['nms_method']
if _net == 'vgg16-ssd':
predictor = create_vgg_ssd_predictor(net,
nms_method=_nms_method,
device=DEVICE)
elif _net == 'mb1-ssd':
predictor = create_mobilenetv1_ssd_predictor(net,
nms_method=_nms_method,
device=DEVICE)
elif _net == 'mb1-ssd-lite':
predictor = create_mobilenetv1_ssd_lite_predictor(
net, nms_method=_nms_method, device=DEVICE)
elif _net == 'sq-ssd-lite':
predictor = create_squeezenet_ssd_lite_predictor(
net, nms_method=_nms_method, device=DEVICE)
elif _net == 'mb2-ssd-lite':
predictor = create_mobilenetv2_ssd_lite_predictor(
net, nms_method=_nms_method, device=DEVICE)
else:
logging.fatal(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
)
parser.print_help(sys.stderr)
sys.exit(1)
results = []
# Predict Bounding Box
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)
# Write the result to file
for class_index, class_name in enumerate(class_names):
if class_index == 0: continue # ignore background
file_name = "det_test_{}.txt".format(class_name)
prediction_path = os.path.join(args['flow_control']['eval_dir'],
file_name)
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.get_annotation(int(sub[i, 0]))
f.write(
str(image_id) + " " + " ".join([str(v)
for v in prob_box]) + "\n")
# image_id = dataset.ids[int(sub[i, 0])]
# print(str(image_id) + " " + " ".join([str(v) for v in prob_box]), file=f)
aps = []
prcs = []
recalls = []
print("\n\nAverage Precision Per-class:")
for class_index, class_name in enumerate(class_names):
if class_index == 0:
continue
file_name = "det_test_{}.txt".format(class_name)
prediction_path = os.path.join(args['flow_control']['eval_dir'],
file_name)
# [email protected] evaluation method
ap, precision, recall = compute_average_precision_per_class(
args, true_case_stat[class_index], all_gb_boxes[class_index],
all_difficult_cases[class_index], prediction_path,
args['flow_control']['iou_threshold'],
args['flow_control']['use_2007_metric'])
# # COCO eval
# ap, precision, recall = coco_ap_per_class(
# true_case_stat[class_index],
# all_gb_boxes[class_index],
# all_difficult_cases[class_index],
# prediction_path,
# args.use_2007_metric
# )
aps.append(ap)
prcs.append(precision)
recalls.append(recall)
print("{}: {}".format(class_name, ap))
print("\nAverage Precision Across All Classes:{}".format(
sum(aps[0:5]) / len(aps[0:5])))
print("\nAverage Precision :{}".format(sum(prcs[0:5]) / len(prcs[0:5])))
print("\nAverage Recall :{}".format(sum(recalls[0:5]) / len(recalls[0:5])))
pretrained_dict = {
k: v
for k, v in pretrained_dict.items() if k in net_dict
}
# # 2. overwrite
net_dict.update(pretrained_dict)
# # 3. load
net.load_state_dict(net_dict)
else:
net.load(args.trained_model)
net = net.to(DEVICE)
print(f'It took {timer.end("Load Model")} seconds to load the model.')
if args.net == 'vgg16-ssd':
predictor = create_vgg_ssd_predictor(
net, nms_method=args.nms_method, device=DEVICE,
config=config) # iou_theresh : nmsのしきい値
elif args.net == 'resnet50-ssd':
predictor = create_resnet50_ssd_predictor(net,
nms_method=args.nms_method,
device=DEVICE,
config=config)
# elif args.net == 'mb1-ssd':
# predictor = create_mobilenetv1_ssd_predictor(net, nms_method=args.nms_method, device=DEVICE)
# elif args.net == 'mb1-ssd-lite':
# predictor = create_mobilenetv1_ssd_lite_predictor(net, nms_method=args.nms_method, device=DEVICE)
# elif args.net == 'sq-ssd-lite':
# predictor = create_squeezenet_ssd_lite_predictor(net,nms_method=args.nms_method, device=DEVICE)
# elif args.net == 'mb2-ssd-lite':
# predictor = create_mobilenetv2_ssd_lite_predictor(net, nms_method=args.nms_method, device=DEVICE)
else:
elif net_type == 'mb2-ssd-lite-xiaomi':
net = create_mobilenetv2_ssd_lite_xiaomi(num_classes, is_test=True)
elif net_type == 'sq-ssd-lite':
net = create_squeezenet_ssd_lite(num_classes, is_test=True)
elif net_type == 'fairnas-a':
net = create_fairnas_a_ssd_lite(num_classes, is_test=True)
elif net_type == 'fairnas-b':
net = create_fairnas_b_ssd_lite(num_classes, is_test=True)
else:
print("The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite.")
sys.exit(1)
net.load(model_path)
net.to(DEVICE)
if net_type == 'vgg16-ssd':
predictor = create_vgg_ssd_predictor(net, candidate_size=100, device=DEVICE)
elif net_type == 'mb1-ssd':
predictor = create_mobilenetv1_ssd_predictor(net, candidate_size=100, device=DEVICE)
elif net_type == 'mb1-ssd-lite':
predictor = create_mobilenetv1_ssd_lite_predictor(net, candidate_size=100, device=DEVICE)
elif net_type == 'mb2-ssd-lite':
predictor = create_mobilenetv2_ssd_lite_predictor(net, candidate_size=100, device=DEVICE)
elif net_type == 'mb2-ssd-lite-xiaomi':
predictor = create_mobilenetv2_ssd_lite_predictor_xiaomi(net, candidate_size=100, device=DEVICE)
elif net_type == 'sq-ssd-lite':
predictor = create_squeezenet_ssd_lite_predictor(net, candidate_size=100, device=DEVICE)
elif net_type == 'fairnas-a':
predictor = create_fairnas_a_ssd_lite_predictor(net, candidate_size=100, device=DEVICE)
elif net_type == 'fairnas-b':
predictor = create_fairnas_b_ssd_lite_predictor(net, candidate_size=100, device=DEVICE)
else:
elif net_type == 'mb1-ssd-lite':
net = create_mobilenetv1_ssd_lite(len(class_names), is_test=True)
elif net_type == 'mb2-ssd-lite':
net = create_mobilenetv2_ssd_lite(len(class_names), is_test=True)
elif net_type == 'sq-ssd-lite':
net = create_squeezenet_ssd_lite(len(class_names), is_test=True)
else:
print(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
)
sys.exit(1)
net.load(model_path)
if net_type == 'vgg16-ssd':
predictor = create_vgg_ssd_predictor(net,
candidate_size=200,
device=device)
elif net_type == 'mb1-ssd':
predictor = create_mobilenetv1_ssd_predictor(net,
candidate_size=200,
device=device)
elif net_type == 'mb1-ssd-lite':
predictor = create_mobilenetv1_ssd_lite_predictor(net,
candidate_size=200,
device=device)
elif net_type == 'mb2-ssd-lite':
predictor = create_mobilenetv2_ssd_lite_predictor(net,
candidate_size=200,
device=device)
elif net_type == 'sq-ssd-lite':
predictor = create_squeezenet_ssd_lite_predictor(net,
def load_model(self,
net="mb1-ssd",
model_path=None,
label_path=None,
use_gpu=True):
net_type = net
self.system_dict["device"] = torch.device(
'cuda' if torch.cuda.is_available() and use_gpu else 'cpu')
self.system_dict["class_names"] = [
name.strip() for name in open(label_path).readlines()
]
self.system_dict["class_names"].insert(0, 'BACKGROUND')
if net_type == 'vgg16-ssd':
self.system_dict["net"] = create_vgg_ssd(
len(self.system_dict["class_names"]),
is_test=True,
device=self.system_dict["device"])
elif net_type == 'mb1-ssd':
self.system_dict["net"] = create_mobilenetv1_ssd(
len(self.system_dict["class_names"]),
is_test=True,
device=self.system_dict["device"])
elif net_type == 'mb1-ssd-lite':
self.system_dict["net"] = create_mobilenetv1_ssd_lite(
len(self.system_dict["class_names"]),
is_test=True,
device=self.system_dict["device"])
elif net_type == 'mb2-ssd-lite':
self.system_dict["net"] = create_mobilenetv2_ssd_lite(
len(self.system_dict["class_names"]),
is_test=True,
device=self.system_dict["device"])
elif net_type == 'sq-ssd-lite':
self.system_dict["net"] = create_squeezenet_ssd_lite(
len(self.system_dict["class_names"]),
is_test=True,
device=self.system_dict["device"])
else:
print(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
)
sys.exit(1)
self.system_dict["net"].load(model_path)
if net_type == 'vgg16-ssd':
self.system_dict["predictor"] = create_vgg_ssd_predictor(
self.system_dict["net"],
candidate_size=200,
device=self.system_dict["device"])
elif net_type == 'mb1-ssd':
self.system_dict["predictor"] = create_mobilenetv1_ssd_predictor(
self.system_dict["net"],
candidate_size=200,
device=self.system_dict["device"])
elif net_type == 'mb1-ssd-lite':
self.system_dict[
"predictor"] = create_mobilenetv1_ssd_lite_predictor(
self.system_dict["net"],
candidate_size=200,
device=self.system_dict["device"])
elif net_type == 'mb2-ssd-lite':
self.system_dict[
"predictor"] = create_mobilenetv2_ssd_lite_predictor(
self.system_dict["net"],
candidate_size=200,
device=self.system_dict["device"])
elif net_type == 'sq-ssd-lite':
self.system_dict[
"predictor"] = create_squeezenet_ssd_lite_predictor(
self.system_dict["net"],
candidate_size=200,
device=self.system_dict["device"])
else:
self.system_dict["predictor"] = create_vgg_ssd_predictor(
self.system_dict["net"],
candidate_size=200,
device=self.system_dict["device"])
def main():
eval_path = pathlib.Path(args.eval_dir)
eval_path.mkdir(exist_ok=True)
timer = Timer()
class_names = {
'BACKGROUND', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus',
'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse',
'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train',
'tvmonitor'
}
if args.dataset_type == "voc":
dataset = VOCDataset(args.dataset, is_test=True)
elif args.dataset_type == 'open_images':
dataset = OpenImagesDataset(args.dataset, dataset_type="test")
# true_case_stat, all_gb_boxes, all_difficult_cases = group_annotation_by_class(dataset)
if args.net == 'vgg16-ssd':
net = create_vgg_ssd(len(class_names), is_test=True)
elif args.net == 'mb1-ssd':
net = create_mobilenetv1_ssd(len(class_names), is_test=True)
elif args.net == 'mb1-ssd-lite':
net = create_mobilenetv1_ssd_lite(len(class_names), is_test=True)
elif args.net == 'sq-ssd-lite':
net = create_squeezenet_ssd_lite(len(class_names), is_test=True)
elif args.net == 'mb2-ssd-lite':
net = create_mobilenetv2_ssd_lite(len(class_names),
width_mult=args.mb2_width_mult,
is_test=True)
else:
logging.fatal(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
)
parser.print_help(sys.stderr)
sys.exit(1)
timer.start("Load Model")
net.load(args.trained_model)
net = net.to(DEVICE)
print('It took {timer.end("Load Model")} seconds to load the model.')
if args.net == 'vgg16-ssd':
predictor = create_vgg_ssd_predictor(net,
nms_method=args.nms_method,
device=DEVICE)
elif args.net == 'mb1-ssd':
predictor = create_mobilenetv1_ssd_predictor(
net, nms_method=args.nms_method, device=DEVICE)
elif args.net == 'mb1-ssd-lite':
predictor = create_mobilenetv1_ssd_lite_predictor(
net, nms_method=args.nms_method, device=DEVICE)
elif args.net == 'sq-ssd-lite':
predictor = create_squeezenet_ssd_lite_predictor(
net, nms_method=args.nms_method, device=DEVICE)
elif args.net == 'mb2-ssd-lite':
predictor = create_mobilenetv2_ssd_lite_predictor(
net, nms_method=args.nms_method, device=DEVICE)
else:
logging.fatal(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
)
parser.print_help(sys.stderr)
sys.exit(1)
print(
ssd_eval(predictor=predictor,
dataset=dataset,
data_path=eval_path,
class_names=class_names))
elif args.net == 'mb2-ssd-lite':
net = create_mobilenetv2_ssd_lite(len(class_names),
width_mult=args.mb2_width_mult,
is_test=True)
else:
logging.fatal(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
)
parser.print_help(sys.stderr)
sys.exit(1)
timer.start("Load Model")
net.ssd.load_weights(args.trained_model, by_name=True)
print(f'It took {timer.end("Load Model")} seconds to load the model.')
if args.net == 'vgg16-ssd':
predictor = create_vgg_ssd_predictor(net, nms_method=args.nms_method)
elif args.net == 'mb1-ssd':
predictor = create_mobilenetv1_ssd_predictor(
net, nms_method=args.nms_method)
elif args.net == 'mb1-ssd-lite':
predictor = create_mobilenetv1_ssd_lite_predictor(
net, nms_method=args.nms_method)
elif args.net == 'sq-ssd-lite':
predictor = create_squeezenet_ssd_lite_predictor(
net, nms_method=args.nms_method)
elif args.net == 'mb2-ssd-lite':
predictor = create_mobilenetv2_ssd_lite_predictor(
net, nms_method=args.nms_method)
else:
logging.fatal(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
def main(args):
net_type = args.net_type
img_folder = args.img_folder
model_path = args.weights_path
label_path = args.label_path
class_names = [name.strip() for name in open(label_path).readlines()]
out_path = args.out_path
if not os.path.exists(out_path):
os.mkdir(out_path)
num_gpus = torch.cuda.device_count()
device = 'cuda' if num_gpus else 'cpu'
if net_type == 'vgg16-ssd':
net = create_vgg_ssd(len(class_names), is_test=True)
elif net_type == 'mb1-ssd':
net = create_mobilenetv1_ssd(len(class_names), is_test=True)
elif net_type == 'mb1-ssd-lite':
net = create_mobilenetv1_ssd_lite(len(class_names), is_test=True)
elif net_type == 'mb2-ssd-lite':
net = create_mobilenetv2_ssd_lite(len(class_names), is_test=True)
elif net_type == 'sq-ssd-lite':
net = create_squeezenet_ssd_lite(len(class_names), is_test=True)
#elif net_type == 'mb3-ssd-lite':
# net = create_mobilenetv3_ssd_lite(len(class_names), is_test=True)
else:
print(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
)
sys.exit(1)
net.load(model_path)
if net_type == 'vgg16-ssd':
predictor = create_vgg_ssd_predictor(net,
candidate_size=20,
device=device)
elif net_type == 'mb1-ssd':
predictor = create_mobilenetv1_ssd_predictor(net,
candidate_size=20,
device=device)
elif net_type == 'mb1-ssd-lite':
predictor = create_mobilenetv1_ssd_lite_predictor(net,
candidate_size=20,
device=device)
elif net_type == 'mb2-ssd-lite':
predictor = create_mobilenetv2_ssd_lite_predictor(net,
candidate_size=20,
device=device)
elif net_type == 'sq-ssd-lite':
predictor = create_squeezenet_ssd_lite_predictor(net,
candidate_size=20,
device=device)
#elif net_type == 'mb3-ssd-lite':
# predictor = create_mobilenetv3_ssd_lite_predictor(net, candidate_size=10)
else:
print(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
)
sys.exit(1)
timer = Timer()
img_names = glob.glob(img_folder + os.sep + '*.jpg')
#result_csv=os.path.join(out_path,'rest_result.csv')
if len(img_names) == 0:
print('No imagesfound in {}'.format(img_folder))
exit(-1)
for img_name in img_names:
image = cv2.imread(img_name)
timer.start()
boxes, labels, probs = predictor.predict(image, 10, 0.3)
interval = timer.end()
print('Time: {:.2f}s, Detect Objects: {:d}.'.format(
interval, labels.size(0)))
label_text = []
for i in range(boxes.size(0)):
box = boxes[i, :]
label = f"{class_names[labels[i]]}: {probs[i]:.2f}"
label_text.append(label)
cv2.rectangle(image, (box[0], box[1]), (box[2], box[3]),
(255, 255, 0), 4)
cv2.putText(image, label, (box[0] + 20, box[1] + 40),
cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 255), 2)
if args.store_result:
new_name = '{}/{}'.format(out_path, img_name.split('/')[-1])
cv2.imwrite(new_name, image)
if not label_text:
result_label = 'empty'
else:
result_label = label_text[0]
with open(os.path.join(out_path, 'rest_result.csv'),
'a+') as result_writer:
result_writer.write(
img_name.split('/')[-1] + ',' + result_label + '\n')
cv2.imshow('result', image)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cv2.destroyAllWindows()
def imwrite(dataset, net_type, epoch, model_path):
num_classes = len(dataset.class_names)
if net_type == 'vgg16-ssd':
net = create_vgg_ssd(num_classes, is_test=True)
elif net_type == 'mb1-ssd':
net = create_mobilenetv1_ssd(num_classes, is_test=True)
elif net_type == 'mb1-ssd-lite':
net = create_mobilenetv1_ssd_lite(num_classes, is_test=True)
elif net_type == 'mb2-ssd-lite':
net = create_mobilenetv2_ssd_lite(num_classes, is_test=True)
elif net_type == 'mb3-large-ssd-lite':
net = create_mobilenetv3_large_ssd_lite(num_classes, is_test=True)
elif net_type == 'mb3-small-ssd-lite':
net = create_mobilenetv3_small_ssd_lite(num_classes, is_test=True)
elif net_type == 'sq-ssd-lite':
net = create_squeezenet_ssd_lite(num_classes, is_test=True)
else:
print(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
)
sys.exit(1)
net.load(model_path)
if net_type == 'vgg16-ssd':
predictor = create_vgg_ssd_predictor(net, candidate_size=200)
elif net_type == 'mb1-ssd':
predictor = create_mobilenetv1_ssd_predictor(net, candidate_size=200)
elif net_type == 'mb1-ssd-lite':
predictor = create_mobilenetv1_ssd_lite_predictor(net,
candidate_size=200)
elif net_type == 'mb2-ssd-lite' or net_type == "mb3-large-ssd-lite" or net_type == "mb3-small-ssd-lite":
predictor = create_mobilenetv2_ssd_lite_predictor(net,
candidate_size=200)
elif net_type == 'sq-ssd-lite':
predictor = create_squeezenet_ssd_lite_predictor(net,
candidate_size=200)
else:
predictor = create_vgg_ssd_predictor(net, candidate_size=200)
for i in range(10):
image, orig_boxes, labels = dataset[i]
orig_image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
boxes, labels, probs = predictor.predict(image, 10, 0.4)
for j in range(boxes.size(0)): # predict
box = boxes[j, :]
box = [int(i) for i in box]
cv2.rectangle(orig_image, (box[0], box[1]), (box[2], box[3]),
(255, 255, 0), 1)
label = f"{probs[j]:.2f}"
cv2.putText(
orig_image,
label,
(box[0] + 3, box[1] + 5),
cv2.FONT_HERSHEY_SIMPLEX,
1, # font scale
(255, 255, 0),
1) # line type
for j in range(orig_boxes.shape[0]): # ground truth
box = orig_boxes[j, :]
cv2.rectangle(orig_image, (box[0], box[1]), (box[2], box[3]),
(255, 0, 255), 1)
path = f"out/{i:02}_{epoch:04}.jpg"
cv2.imwrite(path, orig_image)
net = create_mobilenetv2_ssd_lite(len(class_names), is_test=True)
elif net_type == 'mb3-large-ssd-lite':
net = create_mobilenetv3_large_ssd_lite(len(class_names), is_test=True)
elif net_type == 'mb3-small-ssd-lite':
net = create_mobilenetv3_small_ssd_lite(len(class_names), is_test=True)
elif net_type == 'sq-ssd-lite':
net = create_squeezenet_ssd_lite(len(class_names), is_test=True)
else:
print(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
)
sys.exit(1)
net.load(model_path)
if net_type == 'vgg16-ssd':
predictor = create_vgg_ssd_predictor(net, candidate_size=200)
elif net_type == 'mb1-ssd':
predictor = create_mobilenetv1_ssd_predictor(net, candidate_size=200)
elif net_type == 'mb1-ssd-lite':
predictor = create_mobilenetv1_ssd_lite_predictor(net, candidate_size=200)
elif net_type == 'mb2-ssd-lite' or net_type == "mb3-large-ssd-lite" or net_type == "mb3-small-ssd-lite":
predictor = create_mobilenetv2_ssd_lite_predictor(net, candidate_size=200)
elif net_type == 'sq-ssd-lite':
predictor = create_squeezenet_ssd_lite_predictor(net, candidate_size=200)
else:
predictor = create_vgg_ssd_predictor(net, candidate_size=200)
orig_image = cv2.imread(image_path)
image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)
boxes, labels, probs = predictor.predict(image, 10, 0.4)
def PersonDetector(orig_image, net_type="mb1-ssd"):
class_names = [name.strip() for name in open(label_path).readlines()]
num_classes = len(class_names)
if net_type == 'vgg16-ssd':
net = create_vgg_ssd(len(class_names), is_test=True)
elif net_type == 'mb1-ssd':
net = create_mobilenetv1_ssd(len(class_names), is_test=True)
elif net_type == 'mb1-ssd-lite':
net = create_mobilenetv1_ssd_lite(len(class_names), is_test=True)
elif net_type == 'mb2-ssd-lite':
net = create_mobilenetv2_ssd_lite(len(class_names), is_test=True)
elif net_type == 'sq-ssd-lite':
net = create_squeezenet_ssd_lite(len(class_names), is_test=True)
else:
print(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
)
sys.exit(1)
net.load(model_path)
if net_type == 'vgg16-ssd':
predictor = create_vgg_ssd_predictor(net, candidate_size=200)
elif net_type == 'mb1-ssd':
predictor = create_mobilenetv1_ssd_predictor(net, candidate_size=200)
elif net_type == 'mb1-ssd-lite':
predictor = create_mobilenetv1_ssd_lite_predictor(net,
candidate_size=200)
elif net_type == 'mb2-ssd-lite':
predictor = create_mobilenetv2_ssd_lite_predictor(net,
candidate_size=200)
elif net_type == 'sq-ssd-lite':
predictor = create_squeezenet_ssd_lite_predictor(net,
candidate_size=200)
else:
print(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
)
sys.exit(1)
timer = Timer()
image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)
timer.start()
boxes, labels, probs = predictor.predict(image, 10, 0.4)
interval = timer.end()
print('Time: {:.2f}s, Detect Objects: {:d}.'.format(
interval, labels.size(0)))
max_width = -1
x, y, w, h = None, None, None, None
for i in range(boxes.size(0)):
box = boxes[i, :]
label = f"{class_names[labels[i]]}: {probs[i]:.2f}"
if (max_width < box[2] - box[0]):
x, y = box[0], box[1]
w, h = box[2] - box[0], box[3] - box[1]
max_width = w
if (x is not None and y is not None and w is not None and h is not None):
cv2.rectangle(orig_image, (x, y), (w + x, h + y), (255, 255, 0), 4)
return (x, y, w, h)
# img = cv2.imread("Image/img.jpg")
# PersonDetector(img)
is_test=True)
else:
logging.fatal(
"The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite."
)
parser.print_help(sys.stderr)
sys.exit(1)
timer.start("Load Model")
# net.load(args.trained_model)
net.load_checkpoint(args.trained_model)
net = net.to(DEVICE)
print(f'It took {timer.end("Load Model")} seconds to load the model.')
if args.net == 'vgg16-ssd':
predictor = create_vgg_ssd_predictor(net,
nms_method=args.nms_method,
device=DEVICE)
elif args.net == 'mb1-ssd':
predictor = create_mobilenetv1_ssd_predictor(
net, nms_method=args.nms_method, device=DEVICE)
elif args.net == 'mb1-ssd-lite':
predictor = create_mobilenetv1_ssd_lite_predictor(
net, nms_method=args.nms_method, device=DEVICE)
elif args.net == 'sq-ssd-lite':
predictor = create_squeezenet_ssd_lite_predictor(
net, nms_method=args.nms_method, device=DEVICE)
elif args.net == 'mb2-ssd-lite':
predictor = create_mobilenetv2_ssd_lite_predictor(
net, nms_method=args.nms_method, device=DEVICE)
else:
logging.fatal(
def ssd_vgg(batch_size):
output_f = open('../batch_test.txt', 'w')
output_f.close()
class_names = [
'gun', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car',
'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike',
'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor'
]
model_path = 'models/vgg16-ssd-mp-0_7726.pth'
image_path = 'gun.jpg'
net = create_vgg_ssd(len(class_names), is_test=True)
# print(net)
net.load(model_path)
predictor = create_vgg_ssd_predictor(net, candidate_size=200)
device = torch.cuda.current_device()
print(device)
# Classification
# loaded_model = models.alexnet(pretrained=True)
loaded_model = models.vgg16(pretrained=True)
loaded_model.cuda()
orig_image = cv2.imread(image_path)
ssd_image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)
print('before first pred trans')
trans = PredictionTransform(300, [123, 117, 104], 1.0)
print('after first pred trans')
# print(ssd_trans)
height, width, _ = ssd_image.shape
# print(cur_image.shape)
# print(type(ssd_image))
# print('before trans')
vgg_image = trans(ssd_image)
print(vgg_image.size())
# print('After trans')
# print(type(image))
images = vgg_image.unsqueeze(0)
print(images.size())
# images = images.to(device)
print('before loop')
batch_idx = 0
total_time = 0.0
for i in range(loop_num):
output_f = open('../batch_test.txt', 'a')
output_f.write('New Loop\n')
output_f.flush()
output_f.close()
# print(vgg_image.size())
images = Variable(images).cuda()
# labels = Variable(labels).cuda()
tmp_images = images.clone()
# images = torch.cat((images, ), 0)
# print(images.size())
print(
'-------------------------------------images=-----------------------'
)
print(images.size())
print(images[0][0][0])
t0 = cal_time()
boxes, labels, probs = predictor.predict(ssd_image, batch_size)
t1 = cal_time()
# print(Variable(test_loader[0]).cuda().size())
# print(images.size())
# out = loaded_model(images)
batch_idx += 1
if batch_idx > 5:
total_time += (t1 - t0)
output_f = open('../batch_test.txt', 'a')
output_f.write('Execution time: ' + str(t1 - t0) + '\n')
# print('Execution time: ' + str(t1 - t0))
output_f.flush()
output_f.close()
# if batch_idx > 1:
# break
break
output_f = open('../batch_test.txt', 'a')
output_f.write('Avg time: ' + str(total_time / 25.0) + '\n')
# print('Avg time: ' + str(total_time/25.0))
output_f.flush()
output_f.close()
