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)