def main(_argv):
# import weights
yolo = YoloV3(classes=FLAGS.num_classes)
yolo.load_weights(FLAGS.weights)
logging.info("weights loaded")
# Import classes
class_names = [c.strip() for c in open(FLAGS.classes).readlines()]
logging.info("classes loaded")
# list of time for procces on each frame
times = []
# try to load webcam or a video file
try:
vid = cv2.VideoCapture(int(FLAGS.video))
except:
vid = cv2.VideoCapture(FLAGS.video)
out = None
if FLAGS.output:
#by default VideoCapture returns float instead of int
width = int(vid.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = int(vid.get(cv2.CAP_PROP_FPS))
codec = cv2.VideoWriter_fourcc(*FLAGS.output_format)
out = cv2.VideoWriter(FLAGS.output, codec, fps, (width, height))
while True:
_, img = vid.read()
if img is None:
logging.info("Empty Frame")
time.sleep(0.1)
img_in = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img_in = tf.expand_dims(img_in, 0)
img_in = transform_images(img_in, FLAGS.size)
t1 = time.time()
boxes, scores, classes, nums = yolo.predict(img_in)
for i in nums:
print(i)
t2 = time.time()
times.append(t2 - t1)
times = times[-20:]
img = draw_outputs(img, (boxes, scores, classes, nums), class_names)
img = cv2.putText(
img, "Time: {:.2f}ms".format(sum(times) / len(times) * 1000),
(0, 30), cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, (1, 0, 255), 2)
if FLAGS.output:
out.write(img)
cv2.imshow('Output', img)
if cv2.waitKey(1) == ord('q'):
break
vid.release()
cv2.destroyAllWindows()
def main(_argv):
# 读取 类别 并生成 list
class_names = [c.strip() for c in open(FLAGS.classes).readlines()]
logging.info('classes loaded')
# 读取 TF record格式 的数据集并 随机排序
dataset = load_tfrecord_dataset(FLAGS.dataset, FLAGS.classes, FLAGS.size)
dataset = dataset.shuffle(512)
# 读取第一张图片并可视化并保存
for image, labels in dataset.take(1):
boxes = []
scores = []
classes = []
for x1, y1, x2, y2, label in labels:
if x1 == 0 and x2 == 0:
continue
boxes.append((x1, y1, x2, y2))
scores.append(1)
classes.append(label)
nums = [len(boxes)]
boxes = [boxes]
scores = [scores]
classes = [classes]
logging.info('labels:')
for i in range(nums[0]):
logging.info('\t{}, {}, {}'.format(class_names[int(classes[0][i])],
np.array(scores[0][i]),
np.array(boxes[0][i])))
img = cv2.cvtColor(image.numpy(), cv2.COLOR_RGB2BGR)
img = draw_outputs(img, (boxes, scores, classes, nums), class_names)
cv2.imwrite(FLAGS.output, img)
logging.info('output saved to: {}'.format(FLAGS.output))
def main(_argv):
if FLAGS.tiny:
yolo = YoloV3Tiny()
else:
yolo = YoloV3()
yolo.load_weights(FLAGS.weights)
logging.info('weights loaded from {}'.format(FLAGS.weights))
class_names = voc_label_map_list
img = tf.image.decode_image(open(FLAGS.image, 'rb').read(), channels=3)
img = tf.expand_dims(img, 0)
img = transform_images(img, FLAGS.size)
t1 = time.time()
boxes, scores, classes, nums = yolo(img)
t2 = time.time()
logging.info('time: {}'.format(t2 - t1))
logging.info('detections:')
for i in range(nums[0]):
logging.info('\t{}, {}, {}'.format(class_names[int(classes[0][i])],
np.array(scores[0][i]),
np.array(boxes[0][i])))
img = cv2.imread(FLAGS.image)
img = draw_outputs(img, (boxes, scores, classes, nums), class_names)
cv2.imshow('rr', img)
cv2.waitKey(0)
def main(_argv):
if FLAGS.tiny:
yolo = YoloV3Tiny()
else:
yolo = YoloV3()
yolo.load_weights(FLAGS.weights)
logging.info('weights loaded')
class_names = get_coco_names(FLAGS.classes)
logging.info('classes loaded')
img = tf.image.decode_image(open(FLAGS.image, 'rb').read(), channels=3)
img = tf.expand_dims(img, 0)
img = transform_images(img, FLAGS.size)
t1 = time.time()
boxes, scores, classes, nums = yolo(img)
t2 = time.time()
logging.info('time: {}'.format(t2 - t1))
logging.info('detections:')
for i in range(nums[0]):
logging.info('\t{}, {}, {}'.format(class_names[int(classes[0][i])],
np.array(scores[0][i]),
np.array(boxes[0][i])))
img = cv2.imread(FLAGS.image)
img = draw_outputs(img, (boxes, scores, classes, nums), class_names)
cv2.imwrite(FLAGS.output, img)
logging.info('output saved to: {}'.format(FLAGS.output))
def test_model(model, img, epoch, batch):
os.makedirs('results', exist_ok=True)
class_names = get_coco_names('dataset/coco.prototxt')
assert isinstance(img, tf.Tensor), 'img must be an tensor of image'
out = model(tf.expand_dims(img, axis=0), training=False)
# the output is all 0??? why??
# drawing out on test image
# print(out)
# print(img.numpy().shape)
img = np.array(img.numpy()*255., dtype=np.uint8)
# convert BGR to RGB
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
res = draw_outputs(img, out, class_names)
cv2.imwrite('results/{}_{}_ori.png'.format(epoch, batch), img)
cv2.imwrite('results/{}_{}_pred.png'.format(epoch, batch), res)
def main(_argv):
class_names = [c.strip() for c in open(FLAGS.classes).readlines()]
logging.info('classes loaded')
if FLAGS.tiny:
yolo = YoloV3Tiny(classes=len(class_names))
else:
yolo = YoloV3(classes=len(class_names))
yolo.load_weights(FLAGS.weights).expect_partial()
logging.info('weights loaded')
if not FLAGS.image:
login.error('Detect image must be specified')
return 1
elif os.path.isdir(FLAGS.image):
detect_images = [os.path.join(FLAGS.image, x) for x in os.listdir(FLAGS.image) if x[-3:] == 'jpg']
else:
detect_images = [FLAGS.image]
for image in detect_images:
img_raw = tf.image.decode_image(
open(image, 'rb').read(), channels=3)
img = tf.expand_dims(img_raw, 0)
img = transform_images(img, FLAGS.size)
t1 = time.time()
boxes, scores, classes, nums = yolo(img)
t2 = time.time()
logging.info('time: {}'.format(t2 - t1))
logging.info('detections:')
if nums[0].numpy() == 0:
continue
for i in range(nums[0]):
logging.info('\t{}, {}, {}'.format(class_names[int(classes[0][i])],
np.array(scores[0][i]),
np.array(boxes[0][i])))
img = cv2.cvtColor(img_raw.numpy(), cv2.COLOR_RGB2BGR)
img = draw_outputs(img, (boxes, scores, classes, nums), class_names)
cv2.imshow(image, img)
cv2.waitKey(0)
def main(_argv):
if FLAGS.tiny:
yolo = YoloV3Tiny()
else:
yolo = YoloV3()
yolo.load_weights(FLAGS.weights)
logging.info('weights loaded')
class_names = [c.strip() for c in open(FLAGS.classes).readlines()]
logging.info('classes loaded')
times = []
try:
vid = cv2.VideoCapture(int(FLAGS.video))
except:
vid = cv2.VideoCapture(FLAGS.video)
while True:
_, img = vid.read()
if img is None:
logging.warning("Empty Frame")
time.sleep(0.1)
continue
img_in = tf.expand_dims(img, 0)
img_in = transform_images(img_in, FLAGS.size)
t1 = time.time()
boxes, scores, classes, nums = yolo.predict(img_in)
t2 = time.time()
times.append(t2 - t1)
times = times[-20:]
img = draw_outputs(img, (boxes, scores, classes, nums), class_names)
img = cv2.putText(
img, "Time: {:.2f}ms".format(sum(times) / len(times) * 1000),
(0, 30), cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, (0, 0, 255), 2)
cv2.imshow('output', img)
if cv2.waitKey(1) == ord('q'):
break
cv2.destroyAllWindows()
def main(_argv):
# 打开内存增长
physical_devices = tf.config.experimental.list_physical_devices('GPU')
for physical_device in physical_devices:
tf.config.experimental.set_memory_growth(physical_device, True)
# 创建网络模型
if FLAGS.tiny:
yolo = YoloV3Tiny(classes=FLAGS.num_classes)
else:
yolo = YoloV3(classes=FLAGS.num_classes)
# 加载权重
yolo.load_weights(FLAGS.weights).expect_partial()
logging.info('weights loaded')
# 加载类别列表
class_names = [c.strip() for c in open(FLAGS.classes).readlines()]
logging.info('classes loaded')
# 读取图片数据
if FLAGS.tfrecord:
dataset = load_tfrecord_dataset(FLAGS.tfrecord, FLAGS.classes,
FLAGS.size)
dataset = dataset.shuffle(512)
img_raw, _label = next(iter(dataset.take(1)))
else:
img_raw = tf.image.decode_image(open(FLAGS.image, 'rb').read(),
channels=3)
# 调整图片
img = tf.expand_dims(img_raw, 0)
img = transform_images(img, FLAGS.size)
# 对图片进行目标检测
t1 = time.time()
boxes, scores, classes, nums = yolo(img)
t2 = time.time()
logging.info('time: {}'.format(t2 - t1))
# 可视化输出
logging.info('detections:')
for i in range(nums[0]):
logging.info('\t{}, {}, {}'.format(class_names[int(classes[0][i])],
np.array(scores[0][i]),
np.array(boxes[0][i])))
img = cv2.cvtColor(img_raw.numpy(), cv2.COLOR_RGB2BGR)
img = draw_outputs(img, (boxes, scores, classes, nums), class_names)
cv2.imwrite(FLAGS.output, img)
logging.info('output saved to: {}'.format(FLAGS.output))
def test_model(img, epoch, batch):
os.makedirs('results', exist_ok=True)
class_names = get_coco_names('dataset/coco.prototxt')
assert isinstance(img, tf.Tensor), 'img must be an tensor of image'
model = YoloV3(FLAGS.size, training=False)
latest_cp = tf.train.latest_checkpoint(os.path.dirname(FLAGS.weights))
if latest_cp:
model.load_weights(latest_cp)
logging.info('testing model from: {}'.format(latest_cp))
else:
logging.info('can not find latest model')
out = model(tf.expand_dims(img, axis=0))
# drawing out on test image
# print(out)
img = np.array(img.numpy() * 255., dtype=np.uint8)
res = draw_outputs(img, out, class_names)
cv2.imwrite('results/img_{}_{}.png'.format(epoch, batch), img)
cv2.imwrite('results/predict_{}_{}.png'.format(epoch, batch), res)
del model
def main(_argv):
class_names = [c.strip() for c in open(FLAGS.classes).readlines()]
logging.info('classes loaded')
if not FLAGS.images or not FLAGS.labels:
logging.error('Image directory and label file must be sepcified.')
return
dataset = load_yolo_dataset(FLAGS.images, FLAGS.labels, FLAGS.size)
dataset = dataset.shuffle(512)
for image, labels in dataset:
boxes = []
scores = []
classes = []
for x1, y1, x2, y2, label in labels:
if x1.numpy() == 0 and x2.numpy() == 0:
continue
boxes.append((x1, y1, x2, y2))
scores.append(1)
classes.append(label)
nums = [len(boxes)]
boxes = [boxes]
scores = [scores]
classes = [classes]
logging.info('labels:')
for i in range(nums[0]):
logging.info('\t{}, {}, {}'.format(class_names[int(classes[0][i])],
np.array(scores[0][i]),
np.array(boxes[0][i])))
img = cv2.cvtColor(np.uint8(image.numpy()), cv2.COLOR_RGB2BGR)
img = draw_outputs(img, (boxes, scores, classes, nums), class_names)
cv2.imshow('Dataset Visualizer', img)
key = cv2.waitKey(0)
if key == 27:
return
def main(_argv):
physical_devices = tf.config.experimental.list_physical_devices('GPU')
for physical_device in physical_devices:
tf.config.experimental.set_memory_growth(physical_device, True)
if FLAGS.tiny:
yolo = YoloV3Tiny(classes=FLAGS.num_classes)
else:
yolo = YoloV3(classes=FLAGS.num_classes)
yolo.load_weights(FLAGS.weights)
logging.info('weights loaded')
class_names = [c.strip() for c in open(FLAGS.classes).readlines()]
logging.info('classes loaded')
times = []
try:
vid = cv2.VideoCapture(int(FLAGS.video))
except:
vid = cv2.VideoCapture(FLAGS.video)
out = None
if FLAGS.output:
# by default VideoCapture returns float instead of int
width = int(vid.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = int(vid.get(cv2.CAP_PROP_FPS))
codec = cv2.VideoWriter_fourcc(*FLAGS.output_format)
out = cv2.VideoWriter(FLAGS.output, codec, fps, (width, height))
while True:
_, img = vid.read()
if img is None:
logging.warning("Empty Frame")
time.sleep(0.1)
continue
img_in = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img_in = tf.expand_dims(img_in, 0)
img_in = transform_images(img_in, FLAGS.size)
t1 = time.time()
boxes, scores, classes, nums = yolo.predict(img_in)
t2 = time.time()
times.append(t2 - t1)
times = times[-20:]
img = draw_outputs(img, (boxes, scores, classes, nums), class_names)
img = cv2.putText(
img, "Time: {:.2f}ms".format(sum(times) / len(times) * 1000),
(0, 30), cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, (0, 0, 255), 2)
if FLAGS.output:
out.write(img)
cv2.imshow('output', img)
if cv2.waitKey(1) == ord('q'):
break
cv2.destroyAllWindows()
