def main(args=None):
# parse arguments
if args is None:
args = sys.argv[1:]
args = parse_args(args)
if args.hard_score_rate:
hard_score_rate = float(args.hard_score_rate.lower())
else:
hard_score_rate = 0.5
print("hard_score_rate={}".format(hard_score_rate))
# make sure keras is the minimum required version
check_keras_version()
# optionally choose specific GPU
use_cpu = False
if args.gpu:
gpu_num = args.gpu
else:
gpu_num = str(0)
if use_cpu:
os.environ["CUDA_VISIBLE_DEVICES"] = str(666)
else:
os.environ["CUDA_VISIBLE_DEVICES"] = gpu_num
##newly added line
os.environ["CUDA_VISIBLE_DEVICES"] = gpu_num = str(0)
keras.backend.tensorflow_backend.set_session(get_session())
# make save path if it doesn't exist
if args.save_path is not None and not os.path.exists(args.save_path):
os.makedirs(args.save_path)
# create the generator
generator = create_generator(args)
# load the model
print('Loading model, this may take a second...')
model = models.load_model(os.path.join(root_dir(), args.model),
backbone_name=args.backbone,
convert=args.convert_model,
nms=False)
# start prediction
while True:
predict(generator,
model,
score_threshold=args.score_threshold,
save_path=os.path.join(root_dir(), 'res_images_iou'),
hard_score_rate=hard_score_rate)
value = input("Please enter y to continue\n")
if value != 'y':
break
def merge_detections(image_name, results):
project = 'SKU_dataset'
result_df = pandas.DataFrame()
result_df['x1'] = results[:, 0].astype(int)
result_df['y1'] = results[:, 1].astype(int)
result_df['x2'] = results[:, 2].astype(int)
result_df['y2'] = results[:, 3].astype(int)
result_df['confidence'] = results[:, 4]
result_df['hard_score'] = results[:, 5]
result_df['uuid'] = 'object_label'
result_df['label_type'] = 'object_label'
result_df['project'] = project
result_df['image_name'] = image_name
result_df.reset_index()
result_df['id'] = result_df.index
pixel_data = None
duplicate_merger = DuplicateMerger()
duplicate_merger.multiprocess = False
duplicate_merger.compression_factor = 1
project = result_df['project'].iloc[0]
image_name = result_df['image_name'].iloc[0]
if pixel_data is None:
pixel_data = read_image_bgr(os.path.join(root_dir(), image_name))
filtered_data = duplicate_merger.filter_duplicate_candidates(result_df, pixel_data)
return filtered_data
def main(args=None):
# parse arguments
if args is None:
args = sys.argv[1:]
args = parse_args(args)
if DEBUG_MODE:
args.steps = 10
# create object that stores backbone information
backbone = models.backbone(args.backbone)
# make sure keras is the minimum required version
check_keras_version()
# optionally choose specific GPU
use_cpu = False
if args.gpu:
gpu_num = args.gpu
else:
gpu_num = str(0)
if use_cpu:
os.environ["CUDA_VISIBLE_DEVICES"] = str(666)
else:
os.environ["CUDA_VISIBLE_DEVICES"] = gpu_num
keras.backend.tensorflow_backend.set_session(get_session())
# Weights and logs saves in a new locations
stmp = time.strftime("%c").replace(" ", "_")
args.snapshot_path = os.path.join(args.snapshot_path, stmp)
args.tensorboard_dir = os.path.join(args.tensorboard_dir, stmp)
print("Weights will be saved in {}".format(args.snapshot_path))
print("Logs will be saved in {}".format(args.tensorboard_dir))
create_folder(args.snapshot_path)
create_folder(args.tensorboard_dir)
# create the generators
train_generator, validation_generator = create_generators(args)
print('train_size:{},val_size:{}'.format(train_generator.size(),
validation_generator.size()))
# create the model
if args.snapshot is not None:
print('Loading model, this may take a second...')
model = models.load_model(args.snapshot, backbone_name=args.backbone)
training_model = model
prediction_model = retinanet_bbox(model=model)
else:
weights = os.path.join(os.path.join(root_dir(), args.weights))
# default to imagenet if nothing else is specified
if weights is None and args.imagenet_weights:
weights = backbone.download_imagenet()
print('Creating model, this may take a second...')
model, training_model, prediction_model = create_models(
backbone_retinanet=backbone.retinanet,
num_classes=train_generator.num_classes(),
weights=weights,
multi_gpu=args.multi_gpu,
freeze_backbone=args.freeze_backbone)
# print model summary
# print(model.summary())
# this lets the generator compute backbone layer shapes using the actual backbone model
if 'vgg' in args.backbone or 'densenet' in args.backbone:
compute_anchor_targets = functools.partial(
anchor_targets_bbox, shapes_callback=make_shapes_callback(model))
train_generator.compute_anchor_targets = compute_anchor_targets
if validation_generator is not None:
validation_generator.compute_anchor_targets = compute_anchor_targets
# create the callbacks
callbacks = create_callbacks(
model,
training_model,
prediction_model,
validation_generator,
args,
)
# start training
training_model.fit_generator(generator=train_generator,
steps_per_epoch=args.steps,
epochs=args.epochs,
verbose=1,
callbacks=callbacks,
validation_data=validation_generator,
validation_steps=validation_generator.size())
def parse_args(args):
""" Parse the arguments.
"""
parser = argparse.ArgumentParser(
description='Simple training script for training a RetinaNet network.')
subparsers = parser.add_subparsers(
help='Arguments for specific dataset types.', dest='dataset_type')
subparsers.required = True
coco_parser = subparsers.add_parser('coco')
coco_parser.add_argument('coco_path',
help='Path to dataset directory (ie. /tmp/COCO).')
pascal_parser = subparsers.add_parser('pascal')
pascal_parser.add_argument(
'pascal_path', help='Path to dataset directory (ie. /tmp/VOCdevkit).')
kitti_parser = subparsers.add_parser('kitti')
kitti_parser.add_argument(
'kitti_path', help='Path to dataset directory (ie. /tmp/kitti).')
def csv_list(string):
return string.split(',')
oid_parser = subparsers.add_parser('oid')
oid_parser.add_argument('main_dir', help='Path to dataset directory.')
oid_parser.add_argument('--version',
help='The current dataset version is v4.',
default='v4')
oid_parser.add_argument('--labels-filter',
help='A list of labels to filter.',
type=csv_list,
default=None)
oid_parser.add_argument('--annotation-cache-dir',
help='Path to store annotation cache.',
default='.')
oid_parser.add_argument('--fixed-labels',
help='Use the exact specified labels.',
default=False)
data_dir = annotation_path()
args_annotations = data_dir + '/annotations_train.csv'
args_classes = data_dir + '/class_mappings_train.csv'
args_val_annotations = data_dir + '/annotations_val.csv'
args_snapshot_path = root_dir() + '/snapshot'
args_tensorboard_dir = root_dir() + '/logs'
csv_parser = subparsers.add_parser('csv')
csv_parser.add_argument(
'--annotations',
help='Path to CSV file containing annotations for training.',
default=args_annotations)
csv_parser.add_argument(
'--classes',
help='Path to a CSV file containing class label mapping.',
default=args_classes)
csv_parser.add_argument(
'--val-annotations',
help=
'Path to CSV file containing annotations for validation (optional).',
default=args_val_annotations)
csv_parser.add_argument('--base_dir',
help='Path to base dir for CSV file.',
default=image_path())
group = parser.add_mutually_exclusive_group()
group.add_argument('--snapshot', help='Resume training from a snapshot.')
group.add_argument(
'--imagenet-weights',
help=
'Initialize the model with pretrained imagenet weights. This is the default behaviour.',
action='store_const',
const=True,
default=False)
group.add_argument('--weights',
help='Initialize the model with weights from a file.')
group.add_argument('--no-weights',
help='Don\'t initialize the model with any weights.',
dest='imagenet_weights',
action='store_const',
const=False)
parser.add_argument('--backbone',
help='Backbone model used by retinanet.',
default='resnet50',
type=str)
parser.add_argument('--batch-size',
help='Size of the batches.',
default=1,
type=int)
parser.add_argument(
'--gpu', help='Id of the GPU to use (as reported by nvidia-smi).')
parser.add_argument('--multi-gpu',
help='Number of GPUs to use for parallel processing.',
type=int,
default=0)
parser.add_argument(
'--multi-gpu-force',
help='Extra flag needed to enable (experimental) multi-gpu support.',
action='store_true')
parser.add_argument('--epochs',
help='Number of epochs to train.',
type=int,
default=150)
parser.add_argument('--steps',
help='Number of steps per epoch.',
type=int,
default=10000)
parser.add_argument(
'--snapshot-path',
help=
'Path to store snapshots of models during training (defaults to \'./snapshots\')',
default=args_snapshot_path)
parser.add_argument('--tensorboard-dir',
help='Log directory for Tensorboard output',
default=args_tensorboard_dir)
parser.add_argument('--no-snapshots',
help='Disable saving snapshots.',
dest='snapshots',
action='store_false')
parser.add_argument('--no-evaluation',
help='Disable per epoch evaluation.',
dest='evaluation',
action='store_false')
parser.add_argument('--freeze-backbone',
help='Freeze training of backbone layers.',
action='store_true')
parser.add_argument('--random-transform',
help='Randomly transform image and annotations.',
action='store_true')
parser.add_argument(
'--image-min-side',
help='Rescale the image so the smallest side is min_side.',
type=int,
default=800)
parser.add_argument(
'--image-max-side',
help='Rescale the image if the largest side is larger than max_side.',
type=int,
default=1333)
return check_args(parser.parse_args(args))
def predict(generator,
model,
score_threshold=0.05,
max_detections=9999,
save_path=None,
hard_score_rate=1.):
all_detections = [[None for i in range(generator.num_classes())]
for j in range(generator.size())]
csv_data_lst = []
csv_data_lst.append(
['image_id', 'x1', 'y1', 'x2', 'y2', 'confidence', 'hard_score'])
result_dir = os.path.join(root_dir(), 'results')
create_folder(result_dir)
timestamp = datetime.datetime.utcnow()
res_file = result_dir + '/detections_output_iou_{}_{}.csv'.format(
hard_score_rate, timestamp)
for i in range(generator.size()):
image_name = os.path.join(
generator.image_path(i).split(os.path.sep)[-2],
generator.image_path(i).split(os.path.sep)[-1])
raw_image = generator.load_image(i)
image = generator.preprocess_image(raw_image.copy())
image, scale = generator.resize_image(image)
# run network
boxes, hard_scores, labels, soft_scores = model.predict_on_batch(
np.expand_dims(image, axis=0))
soft_scores = np.squeeze(soft_scores, axis=-1)
soft_scores = hard_score_rate * hard_scores + (
1 - hard_score_rate) * soft_scores
# correct boxes for image scale
boxes /= scale
# select indices which have a score above the threshold
indices = np.where(hard_scores[0, :] > score_threshold)[0]
# select those scores
scores = soft_scores[0][indices]
hard_scores = hard_scores[0][indices]
# find the order with which to sort the scores
scores_sort = np.argsort(-scores)[:max_detections]
# select detections
image_boxes = boxes[0, indices[scores_sort], :]
image_scores = scores[scores_sort]
image_hard_scores = hard_scores[scores_sort]
image_labels = labels[0, indices[scores_sort]]
image_detections = np.concatenate([
image_boxes,
np.expand_dims(image_scores, axis=1),
np.expand_dims(image_labels, axis=1)
],
axis=1)
results = np.concatenate([
image_boxes,
np.expand_dims(image_scores, axis=1),
np.expand_dims(image_hard_scores, axis=1),
np.expand_dims(image_labels, axis=1)
],
axis=1)
filtered_data = EmMerger.merge_detections(image_name, results)
filtered_boxes = []
filtered_scores = []
filtered_labels = []
for ind, detection in filtered_data.iterrows():
box = np.asarray([
detection['x1'], detection['y1'], detection['x2'],
detection['y2']
])
filtered_boxes.append(box)
filtered_scores.append(detection['confidence'])
filtered_labels.append('{0:.2f}'.format(detection['hard_score']))
row = [
image_name, detection['x1'], detection['y1'], detection['x2'],
detection['y2'], detection['confidence'],
detection['hard_score']
]
csv_data_lst.append(row)
if save_path is not None:
create_folder(save_path)
draw_annotations(raw_image,
generator.load_annotations(i),
label_to_name=generator.label_to_name)
draw_detections(raw_image,
np.asarray(filtered_boxes),
np.asarray(filtered_scores),
np.asarray(filtered_labels),
color=(0, 0, 255))
cv2.imwrite(os.path.join(save_path, '{}.png'.format(i)), raw_image)
# copy detections to all_detections
for label in range(generator.num_classes()):
all_detections[i][label] = image_detections[
image_detections[:, -1] == label, :-1]
print('{}/{}'.format(i + 1, generator.size()), end='\r')
# Save annotations csv file
with open(res_file, 'wb') as fl_csv:
writer = csv.writer(fl_csv)
writer.writerows(csv_data_lst)
print("Saved output.csv file")
def main(image_folder_path: str, detection_save_folder: str, model_wights_path: str,
classes=None,
hard_score_rate = 0.5,
backbone = 'resnet50',
gpu = None,
score_threshold = 0.1,
iou_threshold = 0.5,
image_min_side = 800,
image_max_side = 1333,
args = None):
# parse arguments, to fill args with parameters initialized from defaults
if args is None:
args = ["--model", model_wights_path, 'csv']
args = parse_args(args)
if args.hard_score_rate is None:
args.hard_score_rate = hard_score_rate
# hard_score_rate = float(args.hard_score_rate.lower())
if args.model is None:
args.model = model_wights_path
if args.base_dir is None:
args.base_dir = image_folder_path
if args.backbone is None:
args.backbone = backbone
if args.gpu is None:
args.gpu = gpu
if args.score_threshold is None:
args.score_threshold = float(score_threshold)
if args.iou_threshold is None:
args.iou_threshold = float(iou_threshold)
if args.save_path is None:
args.save_path = detection_save_folder
if args.image_max_side is None:
args.image_max_side = int(image_max_side)
if args.image_min_side is None:
args.image_min_side = int(image_min_side)
# make sure keras is the minimum required version
check_keras_version()
# optionally choose specific GPU
use_cpu = False
if args.gpu:
gpu_num = args.gpu
else:
gpu_num = str(0)
if use_cpu:
os.environ["CUDA_VISIBLE_DEVICES"] = str(666)
else:
os.environ["CUDA_VISIBLE_DEVICES"] = gpu_num
keras.backend.tensorflow_backend.set_session(get_session())
# make save path if it doesn't exist
if args.save_path is not None and not os.path.exists(args.save_path):
os.makedirs(args.save_path)
# create the generator
generator = create_generator(args, args.base_dir)
# load the model
print('Loading model, this may take a second...')
model = models.load_model(os.path.join(root_dir(), args.model), backbone_name=args.backbone, convert=args.convert_model, nms=False)
# start prediction
detection_csv_results_file_path = predict(
generator,
model,
score_threshold=args.score_threshold,
# save_path=os.path.join(root_dir(), 'res_images_iou'),
image_save_path = os.path.join(args.save_path, "detection_results_images"),
results_save_path = args.save_path,
hard_score_rate=hard_score_rate)
return detection_csv_results_file_path
