import cv2 from tensorflow.keras import Input, Model from darknet import darknet_base from predict import predict inputs = Input(shape=(None, None, 3)) outputs, config = darknet_base(inputs) model = Model(inputs, outputs) vidcap = cv2.VideoCapture('data/demo.mp4') success, image = vidcap.read() size = (int(vidcap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(vidcap.get(cv2.CAP_PROP_FRAME_HEIGHT))) fourcc = cv2.VideoWriter_fourcc(*'DIVX') out = cv2.VideoWriter('output.avi', fourcc, 20.0, size) while success: output_image = predict(model, image, config) out.write(output_image) success, image = vidcap.read() vidcap.release() out.release()
# NOTE: Because of a bug in TensorFlow, this should be run in the console # NOTE: python tflite.py import os import tensorflow as tf from tensorflow.contrib.lite.python import lite from tensorflow.keras import Input, Model from darknet import darknet_base from definitions import ROOT_DIR inputs = Input(shape=(None, None, 3)) # NOTE: Here, we do not include the YOLO head because TFLite does not # NOTE: support custom layers yet. Therefore, we'll need to implement # NOTE: the YOLO head ourselves. outputs, config = darknet_base(inputs, include_yolo_head=False) model = Model(inputs, outputs) model_path = os.path.join(ROOT_DIR, 'model', 'yolov3.h5') tf.keras.models.save_model(model, model_path, overwrite=True) # Sanity check to see if model loads properly # NOTE: See https://github.com/keras-team/keras/issues/4609#issuecomment-329292173 # on why we have to pass in `tf: tf` in `custom_objects` model = tf.keras.models.load_model(model_path, custom_objects={'tf': tf}) converter = lite.TocoConverter.from_keras_model_file( model_path, input_shapes={'input_1': [1, config['width'], config['height'], 3]}) converter.post_training_quantize = True
import time import cv2 from tensorflow.keras import Input, Model from darknet import darknet_base from predict import predict, predict_with_yolo_head INCLUDE_YOLO_HEAD = True stream = cv2.VideoCapture(0) inputs = Input(shape=(None, None, 3)) outputs, config = darknet_base(inputs, include_yolo_head=INCLUDE_YOLO_HEAD) model = Model(inputs, outputs) while True: # Capture frame-by-frame grabbed, frame = stream.read() if not grabbed: break # Run detection start = time.time() if INCLUDE_YOLO_HEAD: output_image = predict(model, frame, config) else: output_image = predict_with_yolo_head(model, frame, config, confidence=0.3,
from darknet import darknet_base import argparse from tensorflow.keras import Input, Model import tensorflow as tf parser = argparse.ArgumentParser(description='Keras to TF-Lite converter') parser.add_argument('--h5_path', default='../../pretrain_models/yolov3/yolov3.h5', help='Path to Darknet h5 weights file.') parser.add_argument('--output_path', default='../../pretrain_models/yolov3/yolov3_weights.tflite', help='Path to output Keras model file.') args = parser.parse_args() inputs = Input(shape=(None, None, 3)) outputs, config = darknet_base( inputs, yolo_version='yolov3', data_dir='../../pretrain_models/yolov3', include_yolo_head=False) model = Model(inputs, outputs) model_path = args.h5_path tf.keras.models.save_model(model, model_path, overwrite=True) model.summary() # Sanity check to see if model loads properly # NOTE: See https://github.com/keras-team/keras/issues/4609#issuecomment-329292173 # on why we have to pass in `tf: tf` in `custom_objects` model = tf.keras.models.load_model(model_path, custom_objects={'tf': tf}) converter = tf.contrib.lite.TFLiteConverter.from_keras_model_file(model_path, input_shapes={'input_1': [1, config['width'], config['height'], 3]}) tflite_model = converter.convert()
parser = argparse.ArgumentParser(description='Keras to TF-Lite converter') parser.add_argument('--h5_path', default='../../pretrain_models/yolov3_tflite.h5', help='Path to Darknet h5 weights file.') parser.add_argument('--output_path', default='../../pretrain_models/yolov3.tflite', help='Path to output Keras model file.') args = parser.parse_args() inputs = Input(shape=(None, None, 3)) # NOTE: Here, we do not include the YOLO head because TFLite does not # NOTE: support custom layers yet. Therefore, we'll need to implement # NOTE: the YOLO head ourselves. outputs, config = darknet_base( inputs, YOLO_VERSION='yolov3', data_dir='weights', include_yolo_head=False) model = Model(inputs, outputs) model_path = args.h5_path tf.keras.models.save_model(model, model_path, overwrite=True) model.summary() # Sanity check to see if model loads properly # NOTE: See https://github.com/keras-team/keras/issues/4609#issuecomment-329292173 # on why we have to pass in `tf: tf` in `custom_objects` model = tf.keras.models.load_model(model_path, custom_objects={'tf': tf}) converter = tf.contrib.lite.TFLiteConverter.from_keras_model_file(model_path, input_shapes={'input_1': [1, config['width'], config['height'], 3]}) #converter.post_training_quantize = True
parser = argparse.ArgumentParser(description='Keras to TF-Lite converter') parser.add_argument('--h5_path', default='weights/yolo-obj-608-water_15000.h5', help='Path to Darknet h5 weights file.') parser.add_argument('--output_path', default='weights/yolo-obj-608-water_15000.tflite', help='Path to output Keras model file.') args = parser.parse_args() inputs = Input(shape=(None, None, 3)) # NOTE: Here, we do not include the YOLO head because TFLite does not # NOTE: support custom layers yet. Therefore, we'll need to implement # NOTE: the YOLO head ourselves. outputs, config = darknet_base(inputs, YOLO_VERSION='yolo-obj-608-water_15000', data_dir='weights', include_yolo_head=False) model = Model(inputs, outputs) model_path = args.h5_path tf.keras.models.save_model(model, model_path, overwrite=True) model.summary() # Sanity check to see if model loads properly # NOTE: See https://github.com/keras-team/keras/issues/4609#issuecomment-329292173 # on why we have to pass in `tf: tf` in `custom_objects` model = tf.keras.models.load_model(model_path, custom_objects={'tf': tf}) converter = tf.contrib.lite.TFLiteConverter.from_keras_model_file( model_path, input_shapes={'input_1': [1, config['width'], config['height'], 3]})