from keras import backend as K from keras.models import load_model from keras.preprocessing import image from keras.optimizers import Adam from imageio import imread import numpy as np from matplotlib import pyplot as plt import json import argparse import os import time import sys sys.path += [os.path.abspath('ssd_keras-master')] from models.keras_ssd300 import ssd_300 from keras_loss_function.keras_ssd_loss import SSDLoss from keras_layers.keras_layer_AnchorBoxes import AnchorBoxes from keras_layers.keras_layer_DecodeDetections import DecodeDetections from keras_layers.keras_layer_DecodeDetectionsFast import DecodeDetectionsFast from keras_layers.keras_layer_L2Normalization import L2Normalization from ssd_encoder_decoder.ssd_output_decoder import decode_detections, decode_detections_fast from data_generator.object_detection_2d_data_generator import DataGenerator from data_generator.object_detection_2d_photometric_ops import ConvertTo3Channels from data_generator.object_detection_2d_geometric_ops import Resize from data_generator.object_detection_2d_misc_utils import apply_inverse_transforms def get_session(): config = tf.ConfigProto() config.gpu_options.allow_growth = True return tf.Session(config=config) def makedirs(path): try: os.makedirs(path) except OSError: if not os.path.isdir(path): raise def _main(args=None): # parse arguments config_path = args.conf input_path = args.input_path output_path = args.output_path with open(config_path) as config_buffer: config = json.loads(config_buffer.read()) makedirs(args.output_path) ############################### # Parse the annotations ############################### score_threshold = 0.5 labels = config['model']['labels'] categories = {} #categories = {"Razor": 1, "Gun": 2, "Knife": 3, "Shuriken": 4} #la categoría 0 es la background for i in range(len(labels)): categories[labels[i]] = i+1 print('\nTraining on: \t' + str(categories) + '\n') img_height = config['model']['input'] # Height of the model input images img_width = config['model']['input'] # Width of the model input images img_channels = 3 # Number of color channels of the model input images n_classes = len(labels) # Number of positive classes, e.g. 20 for Pascal VOC, 80 for MS COCO classes = ['background'] + labels model_mode = 'training' # TODO: Set the path to the `.h5` file of the model to be loaded. model_path = config['train']['saved_weights_name'] # We need to create an SSDLoss object in order to pass that to the model loader. ssd_loss = SSDLoss(neg_pos_ratio=3, alpha=1.0) K.clear_session() # Clear previous models from memory. model = load_model(model_path, custom_objects={'AnchorBoxes': AnchorBoxes, 'L2Normalization': L2Normalization, 'DecodeDetections': DecodeDetections, 'compute_loss': ssd_loss.compute_loss}) image_paths = [] if os.path.isdir(input_path): for inp_file in os.listdir(input_path): image_paths += [input_path + inp_file] else: image_paths += [input_path] image_paths = [inp_file for inp_file in image_paths if (inp_file[-4:] in ['.jpg', '.png', 'JPEG'])] times = [] for img_path in image_paths: orig_images = [] # Store the images here. input_images = [] # Store resized versions of the images here. print(img_path) # preprocess image for network orig_images.append(imread(img_path)) img = image.load_img(img_path, target_size=(img_height, img_width)) img = image.img_to_array(img) input_images.append(img) input_images = np.array(input_images) # process image start = time.time() y_pred = model.predict(input_images) y_pred_decoded = decode_detections(y_pred, confidence_thresh=score_threshold, iou_threshold=score_threshold, top_k=200, normalize_coords=True, img_height=img_height, img_width=img_width) print("processing time: ", time.time() - start) times.append(time.time() - start) # correct for image scale # visualize detections # Set the colors for the bounding boxes colors = plt.cm.hsv(np.linspace(0, 1, 21)).tolist() plt.figure(figsize=(20,12)) plt.imshow(orig_images[0],cmap = 'gray') current_axis = plt.gca() #print(y_pred) for box in y_pred_decoded[0]: # Transform the predicted bounding boxes for the 300x300 image to the original image dimensions. xmin = box[2] * orig_images[0].shape[1] / img_width ymin = box[3] * orig_images[0].shape[0] / img_height xmax = box[4] * orig_images[0].shape[1] / img_width ymax = box[5] * orig_images[0].shape[0] / img_height color = colors[int(box[0])] label = '{}: {:.2f}'.format(classes[int(box[0])], box[1]) current_axis.add_patch(plt.Rectangle((xmin, ymin), xmax-xmin, ymax-ymin, color=color, fill=False, linewidth=2)) current_axis.text(xmin, ymin, label, size='x-large', color='white', bbox={'facecolor':color, 'alpha':1.0}) #plt.figure(figsize=(15, 15)) #plt.axis('off') save_path = output_path + img_path.split('/')[-1] plt.savefig(save_path) plt.close() file = open(output_path + 'time.txt','w') file.write('Tiempo promedio:' + str(np.mean(times))) file.close() if __name__ == '__main__': argparser = argparse.ArgumentParser(description='train and evaluate ssd model on any dataset') argparser.add_argument('-c', '--conf', help='path to configuration file') argparser.add_argument('-i', '--input_path', help='folder input.', type=str) argparser.add_argument('-o', '--output_path', help='folder output.', default='ouput/', type=str) argparser.add_argument('--score_threshold', help='score threshold detection.', default=0.5, type=float) args = argparser.parse_args() _main(args)