def Init_Net(self, net_type):
print("==========Init Net==========")
if net_type=="mb2-ssd":
model_path = "./models/mobilenet2-ssd.pth"
label_path = "./models/voc-model-labels_mb2.txt"
self.class_names = [name.strip() for name in open(label_path).readlines()]
net = create_mobilenetv2_ssd_lite(len(self.class_names), is_test=True)
net.load(model_path)
self.predictor = create_mobilenetv2_ssd_lite_predictor(net, candidate_size=200)
elif net_type=="mb3-large-ssd":
model_path = "./models/mobilenet3-large-ssd.pth"
label_path = "./models/voc-model-labels_mb3.txt"
self.class_names = [name.strip() for name in open(label_path).readlines()]
net = create_mobilenetv3_ssd_lite("Large", len(self.class_names), is_test=True)
net.load(model_path)
self.predictor = create_mobilenetv3_ssd_lite_predictor(net, candidate_size=200)
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_names = [name.strip() for name in open(label_path).readlines()]
if ~PRED:
net = create_mobilenetv3_ssd_lite(len(class_names), is_test=True)
net.load(model_path)
dummy_input = torch.randn(1, 3, config.image_size, config.image_size)
input_names = ["input0"]
output_names = ["output0"]
# print("strt input")
torch_out = torch.onnx._export(
net, # model being run
dummy_input, # model input (or a tuple for multiple inputs)
"mobilenet_v3_ssd_lite.onnx", # where to save the model (can be a file or file-like object)
verbose=False,
export_params=True,
input_names=input_names,
output_names=output_names
config = vgg_ssd_config
elif args.net == 'mb1-ssd':
create_net = create_mobilenetv1_ssd
config = mobilenetv1_ssd_config
elif args.net == 'mb1-ssd-lite':
create_net = create_mobilenetv1_ssd_lite
config = mobilenetv1_ssd_config
elif args.net == 'sq-ssd-lite':
create_net = create_squeezenet_ssd_lite
config = squeezenet_ssd_config
elif args.net == 'mb2-ssd-lite':
create_net = lambda num: create_mobilenetv2_ssd_lite(
num, width_mult=args.mb2_width_mult)
config = mobilenetv1_ssd_config
elif args.net == 'mb3-ssd-lite':
create_net = lambda num: create_mobilenetv3_ssd_lite(num)
config = mobilenetv1_ssd_config
else:
logging.fatal("The net type is wrong.")
parser.print_help(sys.stderr)
sys.exit(1)
train_transform = TrainAugmentation(config.image_size, config.image_mean,
config.image_std)
target_transform = MatchPrior(config.priors, config.center_variance,
config.size_variance, 0.5)
test_transform = TestTransform(config.image_size, config.image_mean,
config.image_std)
logging.info("Prepare training datasets.")
datasets = []
logger.info(f"Use GPU {DEVICE}({args.gpu_ids}).")
if args.net == 'mb1-ssd':
create_net = create_mobilenetv1_ssd
config = mobilenetv1_ssd_config
elif args.net == 'mb1-ssd-lite':
create_net = create_mobilenetv1_ssd_lite
config = mobilenetv1_ssd_config
elif args.net == 'mb2-ssd-lite':
create_net = lambda num: create_mobilenetv2_ssd_lite(
num, width_mult=args.mb2_width_mult)
config = mobilenetv1_ssd_config
elif args.net == 'mb3s-ssd-lite':
create_net = lambda num: create_mobilenetv3_ssd_lite(
num,
model_mode="SMALL",
width_mult=args.mb3_width_mult,
device=DEVICE)
create_net_test = lambda num: create_mobilenetv3_ssd_lite(
num,
model_mode="SMALL",
width_mult=args.mb3_width_mult,
is_test=True,
device=DEVICE)
config = mobilenetv1_ssd_config
elif args.net == 'mb3l-ssd-lite':
create_net = lambda num: create_mobilenetv3_ssd_lite(
num,
model_mode="LARGE",
width_mult=args.mb3_width_mult,
device=DEVICE)
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,device=device)
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')
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 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)
out_video = args.out_video
Fourcc = cv2.VideoWriter_fourcc('M', 'P', '4', 'V')
writer = cv2.VideoWriter(out_video,
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=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)
elif net_type == 'mb3-ssd-lite':
predictor = create_mobilenetv3_ssd_lite_predictor(net,
candidate_size=10,
device=device)
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:
ret, orig_image = cap.read()
if orig_image is None:
print('END')
break
#continue
image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)
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)))
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()