256 lines
9.3 KiB
Python
Executable File
256 lines
9.3 KiB
Python
Executable File
#! /usr/bin/env python
|
|
|
|
import time
|
|
import os
|
|
import argparse
|
|
import json
|
|
import cv2
|
|
import sys
|
|
sys.path += [os.path.abspath('keras-yolo3-master')]
|
|
|
|
from utils.utils import get_yolo_boxes, makedirs
|
|
from utils.bbox import draw_boxes
|
|
from keras.models import load_model
|
|
from tqdm import tqdm
|
|
import numpy as np
|
|
|
|
|
|
def disconnect(image, boxes, obj_thresh = 0.5, area_min = 400, merge = 0):
|
|
|
|
new_boxes = []
|
|
for num, box in enumerate(boxes):
|
|
|
|
xmin = box.xmin + merge
|
|
xmax = box.xmax - merge
|
|
ymin = box.ymin + merge
|
|
ymax = box.ymax - merge
|
|
|
|
if xmin > 0 and ymin > 0 and xmax < image.shape[1] and ymax < image.shape[0] and box.classes[0] > obj_thresh:
|
|
|
|
area = (ymax - ymin)*(xmax - xmin)
|
|
z_score = np.sum(image[np.int(ymin):np.int(ymax), np.int(xmin):np.int(xmax)]) / area
|
|
|
|
if area > area_min:
|
|
|
|
box.score = z_score
|
|
new_boxes.append(box)
|
|
#boxes_area_score[str(num)] = {'xmin': xmin, 'xmax': xmax, 'ymin': ymin, 'ymax': ymax, 'score' : score, 'area' : area}
|
|
|
|
mean_score = np.mean([box.score for box in new_boxes])
|
|
sd_score = np.std([box.score for box in new_boxes])
|
|
|
|
new_boxes = [box for box in new_boxes if (box.score - mean_score)/sd_score > 2]
|
|
|
|
for box in new_boxes:
|
|
|
|
z_score = (box.score - mean_score)/sd_score
|
|
box.classes[0] = min((z_score-2)*0.5+ 0.5, 1)
|
|
|
|
return new_boxes
|
|
|
|
|
|
def disconnect_plot(image, boxes, obj_thresh = 0.5, area_min = 400, merge = 0):
|
|
|
|
new_boxes = []
|
|
for num, box in enumerate(boxes):
|
|
|
|
xmin = box.xmin + merge
|
|
xmax = box.xmax - merge
|
|
ymin = box.ymin + merge
|
|
ymax = box.ymax - merge
|
|
|
|
if xmin > 0 and ymin > 0 and xmax < image.shape[1] and ymax < image.shape[0] and box.classes[0] > obj_thresh:
|
|
|
|
area = (ymax - ymin)*(xmax - xmin)
|
|
z_score = np.sum(image[np.int(ymin):np.int(ymax), np.int(xmin):np.int(xmax)]) / area
|
|
|
|
if area > area_min:
|
|
|
|
box.score = z_score
|
|
new_boxes.append(box)
|
|
#boxes_area_score[str(num)] = {'xmin': xmin, 'xmax': xmax, 'ymin': ymin, 'ymax': ymax, 'score' : score, 'area' : area}
|
|
|
|
mean_score = np.mean([box.score for box in new_boxes])
|
|
sd_score = np.std([box.score for box in new_boxes])
|
|
|
|
normal_score = ([box.score for box in new_boxes] - mean_score)/sd_score
|
|
# plt.figure()
|
|
# _ = plt.hist(normal_score, bins='auto') # arguments are passed to np.histogram
|
|
# plt.title("Histogram with 'auto' bins")
|
|
# plt.show()
|
|
#
|
|
# plt.figure()
|
|
# mean = np.mean([boxes_area_score[i]['area'] for i in boxes_area_score])
|
|
# sd = np.std([boxes_area_score[i]['area'] for i in boxes_area_score])
|
|
# normal = ([boxes_area_score[i]['area'] for i in boxes_area_score] - mean)/sd
|
|
# _ = plt.hist(normal, bins='auto') # arguments are passed to np.histogram
|
|
# plt.title("Histogram with 'auto' bins")
|
|
# plt.show()
|
|
|
|
new_boxes = [box for box in new_boxes if (box.score - mean_score)/sd_score > 2]
|
|
|
|
for box in new_boxes:
|
|
|
|
z_score = (box.score - mean_score)/sd_score
|
|
box.classes[0] = min((z_score-2)*0.5+ 0.5, 1)
|
|
|
|
|
|
|
|
|
|
colors = plt.cm.brg(np.linspace(0, 1, 21)).tolist()
|
|
plt.figure(figsize=(10,6))
|
|
plt.imshow(I,cmap = 'gray')
|
|
current_axis = plt.gca()
|
|
|
|
for box in new_boxes:
|
|
|
|
color = colors[2]
|
|
|
|
#boxes_area_score[key]['score_norm'] = (boxes_area_score[key]['score'] - mean) / sd
|
|
#z_score = (box.score - mean_score) / sd_score
|
|
#z_score = (boxes_area_score[key]['area'] )
|
|
|
|
### Escribe el z-score
|
|
#if z_score > 1:
|
|
current_axis.text((box.xmin + box.xmax)/2,
|
|
(box.ymin+ box.ymax)/2,
|
|
'%.2f' % box.classes[0], size='x-large',
|
|
color='white', bbox={'facecolor':color, 'alpha':1.0})
|
|
|
|
return new_boxes
|
|
|
|
def _main_(args):
|
|
|
|
config_path = args.conf
|
|
input_path = args.input
|
|
output_path = args.output
|
|
|
|
with open(config_path) as config_buffer:
|
|
config = json.load(config_buffer)
|
|
|
|
makedirs(output_path)
|
|
|
|
###############################
|
|
# Set some parameter
|
|
###############################
|
|
net_h, net_w = 416, 416 # a multiple of 32, the smaller the faster
|
|
obj_thresh, nms_thresh = 0.8, 0.3
|
|
|
|
###############################
|
|
# Load the model
|
|
###############################
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = config['train']['gpus']
|
|
infer_model = load_model(config['train']['saved_weights_name'])
|
|
|
|
###############################
|
|
# Predict bounding boxes
|
|
###############################
|
|
if 'webcam' in input_path: # do detection on the first webcam
|
|
video_reader = cv2.VideoCapture(0)
|
|
|
|
# the main loop
|
|
batch_size = 1
|
|
images = []
|
|
while True:
|
|
ret_val, image = video_reader.read()
|
|
if ret_val == True: images += [image]
|
|
|
|
if (len(images)==batch_size) or (ret_val==False and len(images)>0):
|
|
batch_boxes = get_yolo_boxes(infer_model, images, net_h, net_w, config['model']['anchors'], obj_thresh, nms_thresh)
|
|
|
|
for i in range(len(images)):
|
|
draw_boxes(images[i], batch_boxes[i], config['model']['labels'], obj_thresh)
|
|
cv2.imshow('video with bboxes', images[i])
|
|
images = []
|
|
if cv2.waitKey(1) == 27:
|
|
break # esc to quit
|
|
cv2.destroyAllWindows()
|
|
elif input_path[-4:] == '.mp4': # do detection on a video
|
|
video_out = output_path + input_path.split('/')[-1]
|
|
video_reader = cv2.VideoCapture(input_path)
|
|
|
|
nb_frames = int(video_reader.get(cv2.CAP_PROP_FRAME_COUNT))
|
|
frame_h = int(video_reader.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|
frame_w = int(video_reader.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|
|
|
video_writer = cv2.VideoWriter(video_out,
|
|
cv2.VideoWriter_fourcc(*'MPEG'),
|
|
50.0,
|
|
(frame_w, frame_h))
|
|
# the main loop
|
|
batch_size = 1
|
|
images = []
|
|
start_point = 0 #%
|
|
show_window = False
|
|
for i in tqdm(range(nb_frames)):
|
|
_, image = video_reader.read()
|
|
|
|
if (float(i+1)/nb_frames) > start_point/100.:
|
|
images += [image]
|
|
|
|
if (i%batch_size == 0) or (i == (nb_frames-1) and len(images) > 0):
|
|
# predict the bounding boxes
|
|
batch_boxes = get_yolo_boxes(infer_model, images, net_h, net_w, config['model']['anchors'], obj_thresh, nms_thresh)
|
|
|
|
for i in range(len(images)):
|
|
# draw bounding boxes on the image using labels
|
|
draw_boxes(images[i], batch_boxes[i], config['model']['labels'], obj_thresh)
|
|
|
|
# show the video with detection bounding boxes
|
|
if show_window: cv2.imshow('video with bboxes', images[i])
|
|
|
|
# write result to the output video
|
|
video_writer.write(images[i])
|
|
images = []
|
|
|
|
if show_window and cv2.waitKey(1) == 27: break # esc to quit
|
|
|
|
if show_window: cv2.destroyAllWindows()
|
|
video_reader.release()
|
|
video_writer.release()
|
|
else: # do detection on an image or a set of images
|
|
|
|
|
|
|
|
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'])]
|
|
|
|
# the main loop
|
|
times = []
|
|
|
|
for image_path in image_paths:
|
|
|
|
image = cv2.imread(image_path)
|
|
print(image_path)
|
|
start = time.time()
|
|
# predict the bounding boxes
|
|
boxes = get_yolo_boxes(infer_model, [image], net_h, net_w, config['model']['anchors'], obj_thresh, nms_thresh)[0]
|
|
print('Elapsed time = {}'.format(time.time() - start))
|
|
times.append(time.time() - start)
|
|
# draw bounding boxes on the image using labels
|
|
I = image.copy()
|
|
draw_boxes(I, boxes, config['model']['labels'], obj_thresh)
|
|
|
|
# write the image with bounding boxes to file
|
|
cv2.imwrite(output_path + image_path.split('/')[-1], np.uint8(image))
|
|
|
|
file = open(args.output + '/time.txt','w')
|
|
file.write('Tiempo promedio:' + str(np.mean(times)))
|
|
file.close()
|
|
|
|
if __name__ == '__main__':
|
|
argparser = argparse.ArgumentParser(description='Predict with a trained yolo model')
|
|
argparser.add_argument('-c', '--conf', help='path to configuration file')
|
|
argparser.add_argument('-i', '--input', help='path to an image, a directory of images, a video, or webcam')
|
|
argparser.add_argument('-o', '--output', default='output/', help='path to output directory')
|
|
|
|
args = argparser.parse_args()
|
|
_main_(args)
|