def __init__(self,
root_dir,
ssd,
liste_obj,
num_examples=-1,
clear_data=False,
augmentation=None):
super(VOCDataset, self).__init__()
self.idx_to_name = liste_obj
self.name_to_idx = dict([(v, k)
for k, v in enumerate(self.idx_to_name)])
self.data_dir = os.path.join(root_dir, 'VOC2012')
self.image_dir = os.path.join(self.data_dir, 'JPEGImages')
self.anno_dir = os.path.join(self.data_dir, 'Annotations')
self.ids = list(map(lambda x: x[:-4], os.listdir(self.image_dir)))
self.default_boxes = generate_default_boxes(ssd.config)
self.new_size = ssd.image_size
if clear_data:
list_pop = []
for index in range(len(self.ids)):
img = self._get_image(index)
w, h = img.size
_, labels = self._get_annotation(index, (h, w))
nb = [0 for i in range(len(self.idx_to_name) + 1)]
for label in labels:
if 0 < label < len(self.idx_to_name) + 1:
nb[label] += 1
if not (np.array(nb).any() != 0):
list_pop.append(index)
for i in range(len(list_pop) - 1, 0, -1):
(self.ids).pop(list_pop[i])
if num_examples != -1:
self.ids = self.ids[:num_examples]
self.train_ids = self.ids[:int(len(self.ids) * 0.75)]
self.val_ids = self.ids[int(len(self.ids) * 0.75):]
if augmentation is None:
self.augmentation = ['original']
else:
self.augmentation = augmentation + ['original']
def __init__(self):
self.dataset_name = 'SED-dataset'
self.image_dir = config.image_dir
self.dataset_dir = config.dataset_dir
self.id_to_label = set_id_to_label(config.label_set)
self.input_shape = (config.SSD300['image_size'],
config.SSD300['image_size'])
self.default_boxes = anchor.generate_default_boxes(config)
self.train_annotation_path = os.path.join(self.dataset_dir,
'annotations',
'CNSI-SED-TRAIN.json')
self.val_annotation_path = os.path.join(self.dataset_dir,
'annotations',
'CNSI-SED-VAL.json')
self.eval_annotation_path = os.path.join(self.dataset_dir,
'annotations',
'CNSI-SED-TEST.json')
def predict(self, imgs, **kwargs):
default_boxes = generate_default_boxes(self.config)
confs, locs = self.call(imgs)
confs = tf.squeeze(confs, 0)
locs = tf.squeeze(locs, 0)
confs = tf.math.softmax(confs, axis=-1)
# classes = tf.math.argmax(confs, axis=-1)
# scores = tf.math.reduce_max(confs, axis=-1)
boxes = decode(default_boxes, locs)
out_boxes = []
out_labels = []
out_scores = []
for c in range(1, self.num_classes):
cls_scores = confs[:, c]
score_idx = cls_scores > 0.6
# cls_boxes = tf.boolean_mask(boxes, score_idx)
# cls_scores = tf.boolean_mask(cls_scores, score_idx)
cls_boxes = boxes[score_idx]
cls_scores = cls_scores[score_idx]
nms_idx = compute_nms(cls_boxes, cls_scores, 0.45, 200)
cls_boxes = tf.gather(cls_boxes, nms_idx)
cls_scores = tf.gather(cls_scores, nms_idx)
cls_labels = [c] * cls_boxes.shape[0]
out_boxes.append(cls_boxes)
out_labels.extend(cls_labels)
out_scores.append(cls_scores)
out_boxes = tf.concat(out_boxes, axis=0)
out_scores = tf.concat(out_scores, axis=0)
boxes = tf.clip_by_value(out_boxes, 0.0, 1.0).numpy()
classes = np.array(out_labels)
scores = out_scores.numpy()
return boxes, classes, scores
try:
ssd = network.create_ssd(NUM_CLASSES, config.arch,
config.pretrained_type, config.checkpoint_dir,
config.checkpoint_path)
except Exception as e:
print(e)
print('The program is exiting...')
sys.exit()
os.makedirs('outputs/images', exist_ok=True)
os.makedirs('outputs/detects', exist_ok=True)
visualizer = ImageVisualizer(config.label_set, save_dir='outputs/images')
for i, (filename, imgs, gt_confs, gt_locs) in enumerate(test_generator):
default_boxes = anchor.generate_default_boxes(config)
boxes, classes, scores = predict(imgs, default_boxes)
filename = filename.numpy()[0].decode()
original_image = Image.open(
os.path.join(config.image_dir, '{}'.format(filename)))
boxes *= original_image.size * 2
visualizer.save_image(original_image, boxes, classes,
'{}'.format(filename))
log_file = os.path.join('outputs/detects', '{}.txt')
for cls, box, score in zip(classes, boxes, scores):
cls_name = config.label_set[cls - 1]
with open(log_file.format(cls_name), 'a') as f:
f.write('{} {} {} {} {} {}\n'.format(filename, score,
*[coord
for coord in box]))
return boxes, classes, scores
if __name__ == "__main__":
with open("./config.yml") as f:
cfg = yaml.load(f)
try:
config = cfg[args.arch.upper()]
global BASE_CONFIG
BASE_CONFIG = cfg['BASE']
except AttributeError:
raise ValueError("Unknown architecture: {}".format(args.arch))
default_boxes = generate_default_boxes(config)
batch_generator, info = create_batch_generator(
root_dir=args.data_dir,
year=args.data_year,
default_boxes=default_boxes,
new_size=config["image_size"],
batch_size=BASE_CONFIG['batch_size'],
num_batches=args.num_examples,
labels=cfg['LABELS']['sedssd'],
mode="test"
)
try:
if 'se' in args.arch:
network = create_sedssd(