def data_draw_annotxt(annotation_path, classname_txt, save_dir):
content = open(classname_txt, 'r').readlines()
classnames = [c.strip() for c in content]
color_table = get_color_table(len(classnames))
anno = open(annotation_path, 'r')
for line in anno:
s = line.strip().split(' ')
imgpath = s[0]
dst = os.path.join(save_dir, os.path.basename(imgpath))
s = s[1:]
box_cnt = len(s) // 5
box_names = []
for i in range(box_cnt):
x_min, y_min, x_max, y_max = float(s[i * 5 + 1]), float(
s[i * 5 + 2]), float(s[i * 5 + 3]), float(s[i * 5 + 4])
idx = int(s[i * 5])
box_names.append([x_min, y_min, x_max, y_max, classnames[idx]])
if len(box_names) == 0:
continue
save_labeled_img(imgpath,
box_names,
dst,
classnames,
color_table=color_table)
print(imgpath, " is saved.")
def __init__(self):
self.counter = 0 # number of people
self.violation = 0 # number of violations
# The path of the anchor txt file.
anchor_path = "./data/yolo_anchors.txt"
# Resize the input image with `new_size`, size format: [width, height]
self.new_size = [416, 416]
# Whether to use the letterbox resize.
self.letterbox_resizes = True
# The path of the class names
class_name_path = "./data/coco.names"
# The path of the weights to restore
restore_path = "./checkpoint/best_model_Epoch_75_step_29487_mAP_0.8609_loss_5.4903_lr_1e-05"
# Whether to save the video detection results.
self.save_video = True
self.anchors = parse_anchors(anchor_path)
self.classes = read_class_names(class_name_path)
self.num_class = len(self.classes)
self.color_table = get_color_table(
self.num_class) # color for each label
self.tracker = Sort()
self.memory = {}
self.COLORS = np.random.randint(0, 255, size=(200, 3),
dtype="uint8") # tracker color
self.sess = tf.Session()
self.input_data = tf.compat.v1.placeholder(
tf.float32, [1, self.new_size[1], self.new_size[0], 3],
name='input_data')
yolo_model = yolov3(self.num_class, self.anchors)
with tf.compat.v1.variable_scope('yolov3'):
pred_feature_maps = yolo_model.forward(self.input_data, False)
pred_boxes, pred_confs, pred_probs = yolo_model.predict(
pred_feature_maps)
pred_scores = pred_confs * pred_probs
self.boxes, self.scores, self.labels = gpu_nms(pred_boxes,
pred_scores,
self.num_class,
max_boxes=200,
score_thresh=0.6,
nms_thresh=0.01)
saver = tf.compat.v1.train.Saver()
saver.restore(self.sess, restore_path)
def __init__(self):
self.new_size = cfg.img_size
self.resize = True
self.class_name_path = './data/mydata.names'
self.restore_path = cfg.restore_path
self.classes = read_class_names(self.class_name_path)
self.num_class = len(self.classes)
self.color_table = get_color_table(self.num_class)
self.sess = tf.Session()
self.init_params()
def __init(self):
self.anchor_path = "./data/yolo_anchors.txt"
self.new_size = [416, 416]
self.class_name_path = "./data/my_data/data.names"
self.restore_path = "./checkpoint/model-step_17500_loss_0.003654_lr_0.0004995866"
self.anchors = parse_anchors(self.anchor_path)
self.classes = read_class_names(self.class_name_path)
self.num_class = len(self.classes)
self.color_tabel = get_color_table(self.num_class)
self.img_ori = cv2.imread(self.input_image)
self.weight_ori, self.width_ori = self.img_ori.shape[:2]
self.img = cv2.resize(img_ori, tuple(self.new_size))
self.img = cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB)
self.img = np.asarray(self.img, np.float32)
self.img = self.img[np.newaxis, :] / 255.
#config = tf.ConfigProto()
#config.gpu_options.allow_growth = True
self.__sess = tf.Session()
self.__sess.run(tf.global_variables_initializer())
yolo_model = yolov3(self.num_class, self.anchors)
with tf.variable_scope('yolov3'):
pred_feature_maps = yolo_model.forward(input_data, False)
pred_boxes, pred_confs, pred_probs = yolo_model.predict(
pred_feature_maps)
pred_scores = pred_confs * pred_probs
self.boxes, self.scores, self.labels = gpu_nms(pred_boxes,
pred_scores,
args.num_class,
max_boxes=100,
score_thresh=0.4,
iou_thresh=0.5)
self.__saver = tf.train.Saver()
self.__saver.restore(self.__sess, self.restore_path)
self.input_data = tf.placeholder(
tf.float32, [1, self.new_size[1], self.new_size[0], 3],
name='input_data')
def show_image(img_ori, boxes_, scores_, classes, num_class, labels_, width_ori, height_ori, new_size):
# rescale the coordinates to the original image
color_table = get_color_table(num_class)
boxes_[:, 0] *= (width_ori / float(new_size[0]))
boxes_[:, 2] *= (width_ori / float(new_size[0]))
boxes_[:, 1] *= (height_ori / float(new_size[1]))
boxes_[:, 3] *= (height_ori / float(new_size[1]))
print("box coords:")
print(boxes_)
print('*' * 30)
print("scores:")
print(scores_)
print('*' * 30)
print("labels:")
print(labels_)
for i in range(len(boxes_)):
x0, y0, x1, y1 = boxes_[i]
scores=scores_[i]
plot_one_box(img_ori, [x0, y0, x1, y1], label=classes[labels_[i]]+':'+str(scores)[:6], color=color_table[labels_[i]])
cv2.imshow('Detection result', img_ori)
# cv2.imwrite('detection_result.jpg', img_ori)
cv2.waitKey(30)
def yolodet(anchor_path, image_path, new_size, letterbox, class_name_path, restore_path):
anchors = parse_anchors(anchor_path)
classes = read_class_names(class_name_path)
num_class = len(classes)
color_table = get_color_table(num_class)
img_ori = cv2.imread(image_path)
if letterbox:
img, resize_ratio, dw, dh = letterbox_resize(img_ori, new_size[0], new_size[1])
else:
height_ori, width_ori = img_ori.shape[:2]
img = cv2.resize(img_ori, tuple(new_size))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = np.asarray(img, np.float32)
img = img[np.newaxis, :] / 255.
with tf.Session() as sess:
input_data = tf.placeholder(tf.float32, [1, new_size[1], new_size[0], 3], name='input_data')
yolo_model = yolov3(num_class, anchors)
with tf.variable_scope('yolov3'):
pred_feature_maps = yolo_model.forward(input_data, False)
pred_boxes, pred_confs, pred_probs = yolo_model.predict(pred_feature_maps)
pred_scores = pred_confs * pred_probs
boxes, scores, labels = gpu_nms(pred_boxes, pred_scores, num_class, max_boxes=200, score_thresh=0.3, nms_thresh=0.45)
saver = tf.train.Saver()
saver.restore(sess, restore_path)
boxes_, scores_, labels_ = sess.run([boxes, scores, labels], feed_dict={input_data: img})
# rescale the coordinates to the original image
if letterbox:
boxes_[:, [0, 2]] = (boxes_[:, [0, 2]] - dw) / resize_ratio
boxes_[:, [1, 3]] = (boxes_[:, [1, 3]] - dh) / resize_ratio
else:
boxes_[:, [0, 2]] *= (width_ori/float(new_size[0]))
boxes_[:, [1, 3]] *= (height_ori/float(new_size[1]))
# print("box coords:")
# print(boxes_)
# print('*' * 30)
# print("scores:")
# print(scores_)
# print('*' * 30)
# print("labels:")
# print(labels_)
#
# for i in range(len(boxes_)):
# x0, y0, x1, y1 = boxes_[i]
# plot_one_box(img_ori, [x0, y0, x1, y1], label=classes[labels_[i]] + ', {:.2f}%'.format(scores_[i] * 100), color=color_table[labels_[i]])
# cv2.imshow('Detection result', img_ori)
# cv2.imwrite('detection_result.jpg', img_ori)
# cv2.waitKey(0)
tf.reset_default_graph()
return boxes_, scores_, labels_
from modelv4tiny import yolov4tiny
from utils.data_aug import letterbox_resize
from utils.misc_utils import parse_anchors, read_class_names
from utils.nms_utils import gpu_nms
from utils.plot_utils import get_color_table, plot_one_box
config = tf.compat.v1.ConfigProto(gpu_options=tf.compat.v1.GPUOptions(
allow_growth=True))
anchor_path = "./data/yolo_tiny_anchors.txt"
class_name_path = "./data/coco.names"
anchors = parse_anchors(anchor_path)
classes = read_class_names(class_name_path)
class_num = len(classes)
color_table = get_color_table(class_num)
img_size = [416, 416]
yolo_model = yolov4tiny(class_num, anchors)
yolo_model.set_img_size(np.asarray(img_size))
sess = tf.Session(config=config)
with gfile.FastGFile("./pb_model/frozen_model_v4tiny.pb", 'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
sess.graph.as_default()
tf.import_graph_def(graph_def, name='')
sess.run(tf.global_variables_initializer())
input = sess.graph.get_tensor_by_name('image:0')
feature_map_1 = sess.graph.get_tensor_by_name(
'yolov4tiny/head/feature_map_1:0')
feature_map_2 = sess.graph.get_tensor_by_name(
cv2.rectangle(npimg,(int(boxes_[j][0]),int(boxes_[j][1])),(int(boxes_[j][2]),int(boxes_[j][3])),(0,0,255), 2)
print('Existing Smoke!')
#cv2.imwrite(img_name, npimg)
# boxes_[j][0]+=person_left
# boxes_[j][2]+=person_left
# boxes_[j][1]+=person_up
# boxes_[j][3]+=person_up
smoke_boxes.append(boxes_[j])
smoke_scores.append(scores_[j])
else:
print('*'*30)
print('No existing Smoke!')
else:
print('No existing Smoke!')
classes = {0:'smoke'}
color_table = get_color_table(1)
for i in range(len(smoke_boxes)):
x0, y0, x1, y1 = smoke_boxes[i]
print('change_box:', smoke_boxes[i])
plot_one_box(img_ori, [x0, y0, x1, y1],
label='smoke' + ', {:.2f}%'.format(smoke_scores[i] * 100),
color=color_table[0])
cv2.imwrite('final_result.jpg', img_ori)
# try:
# import matplotlib.pyplot as plt
#
# fig = plt.figure()
# a = fig.add_subplot(2, 2, 1)
# a.set_title('Result')
# #plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
# img1 = cv2.cvtColor(image,cv2.COLOR_BGR2RGB)
def json2txt(self, noannodst_savedir=None, labeled_img_savedir=None):
# setting loggers
logger = self.logger
json_dir = self.json_dir
annotation_txt = self.annotation_txt_path
class_name_path = self.class_name_path
json_names = os.listdir(json_dir)
annotations = []
classnames = []
f_anno = open(annotation_txt, 'w')
wcount = 0
for json_name in json_names:
imgname = json_name[:-5] + ".jpg"
imgpath = os.path.join(self.img_dir, imgname)
line = imgpath
jpath = os.path.join(json_dir, json_name)
f = open(jpath, encoding='utf-8')
j = json.load(f)
box_names = []
if j['labeled'] is False:
if noannodst_savedir is not None:
dst_file = os.path.join(noannodst_savedir, imgname)
shutil.copyfile(imgpath, dst_file)
info = imgname + ' is not labeled.'
logger.warning(info)
wcount += 1
continue
else:
w, h = j['size']['width'], j['size']['height']
ls = j['outputs']['object']
for obj in ls:
name, box = '', {}
for objk in obj:
if objk == 'name':
name = obj[objk]
else:
box = obj[objk]
if name not in classnames:
classnames.append(name)
# if name=='power':
# print('power:',imgname)
if 'xmin' in box:
xmin, ymin, xmax, ymax = max(box['xmin'], 0), max(
box['ymin'],
0), min(box['xmax'],
w - 1), min(box['ymax'], h - 1)
else:
xmin, ymin, xmax, ymax = 10000, 10000, 0, 0
for key in box:
if key[0] == 'x':
if box[key] < xmin:
xmin = box[key]
if box[key] > xmax:
xmax = box[key]
elif key[0] == 'y':
if box[key] < ymin:
ymin = box[key]
if box[key] > ymax:
ymax = box[key]
else:
info = json_name + ' has wrong box key(not x or y).'
logger.error(info)
xmin, ymin, xmax, ymax = max(xmin, 0), max(
ymin, 0), min(xmax, w - 1), min(ymax, h - 1)
if xmin >= xmax or ymin >= ymax:
info = json_name + ' has wrong annotation.'
logger.error(info)
continue
box_names.append([xmin, ymin, xmax, ymax, name])
line = line + ' ' + str(
classnames.index(name)) + ' ' + str(xmin) + ' ' + str(
ymin) + ' ' + str(xmax) + ' ' + str(ymax)
if len(box_names) == 0:
continue
annotations.append(line)
f_anno.write(line)
# logger.info(line)
f_anno.write('\n')
print(line)
if labeled_img_savedir is not None:
dst_file = os.path.join(labeled_img_savedir, imgname)
color_table = get_color_table(len(classnames))
# shutil.copyfile(imgpath, dst_file)
save_labeled_img(imgpath,
box_names,
dst_file,
classnames,
color_table=color_table)
f_anno.close()
info = 'wrong pics:{}'.format(wcount)
logger.info(info)
print(info)
f_class = open(class_name_path, 'w')
for cla in classnames:
f_class.write(cla)
f_class.write('\n')
f_class.close()
logger.info('class.names has been saved.')
def single_image_test(imgname):
parser = argparse.ArgumentParser(
description="YOLO-V3 test single image test procedure.")
parser.add_argument("--input_image",
type=str,
default="./static/uploads/beforeimg/" + imgname,
help="The path of the input image.")
parser.add_argument("--anchor_path",
type=str,
default="./data/yolo_anchors.txt",
help="The path of the anchor txt file.")
parser.add_argument(
"--new_size",
nargs='*',
type=int,
default=[416, 416],
help=
"Resize the input image with `new_size`, size format: [width, height]")
parser.add_argument("--letterbox_resize",
type=lambda x: (str(x).lower() == 'true'),
default=True,
help="Whether to use the letterbox resize.")
parser.add_argument("--class_name_path",
type=str,
default="./data/coco.names",
help="The path of the class names.")
parser.add_argument("--restore_path",
type=str,
default="./data/darknet_weights/yolov3.ckpt",
help="The path of the weights to restore.")
args = parser.parse_args()
args.anchors = parse_anchors(args.anchor_path)
args.classes = read_class_names(args.class_name_path)
args.num_class = len(args.classes)
color_table = get_color_table(args.num_class)
img_ori = cv2.imread(args.input_image)
if args.letterbox_resize:
img, resize_ratio, dw, dh = letterbox_resize(img_ori, args.new_size[0],
args.new_size[1])
else:
height_ori, width_ori = img_ori.shape[:2]
img = cv2.resize(img_ori, tuple(args.new_size))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = np.asarray(img, np.float32)
img = img[np.newaxis, :] / 255.
with tf.Session() as sess:
input_data = tf.placeholder(tf.float32,
[1, args.new_size[1], args.new_size[0], 3],
name='input_data')
yolo_model = yolov3(args.num_class, args.anchors)
with tf.variable_scope('yolov3'):
pred_feature_maps = yolo_model.forward(input_data, False)
pred_boxes, pred_confs, pred_probs = yolo_model.predict(
pred_feature_maps)
pred_scores = pred_confs * pred_probs
boxes, scores, labels = gpu_nms(pred_boxes,
pred_scores,
args.num_class,
max_boxes=200,
score_thresh=0.3,
nms_thresh=0.45)
saver = tf.train.Saver()
saver.restore(sess, args.restore_path)
boxes_, scores_, labels_ = sess.run([boxes, scores, labels],
feed_dict={input_data: img})
# rescale the coordinates to the original image
if args.letterbox_resize:
boxes_[:, [0, 2]] = (boxes_[:, [0, 2]] - dw) / resize_ratio
boxes_[:, [1, 3]] = (boxes_[:, [1, 3]] - dh) / resize_ratio
else:
boxes_[:, [0, 2]] *= (width_ori / float(args.new_size[0]))
boxes_[:, [1, 3]] *= (height_ori / float(args.new_size[1]))
print("box coords:")
print(boxes_)
print('*' * 30)
print("scores:")
print(scores_)
print('*' * 30)
print("labels:")
print(labels_)
for i in range(len(boxes_)):
x0, y0, x1, y1 = boxes_[i]
plot_one_box(img_ori, [x0, y0, x1, y1],
label=args.classes[labels_[i]] +
', {:.2f}%'.format(scores_[i] * 100),
color=color_table[labels_[i]])
#cv2.imshow('Detection result', img_ori)
cv2.imwrite('static/uploads/afterimg/' + imgname, img_ori)
#cv2.waitKey(0)
doc = []
doc.append("发现：")
item = ["安全帽", "未带安全帽的人"]
if (len(labels_) == 0):
doc.append("什么都没有发现。")
else:
for i in range(len(labels_)):
doc.append(item[labels_[i]] + ",范围：" + str(boxes_[i]) +
",可能性为：" + str(scores_[i]))
return doc
def yolodet(image_path,
anchor_path=rootpath + "/yolo/data/yolo_anchors.txt",
new_size=[416, 416],
letterbox=True,
class_name_path=rootpath + "/yolo/data/coco.names",
restore_path=rootpath + "/yolo/data/best_model"):
anchors = parse_anchors(anchor_path)
classes = read_class_names(class_name_path)
num_class = len(classes)
color_table = get_color_table(num_class)
img_ori = cv2.imread(image_path)
if letterbox:
img, resize_ratio, dw, dh = letterbox_resize(img_ori, new_size[0],
new_size[1])
else:
height_ori, width_ori = img_ori.shape[:2]
img = cv2.resize(img_ori, tuple(new_size))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = np.asarray(img, np.float32)
img = img[np.newaxis, :] / 255.
with tf.Session() as sess:
input_data = tf.placeholder(tf.float32,
[1, new_size[1], new_size[0], 3],
name='input_data')
yolo_model = yolov3(num_class, anchors)
with tf.variable_scope('yolov3'):
pred_feature_maps = yolo_model.forward(input_data, False)
pred_boxes, pred_confs, pred_probs = yolo_model.predict(
pred_feature_maps)
pred_scores = pred_confs * pred_probs
boxes, scores, labels = gpu_nms(pred_boxes,
pred_scores,
num_class,
max_boxes=200,
score_thresh=0.3,
nms_thresh=0.45)
saver = tf.train.Saver()
saver.restore(sess, restore_path)
boxes_, scores_, labels_ = sess.run([boxes, scores, labels],
feed_dict={input_data: img})
# rescale the coordinates to the original image
if letterbox:
boxes_[:, [0, 2]] = (boxes_[:, [0, 2]] - dw) / resize_ratio
boxes_[:, [1, 3]] = (boxes_[:, [1, 3]] - dh) / resize_ratio
else:
boxes_[:, [0, 2]] *= (width_ori / float(new_size[0]))
boxes_[:, [1, 3]] *= (height_ori / float(new_size[1]))
tf.reset_default_graph()
#transform detections into 1 line (#1class,#1conf,#1xmin,#1ymin,#1max,#1ymax,#2class,#2conf,...)
boxes = []
for i in range(np.shape(boxes_)[0]):
boxes.append(labels_[i])
boxes.append(scores_[i])
boxes.extend(boxes_[i, :])
return boxes
return boxes, classes.astype(np.int64), scores, num_detections.astype(
np.int64)
if __name__ == '__main__':
import os
import glob
import numpy as np
import cv2
from utils.plot_utils import get_color_table, plot_one_box
from utils.misc_utils import parse_anchors, read_class_names
model = YoloV3('./data/darknet_weights/yolov3_frozen_graph_batch.pb')
classes = read_class_names("./data/coco.names")
color_table = get_color_table(80)
files = glob.glob('./data/demo_data/*.jpg')
images = []
vis_images = []
for file in files:
image = cv2.imread(file)
image = cv2.resize(image, (640, 640))
vis_images.append(image)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) / 255. #important!
images.append(image)
images = np.array(images)
# inference
boxes_, labels_, scores_, num_dect_ = model.run(images)
# visualize
for idx, image in enumerate(vis_images):
for i in range(len(boxes_[idx])):
import cv2
from utils.misc_utils import parse_anchors, read_class_names
from utils.nms_utils import gpu_nms
from utils.plot_utils import get_color_table, plot_one_box
from model import yolov3
input_image = "/home/yhq/Desktop/yolo_v3/data/demo_data/kite.jpg" # 输入测试图片
anchor_path = "/home/yhq/Desktop/yolo_v3/data/yolo_anchors.txt" # anchor文件路径
new_size = [416, 416] # 模型输入图片大小
class_name_path = "/home/yhq/Desktop/yolo_v3/data/coco.names" # 类别名称文件路径
restore_path = "/home/yhq/Desktop/yolo_v3/data/darknet_weights/yolov3.ckpt" # 检查点文件
anchors = parse_anchors(anchor_path)
classes = read_class_names(class_name_path)
num_class = len(classes)
color_table = get_color_table(num_class) # 每个类别对应一个color
# 读图片并进行一些简单处理
img_ori = cv2.imread(input_image)
height_ori, width_ori = img_ori.shape[:2]
img = cv2.resize(img_ori, tuple(new_size))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = np.asarray(img, np.float32)
img = img[np.newaxis, :] / 255.
# 识别工作
with tf.Session() as sess:
input_data = tf.placeholder(tf.float32, [1, new_size[1], new_size[0], 3],
name='input_data')
yolo_model = yolov3(num_class, anchors)
with tf.variable_scope('yolov3'):
pred_feature_maps = yolo_model.forward(input_data, False)
saver = tf.train.Saver() # 变量初始化后建立
def recognize(jpg_path, pb_file_path):
anchors = parse_anchors("./data/yolo_anchors.txt")
classes = read_class_names("./data/coco.names")
num_class = len(classes)
color_table = get_color_table(num_class)
img_ori = cv2.imread(jpg_path)
height_ori, width_ori = img_ori.shape[:2]
img = cv2.resize(img_ori, tuple([IMAGE_SIZE, IMAGE_SIZE]))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = np.asarray(img, np.float32)
img = img[np.newaxis, :] / 255.
with tf.Graph().as_default():
output_graph_def = tf.GraphDef()
print("Load Frozen_Graph File ...")
with open(pb_file_path, "rb") as f:
output_graph_def.ParseFromString(f.read())
tf.import_graph_def(output_graph_def, name="")
print("Finished")
# GPU_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)
config = tf.ConfigProto() # gpu_options=GPU_options)
config.gpu_options.allow_growth = True
with tf.Session(config=config) as sess:
# Define Input and Outputs
input_x = sess.graph.get_tensor_by_name("Placeholder:0")
feature_map_1 = sess.graph.get_tensor_by_name(
"yolov3/yolov3_head/feature_map_1:0")
feature_map_2 = sess.graph.get_tensor_by_name(
"yolov3/yolov3_head/feature_map_2:0")
feature_map_3 = sess.graph.get_tensor_by_name(
"yolov3/yolov3_head/feature_map_3:0")
features = feature_map_1, feature_map_2, feature_map_3
# yolo config
yolo_model = yolov3(num_class, anchors)
yolo_model.pb_forward(input_x)
# # use frozen_graph to inference
# print "RUN Graph ..."
# features = sess.run(features, feed_dict={input_x:np.reshape(img, [-1, IMAGE_SIZE, IMAGE_SIZE, 3])})
# print "Finished"
# feature1, feature2, feature3 = features
# feature1 = tf.convert_to_tensor(feature1)
# feature2 = tf.convert_to_tensor(feature2)
# feature3 = tf.convert_to_tensor(feature3)
# features = feature1, feature2, feature3
print "Predicting ..."
pred_boxes, pred_confs, pred_probs = yolo_model.predict(features)
pred_scores = pred_confs * pred_probs
boxes, scores, labels = gpu_nms(pred_boxes,
pred_scores,
num_class,
max_boxes=30,
score_thresh=0.4,
iou_thresh=0.5)
t0 = time.time()
boxes_, scores_, labels_ = sess.run(
[boxes, scores, labels],
feed_dict={
input_x: np.reshape(img, [-1, IMAGE_SIZE, IMAGE_SIZE, 3])
})
t1 = time.time()
print "Finished"
# rescale the coordinates to the original image
boxes_[:, 0] *= (width_ori / float(IMAGE_SIZE))
boxes_[:, 2] *= (width_ori / float(IMAGE_SIZE))
boxes_[:, 1] *= (height_ori / float(IMAGE_SIZE))
boxes_[:, 3] *= (height_ori / float(IMAGE_SIZE))
print("box coords:")
print(boxes_)
print('*' * 30)
print("scores:")
print(scores_)
print('*' * 30)
print("labels:")
print(labels_)
print("runtime:")
print(t1 - t0)
for i in range(len(boxes_)):
x0, y0, x1, y1 = boxes_[i]
plot_one_box(img_ori, [x0, y0, x1, y1],
label=classes[labels_[i]],
color=color_table[labels_[i]])
#cv2.imshow('Detection result', img_ori)
cv2.imwrite('pb_result.jpg', img_ori)
#cv2.waitKey(0)
num_samples = 50
t0 = time.time()
for i in range(num_samples):
boxes_, scores_, labels_ = sess.run(
[boxes, scores, labels],
feed_dict={
input_x: np.reshape(img,
[-1, IMAGE_SIZE, IMAGE_SIZE, 3])
})
t1 = time.time()
print('Average runtime: %f seconds' %
(float(t1 - t0) / num_samples))
def recognize(jpg_path, pb_file_path):
anchors = parse_anchors("./data/yolo_anchors.txt")
classes = read_class_names("./data/coco.names")
num_class = len(classes)
color_table = get_color_table(num_class)
img_ori = cv2.imread(jpg_path)
height_ori, width_ori = img_ori.shape[:2]
img = cv2.resize(img_ori, tuple([IMAGE_SIZE, IMAGE_SIZE]))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = np.asarray(img, np.float32)
# img = img[np.newaxis, :] / 255.
# img_resized = np.reshape(img, [-1, IMAGE_SIZE, IMAGE_SIZE, 3])
with tf.Graph().as_default():
tf_config = tf.ConfigProto()
tf_config.gpu_options.allow_growth = True
with tf.Session(config=tf_config) as sess:
print("Load TRT_Graph File ...")
with open(pb_file_path, "rb") as f:
output_graph_def = tf.GraphDef()
output_graph_def.ParseFromString(f.read())
print("Finished")
input_name = "import/Placeholder"
output_name1 = "import/yolov3/yolov3_head/feature_map_1"
output_name2 = "import/yolov3/yolov3_head/feature_map_2"
output_name3 = "import/yolov3/yolov3_head/feature_map_3"
output_names = [output_name1, output_name2, output_name3]
yolo_model = yolov3(num_class, anchors)
print("Import TRT Graph ...")
output_node = tf.import_graph_def(
output_graph_def,
return_elements=[
"yolov3/yolov3_head/feature_map_1",
"yolov3/yolov3_head/feature_map_2",
"yolov3/yolov3_head/feature_map_3"
])
print("Finished")
# for op in tf.get_default_graph().as_graph_def().node:
# print(op.name)
tf_input = sess.graph.get_tensor_by_name(input_name + ':0')
feature_map_1 = sess.graph.get_tensor_by_name(output_name1 + ":0")
feature_map_2 = sess.graph.get_tensor_by_name(output_name2 + ":0")
feature_map_3 = sess.graph.get_tensor_by_name(output_name3 + ":0")
features = feature_map_1, feature_map_2, feature_map_3
sess.run(output_node, feed_dict={tf_input: img[None, ...]})
print("1111111")
yolo_model.pb_forward(tf_input)
pred_boxes, pred_confs, pred_probs = yolo_model.predict(features)
pred_scores = pred_confs * pred_probs
print("Detection ......")
boxes, scores, labels = gpu_nms(pred_boxes,
pred_scores,
num_class,
max_boxes=30,
score_thresh=0.4,
iou_thresh=0.5)
boxes_, scores_, labels_ = sess.run(
[boxes, scores, labels], feed_dict={tf_input: img[None, ...]})
# rescale the coordinates to the original image
boxes_[:, 0] *= (width_ori / float(IMAGE_SIZE))
boxes_[:, 2] *= (width_ori / float(IMAGE_SIZE))
boxes_[:, 1] *= (height_ori / float(IMAGE_SIZE))
boxes_[:, 3] *= (height_ori / float(IMAGE_SIZE))
print("box coords:")
print(boxes_)
print('*' * 30)
print("scores:")
print(scores_)
print('*' * 30)
print("labels:")
print(labels_)
for i in range(len(boxes_)):
x0, y0, x1, y1 = boxes_[i]
plot_one_box(img_ori, [x0, y0, x1, y1],
label=classes[labels_[i]],
color=color_table[labels_[i]])
# cv2.imshow('Detection result', img_ori)
cv2.imwrite('detection_result.jpg', img_ori)
def estimatePose():
parser = argparse.ArgumentParser(
description="YOLO-V3 video test procedure.")
# parser.add_argument("input_video", type=str,
# help="The path of the input video.")
parser.add_argument("--anchor_path",
type=str,
default="./data/yolo_anchors.txt",
help="The path of the anchor txt file.")
parser.add_argument(
"--new_size",
nargs='*',
type=int,
default=[416, 416],
help=
"Resize the input image with `new_size`, size format: [width, height]")
parser.add_argument("--letterbox_resize",
type=lambda x: (str(x).lower() == 'true'),
default=True,
help="Whether to use the letterbox resize.")
parser.add_argument("--class_name_path",
type=str,
default="./data/my_data/YOLOPose.names",
help="The path of the class names.")
parser.add_argument("--restore_path",
type=str,
default="./data/pose_weights/lunge_best",
help="The path of the weights to restore.")
parser.add_argument("--save_video",
type=lambda x: (str(x).lower() == 'true'),
default=True,
help="Whether to save the video detection results.")
args = parser.parse_args()
args.anchors = parse_anchors(args.anchor_path)
args.classes = read_class_names(args.class_name_path)
args.num_class = len(args.classes)
color_table = get_color_table(args.num_class)
# vid = cv2.VideoCapture(args.input_video)
vid = cv2.VideoCapture('./data/demo/lunge_03.mp4')
# vid = cv2.VideoCapture(r'C:\Users\soma\SMART_Referee\SMART_Referee_DL\data\lunge\video\lunge_03.mp4')
video_frame_cnt = int(vid.get(7))
video_width = int(vid.get(3))
video_height = int(vid.get(4))
video_fps = int(vid.get(5))
trainer_pose = pd.read_csv('./data/ground_truth/output_right.csv',
header=None)
trainer_pose = trainer_pose.loc[:, [
0, 1, 2, 3, 4, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28
]]
pca_df = trainer_pose.loc[:, [
1, 2, 3, 4, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28
]]
pca_df.loc[:,
[c for c in pca_df.columns if c % 2 ==
1]] = pca_df.loc[:, [c for c in pca_df.columns
if c % 2 == 1]] * video_width / 416
pca_df.loc[:,
[c for c in pca_df.columns if c % 2 ==
0]] = pca_df.loc[:, [c for c in pca_df.columns
if c % 2 == 0]] * video_height / 416
pca_df = pca_df.astype(int)
pca_df = pca_df.replace(0, np.nan)
pca_df = pca_df.dropna()
pca_df.describe()
pca = PCA(n_components=1)
pca.fit(pca_df)
size = [video_width, video_height]
list_p = []
waist_err = 0
critical_point = 0
past_idx = 0
startTrig = 0
cntdown = 90
t = 0
TRLEN = len(trainer_pose)
modify_ankle = pca_df.iloc[0, :].values
base_rect = [(int(video_width / 4), int(video_height / 10)),
(int(video_width * 3 / 4), int(video_height * 19 / 20))]
c_knee = c_waist = c_speed = 0
if args.save_video:
fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
videoWriter = cv2.VideoWriter('video_result.mp4', fourcc, video_fps,
(video_width, video_height))
with tf.Session() as sess:
input_data = tf.placeholder(tf.float32,
[1, args.new_size[1], args.new_size[0], 3],
name='input_data')
yolo_model = yolov3(args.num_class, args.anchors)
with tf.variable_scope('yolov3'):
pred_feature_maps = yolo_model.forward(input_data, False)
pred_boxes, pred_confs, pred_probs = yolo_model.predict(
pred_feature_maps)
pred_scores = pred_confs * pred_probs
boxes, scores, labels = gpu_nms(pred_boxes,
pred_scores,
args.num_class,
max_boxes=200,
score_thresh=0.3,
nms_thresh=0.45)
saver = tf.train.Saver()
saver.restore(sess, args.restore_path)
for i in range(video_frame_cnt):
ret, img_ori = vid.read()
if args.letterbox_resize:
img, resize_ratio, dw, dh = letterbox_resize(
img_ori, args.new_size[0], args.new_size[1])
else:
height_ori, width_ori = img_ori.shape[:2]
img = cv2.resize(img_ori, tuple(args.new_size))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = np.asarray(img, np.float32)
img = img[np.newaxis, :] / 255.
start_time = time.time()
boxes_, scores_, labels_ = sess.run([boxes, scores, labels],
feed_dict={input_data: img})
# rescale the coordinates to the original image
if args.letterbox_resize:
boxes_[:, [0, 2]] = (boxes_[:, [0, 2]] - dw) / resize_ratio
boxes_[:, [1, 3]] = (boxes_[:, [1, 3]] - dh) / resize_ratio
else:
boxes_[:, [0, 2]] *= (width_ori / float(args.new_size[0]))
boxes_[:, [1, 3]] *= (height_ori / float(args.new_size[1]))
people_pose = get_people_pose(boxes_, labels_,
base_rect) # list-dict
people_pose = np.array([p[1] for p in people_pose[0]
]).flatten() # dict-tuple -> list
people_pose = people_pose[[
0, 1, 2, 3, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27
]]
# Start Trigger
if startTrig == 2:
pass
elif startTrig == 0: # start
# 기준 박스
cv2.rectangle(img_ori, base_rect[0], base_rect[1], (0, 0, 255),
2)
if isInBox(people_pose, base_rect[0], base_rect[1]):
# t_resize_pose = resize_pose(people_pose, trainer_pose.iloc[0, 1:].values)
t_resize_pose = resize_pose(people_pose,
pca_df.iloc[0, :].values)
img_ori = draw_ground_truth(img_ori, t_resize_pose)
# img_ori = draw_ground_truth(img_ori, pca_df.iloc[0, :].values)
startTrig = isStart(people_pose,
trainer_pose.iloc[0, 1:].values, size)
cv2.imshow('image', img_ori)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
continue
else:
print("박스안에 들어와주세요!!")
continue
elif startTrig == 1:
img_ori = draw_ground_truth(img_ori, pca_df.iloc[0, :].values)
cv2.putText(img_ori, str(int(cntdown / 30)), (100, 300),
cv2.FONT_HERSHEY_SIMPLEX, 10, (255, 0, 0), 10)
cv2.imshow('image', img_ori)
cntdown -= 1
if cntdown == 0:
startTrig = 2
if cv2.waitKey(1) & 0xFF == ord('q'):
break
continue
'''check ankle : 편차 40이상 발생시 전에 값 으로 업데이트'''
people_pose = check_ankle(list_p, people_pose, modify_ankle, size)
# f = open('user.csv', 'a', encoding='utf-8', newline='')
# wr = csv.writer(f)
# wr.writerow(people_pose)
# ground truth 그리기
list_p.append(people_pose)
img_ori = draw_ground_truth(img_ori, pca_df.iloc[t, :].values)
if check_waist(people_pose):
waist_err += 1
if waist_err is 60: # waist_err는 60번 틀리면 피드백함
feedback_waist()
c_waist += 1
waist_err = 0
if trainer_pose.iloc[t, 0] == 1: # t는 특정 시점 + i frame
critical_point += 1
if critical_point % 2 == 0:
my_pose = makeMypose_df(list_p)
c_speed = check_speed(
my_pose, trainer_pose.iloc[past_idx:t + 1, 1:], pca,
c_speed)
c_knee = check_knee(people_pose, c_knee)
modify_ankle = list_p[-1]
list_p = []
past_idx = t
t += 1
if t == TRLEN:
break
# img_ori = draw_body(img_ori, boxes_, labels_)
# for i in range(len(boxes_)):
# x0, y0, x1, y1 = boxes_[i]
# plot_one_box(img_ori, [x0, y0, x1, y1], label=args.classes[labels_[i]] + ', {:.2f}%'.format(scores_[i] * 100), color=color_table[labels_[i]])
# 사용자 자세 그리기
# img_ori = draw_truth(img_ori, people_pose)
end_time = time.time()
cv2.putText(img_ori,
'{:.2f}ms'.format((end_time - start_time) * 1000),
(40, 40),
0,
fontScale=1,
color=(0, 255, 0),
thickness=2)
cv2.imshow('image', img_ori)
if args.save_video:
videoWriter.write(img_ori)
if cv2.waitKey(1) & 0xFF == ord('q'):
cv2.destroyAllWindows()
break
vid.release()
cv2.destroyAllWindows()
if args.save_video:
videoWriter.release()
f = open('./data/score/result.csv', 'a', encoding='utf-8', newline='')
wr = csv.writer(f)
d = datetime.today().strftime("%Y/%m/%d")
t = datetime.today().strftime("%H:%M:%S")
wr.writerow([d, t, c_knee, c_waist, c_speed])
from model import yolov3
lines_gender = ['female', 'male']
input_video = "./data/demo_data/test.mp4"
anchor_path = "./data/yolo_anchors.txt"
new_size = [416, 416]
class_name_path = "./data/coco.names"
restore_path = "./data/darknet_weights/yolov3.ckpt"
anchors = parse_anchors(anchor_path)
classes = read_class_names(class_name_path)
num_class = len(classes)
color_table = get_color_table(num_class)
vid = cv2.VideoCapture(input_video)
video_frame_cnt = int(vid.get(7))
video_width = int(vid.get(3))
video_height = int(vid.get(4))
video_fps = int(vid.get(5))
fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
videoWriter = cv2.VideoWriter('video_result.mp4', fourcc, video_fps,
(video_width, video_height))
with tf.Session() as sess:
input_data = tf.placeholder(tf.float32, [1, new_size[1], new_size[0], 3],
name='input_data')
yolo_model = yolov3(num_class, anchors)
def recognize(jpg_path, pb_file_path):
anchors = parse_anchors("./data/yolo_anchors.txt")
classes = read_class_names("./data/coco.names")
num_class = len(classes)
color_table = get_color_table(num_class)
img_ori = cv2.imread(jpg_path)
height_ori, width_ori = img_ori.shape[:2]
img = cv2.resize(img_ori, tuple([IMAGE_SIZE, IMAGE_SIZE]))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = np.asarray(img, np.float32)
img = img[np.newaxis, :] / 255.
with tf.Graph().as_default():
output_graph_def = tf.GraphDef()
with open(pb_file_path, "rb") as f:
output_graph_def.ParseFromString(f.read())
_ = tf.import_graph_def(output_graph_def, name="")
tf_config = tf.ConfigProto()
tf_config.gpu_options.allow_growth = True
with tf.Session(config=tf_config) as sess:
init = tf.global_variables_initializer()
sess.run(init)
input_name = "Placeholder"
output_name1 = "yolov3/yolov3_head/feature_map_1"
output_name2 = "yolov3/yolov3_head/feature_map_2"
output_name3 = "yolov3/yolov3_head/feature_map_3"
output_names = [output_name1, output_name2, output_name3]
yolo_model = yolov3(num_class, anchors)
input_data = tf.placeholder(tf.float32,
[1, IMAGE_SIZE, IMAGE_SIZE, 3],
name='input_data')
trt_graph = trt.create_inference_graph(
input_graph_def=output_graph_def,
outputs=output_names,
max_batch_size=1,
max_workspace_size_bytes=1 << 25,
precision_mode='FP16',
minimum_segment_size=5)
with open('./data/yolov3_trt.pb', 'wb') as f:
f.write(trt_graph.SerializeToString())
tf.import_graph_def(trt_graph, name='')
tf_input = sess.graph.get_tensor_by_name(input_name + ':0')
feature_map_1 = sess.graph.get_tensor_by_name(output_name1 + ":0")
feature_map_2 = sess.graph.get_tensor_by_name(output_name2 + ":0")
feature_map_3 = sess.graph.get_tensor_by_name(output_name3 + ":0")
features = feature_map_1, feature_map_2, feature_map_3
# tf_scores = tf_sess.graph.get_tensor_by_name('detection_scores:0')
# tf_boxes = tf_sess.graph.get_tensor_by_name('detection_boxes:0')
# tf_classes = tf_sess.graph.get_tensor_by_name('detection_classes:0')
# tf_num_detections = tf_sess.graph.get_tensor_by_name('num_detections:0')
# features = sess.run(features, feed_dict={tf_input:np.reshape(img, [-1, IMAGE_SIZE, IMAGE_SIZE, 3])})
# feature1, feature2, feature3 = features
# feature1 = tf.convert_to_tensor(feature1)
# feature2 = tf.convert_to_tensor(feature2)
# feature3 = tf.convert_to_tensor(feature3)
# features = feature1, feature2, feature3
yolo_model.pb_forward(input_data)
pred_boxes, pred_confs, pred_probs = yolo_model.predict(features)
pred_scores = pred_confs * pred_probs
boxes, scores, labels = gpu_nms(pred_boxes,
pred_scores,
num_class,
max_boxes=30,
score_thresh=0.4,
iou_thresh=0.5)
boxes_, scores_, labels_ = sess.run([boxes, scores, labels],
feed_dict={input_data: img})
# rescale the coordinates to the original image
boxes_[:, 0] *= (width_ori / float(IMAGE_SIZE))
boxes_[:, 2] *= (width_ori / float(IMAGE_SIZE))
boxes_[:, 1] *= (height_ori / float(IMAGE_SIZE))
boxes_[:, 3] *= (height_ori / float(IMAGE_SIZE))
print("box coords:")
print(boxes_)
print('*' * 30)
print("scores:")
print(scores_)
print('*' * 30)
print("labels:")
print(labels_)
for i in range(len(boxes_)):
x0, y0, x1, y1 = boxes_[i]
plot_one_box(img_ori, [x0, y0, x1, y1],
label=classes[labels_[i]],
color=color_table[labels_[i]])
# cv2.imshow('Detection result', img_ori)
cv2.imwrite('detection_result.jpg', img_ori)
type=str,
default="./data/my_data/voc.names",
help="The path of the class names.")
parser.add_argument(
"--restore_path",
type=str,
default=
"/media/ubutnu/fc1a3be7-9b03-427e-9cc9-c4b242cefbff/YOLOv3_TensorFlow/checkpoint/model-epoch_90_step_175083_loss_0.4213_lr_1e-05",
help="The path of the weights to restore.")
args = parser.parse_args()
args.anchors = parse_anchors(args.anchor_path)
args.classes = read_class_names(args.class_name_path)
args.num_class = len(args.classes)
color_table = get_color_table(args.num_class)
img_ori = cv2.imread(args.input_image)
if args.letterbox_resize:
img, resize_ratio, dw, dh = letterbox_resize(img_ori, args.new_size[0],
args.new_size[1])
else:
height_ori, width_ori = img_ori.shape[:2]
img = cv2.resize(img_ori, tuple(args.new_size))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = np.asarray(img, np.float32)
img = img[np.newaxis, :] / 255.
with tf.Session() as sess:
input_data = tf.placeholder(tf.float32,
[1, args.new_size[1], args.new_size[0], 3],
# -*- coding:utf-8 -*- from utils.misc_utils import parse_anchors, read_class_names from utils.plot_utils import get_color_table """ 检测配置 """ detect_object = 'img' # 默认检测对象 input_image = './data/demo_data/person.jpg' # 默认图片路径 input_video = './data/demo_data/video_demo.mp4' # 默认视频路径 output_image = './data/demo_data/result/result.jpg' # 保存图片路径 output_video = './data/demo_data/result/result.mp4' # 保存视频路径 anchor_path = './data/yolo_anchors.txt' # anchor 文件路径 anchors = parse_anchors(anchor_path) # anchor内容 weight_path = './data/darknet_weights/yolov3.ckpt' # weights路径 class_name_path = './data/coco.names' # 类别文件路径 classes = read_class_names(class_name_path) # 类别文件list num_class = len(classes) # 类别数量 new_size = [416, 416] # 图片改变后的大小 use_letterbox_resize = True # 是否使用letterbox color_table = get_color_table(num_class) # 根据类别数生成颜色列表
