167 lines
6.2 KiB
Python
167 lines
6.2 KiB
Python
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
|
|
|
|
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)
|