def __init__(self, info):
super(SSD, self).__init__()
self.infer = info['infer']
detection_metadata = info['metadatas'][1]
if detection_metadata and 'Categories' in detection_metadata:
self.categories = detection_metadata['Categories']
else:
self.categories = ['object']
self.num_classes = len(self.categories) + 1
lib.eprint('ssd: set num_classes={}'.format(self.num_classes))
self.mode = info['params'].get('mode', 'mb2-ssd-lite')
mb2_width_mult = info['params'].get('mb2_width_mult', 1.0)
# adapt from train_ssd.py
if self.mode == 'vgg16-ssd':
create_net = create_vgg_ssd
config = vgg_ssd_config
elif self.mode == 'mb1-ssd':
create_net = create_mobilenetv1_ssd
config = mobilenetv1_ssd_config
elif self.mode == 'mb1-ssd-lite':
create_net = create_mobilenetv1_ssd_lite
config = mobilenetv1_ssd_config
elif self.mode == 'sq-ssd-lite':
create_net = create_squeezenet_ssd_lite
config = squeezenet_ssd_config
elif self.mode == 'mb2-ssd-lite':
create_net = lambda num, is_test: create_mobilenetv2_ssd_lite(
num, width_mult=mb2_width_mult, is_test=is_test)
config = mobilenetv1_ssd_config
elif self.mode == 'mb3-large-ssd-lite':
create_net = lambda num: create_mobilenetv3_large_ssd_lite(
num, is_test=is_test)
config = mobilenetv1_ssd_config
elif self.mode == 'mb3-small-ssd-lite':
create_net = lambda num: create_mobilenetv3_small_ssd_lite(
num, is_test=is_test)
config = mobilenetv1_ssd_config
config.iou_threshold = info['params'].get(
'iou_threshold', config.iou_threshold)
self.prob_threshold = info['params'].get(
'confidence_threshold', 0.01)
self.config = config
self.model = create_net(self.num_classes, is_test=self.infer)
self.criterion = MultiboxLoss(config.priors,
iou_threshold=0.5,
neg_pos_ratio=3,
center_variance=0.1,
size_variance=0.2,
device=info['device'])
self.match_prior = MatchPrior(config.priors,
config.center_variance,
config.size_variance, 0.5)
self.image_mean = torch.tensor(self.config.image_mean,
dtype=torch.float32).reshape(
1, 3, 1,
1).to(info['device'])
class_names = [name.strip() for name in open(label_path).readlines()]
print(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 == '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)
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 == 'mb2-ssd':
create_net = mb2_ssd
config = mobilenetv1_ssd_config
elif args.net == 'mb3-large-ssd-lite':
create_net = lambda num: create_mobilenetv3_large_ssd_lite(num)
config = mobilenetv1_ssd_config
elif args.net == 'mb3-small-ssd-lite':
create_net = lambda num: create_mobilenetv3_small_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 = []
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)