def main(_argv):
input_layer = tf.keras.layers.Input([FLAGS.size, FLAGS.size, 3])
feature_maps = YOLOv3(input_layer)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
# model.summary()
utils.load_weights(model, FLAGS.weights)
test_img = tf.image.decode_image(open(FLAGS.image, 'rb').read(),
channels=3)
img_size = test_img.shape[:2]
test_img = tf.expand_dims(test_img, 0)
test_img = utils.transform_images(test_img, FLAGS.size)
pred_bbox = model.predict(test_img)
pred_bbox = [tf.reshape(x, (-1, tf.shape(x)[-1])) for x in pred_bbox]
pred_bbox = tf.concat(pred_bbox, axis=0)
boxes = utils.postprocess_boxes(pred_bbox, img_size, FLAGS.size, 0.3)
boxes = utils.nms(boxes, 0.45, method='nms')
original_image = cv2.imread(FLAGS.image)
img = utils.draw_outputs(original_image, boxes)
cv2.imwrite(FLAGS.output, img)
def main(_argv):
input_layer = tf.keras.layers.Input([FLAGS.size, FLAGS.size, 3])
feature_maps = YOLOv3(input_layer)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
# model.summary()
utils.load_weights(model, FLAGS.weights)
times = []
try:
vid = cv2.VideoCapture(int(FLAGS.video))
except:
vid = cv2.VideoCapture(FLAGS.video)
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_size = img.shape[:2]
img_in = tf.expand_dims(img, 0)
img_in = utils.transform_images(img_in, FLAGS.size)
t1 = time.time()
pred_bbox = model.predict(img_in)
t2 = time.time()
times.append(t2 - t1)
times = times[-20:]
pred_bbox = [tf.reshape(x, (-1, tf.shape(x)[-1])) for x in pred_bbox]
pred_bbox = tf.concat(pred_bbox, axis=0)
boxes = utils.postprocess_boxes(pred_bbox, img_size, FLAGS.size, 0.3)
boxes = utils.nms(boxes, 0.45, method='nms')
img = utils.draw_outputs(img, boxes)
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()
def load_model(model_name, weight_path, input_size, framework):
assert model_name in ['yolov3_tiny', 'yolov3', 'yolov4']
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
if framework == 'tf':
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
if model_name == 'yolov3_tiny':
feature_maps = YOLOv3_tiny(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode(fm, NUM_CLASS)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
elif model_name == 'yolov3':
feature_maps = YOLOv3(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode(fm, NUM_CLASS)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
elif model_name == 'yolov4':
feature_maps = YOLOv4(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode(fm, NUM_CLASS)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
else:
model = None
raise ValueError
if weight_path.split(".")[-1] == "weights":
if model_name == 'yolov3_tiny':
utils.load_weights_tiny(model, weight_path)
print('load yolo tiny 3')
elif model_name == 'yolov3':
utils.load_weights_v3(model, weight_path)
print('load yolo 3')
elif model_name == 'yolov4':
utils.load_weights(model, weight_path)
print('load yolo 4')
else:
raise ValueError
else:
model.load_weights(weight_path).expect_partial()
print('Restoring weights from: %s ' % weight_path)
return model
def _main(args):
kwargs = {}
classes_path = args.classes_path if args.classes_path else 'data/coco_classes.txt'
yolo = YOLOv3(initial_weights_path=str(args.initial_weights_path),
annotations_path=str(args.annotations_path),
is_training=True,
anchors_path='data/yolo_anchors.txt',
classes_path=classes_path,
log_dir='log',
**kwargs)
yolo.train(use_focal_loss=args.use_focal_loss)
def transfer_tflite(model_name, weight_path, output, input_size):
assert model_name in ['yolov3_tiny', 'yolov3', 'yolov4']
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
if model_name == 'yolov3_tiny':
feature_maps = YOLOv3_tiny(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode(fm, NUM_CLASS)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
elif model_name == 'yolov3':
feature_maps = YOLOv3(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode(fm, NUM_CLASS)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
elif model_name == 'yolov4':
feature_maps = YOLOv4(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode(fm, NUM_CLASS)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
else:
model = None
raise ValueError
if weight_path.split(".")[-1] == "weights":
if model_name == 'yolov3_tiny':
utils.load_weights_tiny(model, weight_path)
elif model_name == ' yolov3':
utils.load_weights_v3(model, weight_path)
elif model_name == 'yolov4':
utils.load_weights(model, weight_path)
else:
raise ValueError
else:
model.load_weights(weight_path).expect_partial()
print('Restoring weights from: %s ... ' % weight_path)
converter = tf.lite.TFLiteConverter.from_keras_model(model)
# converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_SIZE]
tflite_model = converter.convert()
open(output, 'wb').write(tflite_model)
def _main(args):
assert len(args.source_images_dir
) > 0, "source images directory can not be empty!"
assert len(args.output_dir
) > 0, "detected image output directory can not be empty!"
assert len(args.weights_path) > 0, "weights path can not be empty!"
output_dir = args.output_dir if args.output_dir.endswith(
'/') else args.output_dir + '/'
source_images_dir = args.source_images_dir if args.source_images_dir.endswith(
'/') else args.source_images_dir + '/'
os.makedirs(output_dir, exist_ok=True)
classes_path = args.classes_path if args.classes_path else 'data/coco_classes.txt'
kwargs = {}
yolo = YOLOv3(initial_weights_path=str(args.weights_path),
is_training=False,
anchors_path='data/yolo_anchors.txt',
classes_path=classes_path,
log_dir='log',
**kwargs)
source_images = []
print(source_images_dir)
print('######### fetching all filenames ##########')
for (_, _, filename) in walk(source_images_dir):
files = []
for f in filename:
if f.lower().endswith('.jpg') or f.lower().endswith(
'.png') or f.lower().endswith('.jpeg'):
files.append(source_images_dir + f)
source_images.extend(files)
print(source_images)
for source_image_path in source_images:
try:
image = Image.open(source_image_path)
except:
continue
print('######### detecting image with name: {} ##########'.format(
source_image_path))
detected_source_image = yolo.detect_image(image)
print('######### detected image saved to: {} ##########'.format(
output_dir + source_image_path.split('/')[-1]))
detected_source_image.save(
output_dir + source_image_path.split('/')[-1], 'JPEG')
print('######### finishing image detecting ##########')
yolo.close_session()
def detect(image_path, weight_path, input_size):
STRIDES = np.array(cfg.YOLO.STRIDES)
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS, tiny=False)
original_image = cv2.imread(image_path)
original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
original_image_size = original_image.shape[:2]
image_data = utils.image_preprocess(np.copy(original_image),
[input_size, input_size])
image_data = image_data[np.newaxis, ...].astype(np.float32)
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
feature_maps = YOLOv3(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode(fm, NUM_CLASS)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
if weight_path:
weight = np.load(weight_path, allow_pickle=True)
model.set_weights(weight)
print('Restoring weights from: %s ' % weight_path)
pred_bbox = model.predict(image_data)
pred_bbox = utils.postprocess_bbbox(pred_bbox, ANCHORS, STRIDES)
bboxes = utils.postprocess_boxes(pred_bbox, original_image_size,
input_size, 0.25)
bboxes = utils.nms(bboxes, 0.2, method='nms')
image = visualize.draw_bbox(original_image,
bboxes,
classes=utils.read_class_names(
cfg.YOLO.CLASSES))
image = Image.fromarray(image)
image.show()
def detect(model_name, weight_path, input_size, image_path, framework):
assert model_name in ['yolov3_tiny', 'yolov3', 'yolov4']
if model_name == 'yolov3_tiny':
STRIDES = np.array(cfg.YOLO.STRIDES_TINY)
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_TINY, True)
elif model_name == 'yolov3':
STRIDES = np.array(cfg.YOLO.STRIDES)
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_V3, False)
elif model_name == 'yolov4':
STRIDES = np.array(cfg.YOLO.STRIDES)
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS, False)
else:
raise ValueError
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
XYSCALE = cfg.YOLO.XYSCALE
original_image = cv2.imread(image_path)
original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
original_image_size = original_image.shape[:2]
image_data = utils.image_preprocess(np.copy(original_image),
[input_size, input_size])
image_data = image_data[np.newaxis, ...].astype(np.float32)
if framework == 'tf':
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
if model_name == 'yolov3_tiny':
feature_maps = YOLOv3_tiny(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode(fm, NUM_CLASS)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
elif model_name == 'yolov3':
feature_maps = YOLOv3(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode(fm, NUM_CLASS)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
elif model_name == 'yolov4':
feature_maps = YOLOv4(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode(fm, NUM_CLASS)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
else:
model = None
raise ValueError
if weight_path.split(".")[-1] == "weights":
if model_name == 'yolov3_tiny':
utils.load_weights_tiny(model, weight_path)
# utils.extract_weights_tiny(model, weight_path)
print('load yolo tiny 3')
elif model_name == 'yolov3':
utils.load_weights_v3(model, weight_path)
print('load yolo 3')
elif model_name == 'yolov4':
utils.load_weights(model, weight_path)
print('load yolo 4')
else:
raise ValueError
elif weight_path.split(".")[-1] == "npy":
if model_name == 'yolov3_tiny':
# utils.load_weights_tiny_npy(model, weight_path)
print('load yolo tiny 3 npy')
else:
model.load_weights(weight_path)
print('Restoring weights from: %s ' % weight_path)
# weight = np.load('D:\\coursera\\YoLoSerirs\\checkpoint\\yolo3_tiny.npy', allow_pickle=True)
# model.set_weights(weight)
# model.summary()
start_time = time.time()
pred_bbox = model.predict(image_data)
print(time.time() - start_time)
else:
# Load TFLite model and allocate tensors.
interpreter = tf.lite.Interpreter(model_path=weight_path)
interpreter.allocate_tensors()
# Get input and output tensors.
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
print(input_details)
print(output_details)
interpreter.set_tensor(input_details[0]['index'], image_data)
start_time = time.time()
interpreter.invoke()
pred_bbox = [
interpreter.get_tensor(output_details[i]['index'])
for i in range(len(output_details))
]
print(time.time() - start_time)
if model_name == 'yolov4':
pred_bbox = utils.postprocess_bbbox(pred_bbox, ANCHORS, STRIDES,
XYSCALE)
else:
pred_bbox = utils.postprocess_bbbox(pred_bbox, ANCHORS, STRIDES)
bboxes = utils.postprocess_boxes(pred_bbox, original_image_size,
input_size, 0.5)
bboxes = utils.nms(bboxes, 0.3, method='nms')
image = visualize.draw_bbox(original_image, bboxes)
image = Image.fromarray(image)
image.show()
def train(model_name, weight_path, save_path, logdir=None):
assert model_name in ['yolov3_tiny', 'yolov3', 'yolov4']
physical_devices = tf.config.experimental.list_physical_devices('GPU')
if len(physical_devices) > 0:
tf.config.experimental.set_memory_growth(physical_devices[0], True)
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
STRIDES = np.array(cfg.YOLO.STRIDES)
IOU_LOSS_THRESH = cfg.YOLO.IOU_LOSS_THRESH
XYSCALE = cfg.YOLO.XYSCALE
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS)
trainset = Dataset('train')
testset = Dataset('test')
isfreeze = False
steps_per_epoch = len(trainset)
first_stage_epochs = cfg.TRAIN.FISRT_STAGE_EPOCHS
second_stage_epochs = cfg.TRAIN.SECOND_STAGE_EPOCHS
global_steps = tf.Variable(1, trainable=False, dtype=tf.int64)
warmup_steps = cfg.TRAIN.WARMUP_EPOCHS * steps_per_epoch
total_steps = (first_stage_epochs + second_stage_epochs) * steps_per_epoch
input_layer = tf.keras.layers.Input([cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3])
if model_name=='yolov3_tiny':
feature_maps = YOLOv3_tiny(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode_train(fm, NUM_CLASS, STRIDES, ANCHORS, i)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
elif model_name=='yolov3':
feature_maps = YOLOv3(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode_train(fm, NUM_CLASS, STRIDES, ANCHORS, i)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
elif model_name=='yolov4':
feature_maps = YOLOv4(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode_train(fm, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
else:
raise ValueError
# for name in ['conv2d_93', 'conv2d_101', 'conv2d_109']:
# layer = model.get_layer(name)
# print(layer.name, layer.output_shape)
if weight_path:
if weight_path.split(".")[-1] == "weights":
if model_name == 'yolov3_tiny':
utils.load_weights_tiny(model, weight_path)
elif model_name=='yolov3':
utils.load_weights_v3(model, weight_path)
elif model_name=='yolov4':
utils.load_weights(model, weight_path)
else:
raise ValueError
else:
model.load_weights(weight_path)
print('Restoring weights from: %s ... ' % weight_path)
optimizer = tf.keras.optimizers.Adam()
if logdir:
if os.path.exists(logdir):
shutil.rmtree(logdir)
writer = tf.summary.create_file_writer(logdir)
else:
writer = None
def train_step(image_data, target):
with tf.GradientTape() as tape:
pred_result = model(image_data, training=True)
giou_loss = conf_loss = prob_loss = 0
# optimizing process
for i in range(3):
conv, pred = pred_result[i * 2], pred_result[i * 2 + 1]
loss_items = ops.compute_loss(pred, conv, target[i][0], target[i][1],
STRIDES=STRIDES, NUM_CLASS=NUM_CLASS,
IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i)
giou_loss += loss_items[0]
conf_loss += loss_items[1]
prob_loss += loss_items[2]
total_loss = giou_loss + conf_loss + prob_loss
gradients = tape.gradient(total_loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
tf.print("=> STEP %4d lr: %.6f giou_loss: %4.2f conf_loss: %4.2f "
"prob_loss: %4.2f total_loss: %4.2f" % (global_steps, optimizer.lr.numpy(),
giou_loss, conf_loss,
prob_loss, total_loss))
# update learning rate
global_steps.assign_add(1)
if global_steps < warmup_steps:
lr = global_steps / warmup_steps * cfg.TRAIN.LR_INIT
else:
lr = cfg.TRAIN.LR_END + \
0.5*(cfg.TRAIN.LR_INIT - cfg.TRAIN.LR_END) * \
((1 + tf.cos((global_steps - warmup_steps) / (total_steps - warmup_steps) * np.pi)))
optimizer.lr.assign(lr.numpy())
# if writer:
# # writing summary data
# with writer.as_default():
# tf.summary.scalar("lr", optimizer.lr, step=global_steps)
# tf.summary.scalar("loss/total_loss", total_loss, step=global_steps)
# tf.summary.scalar("loss/giou_loss", giou_loss, step=global_steps)
# tf.summary.scalar("loss/conf_loss", conf_loss, step=global_steps)
# tf.summary.scalar("loss/prob_loss", prob_loss, step=global_steps)
# writer.flush()
def test_step(image_data, target):
pred_result = model(image_data, training=True)
giou_loss = conf_loss = prob_loss = 0
# optimizing process
for i in range(3):
conv, pred = pred_result[i * 2], pred_result[i * 2 + 1]
loss_items = ops.compute_loss(pred, conv, target[i][0], target[i][1],
STRIDES=STRIDES, NUM_CLASS=NUM_CLASS,
IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i)
giou_loss += loss_items[0]
conf_loss += loss_items[1]
prob_loss += loss_items[2]
total_loss = giou_loss + conf_loss + prob_loss
tf.print("=> TEST STEP %4d giou_loss: %4.2f conf_loss: %4.2f "
"prob_loss: %4.2f total_loss: %4.2f" % (global_steps, giou_loss, conf_loss,
prob_loss, total_loss))
for epoch in range(first_stage_epochs + second_stage_epochs):
if epoch < first_stage_epochs:
if not isfreeze:
isfreeze = True
for name in ['conv2d_93', 'conv2d_101', 'conv2d_109']:
freeze = model.get_layer(name)
ops.freeze_all(freeze)
elif epoch >= first_stage_epochs:
if isfreeze:
isfreeze = False
for name in ['conv2d_93', 'conv2d_101', 'conv2d_109']:
freeze = model.get_layer(name)
ops.unfreeze_all(freeze)
for image_data, target in trainset:
train_step(image_data, target)
for image_data, target in testset:
test_step(image_data, target)
if save_path:
model.save_weights(save_path)
param = None
# 输入字典
feed_vars = [('image', inputs), ]
feed_vars = OrderedDict(feed_vars)
if algorithm == 'YOLOv4':
if postprocess == 'fastnms':
boxes, scores, classes = YOLOv4(inputs, num_classes, num_anchors, is_test=False, trainable=True, export=True, postprocess=postprocess, param=param)
test_fetches = {'boxes': boxes, 'scores': scores, 'classes': classes, }
if postprocess == 'multiclass_nms':
pred = YOLOv4(inputs, num_classes, num_anchors, is_test=False, trainable=True, export=True, postprocess=postprocess, param=param)
test_fetches = {'pred': pred, }
elif algorithm == 'YOLOv3':
backbone = Resnet50Vd()
head = YOLOv3Head(keep_prob=1.0) # 一定要设置keep_prob=1.0, 为了得到一致的推理结果
yolov3 = YOLOv3(backbone, head)
if postprocess == 'fastnms':
boxes, scores, classes = yolov3(inputs, export=True, postprocess=postprocess, param=param)
test_fetches = {'boxes': boxes, 'scores': scores, 'classes': classes, }
if postprocess == 'multiclass_nms':
pred = yolov3(inputs, export=True, postprocess=postprocess, param=param)
test_fetches = {'pred': pred, }
infer_prog = infer_prog.clone(for_test=True)
place = fluid.CPUPlace()
exe = fluid.Executor(place)
exe.run(startup_prog)
logger.info("postprocess: %s" % postprocess)
load_params(exe, infer_prog, model_path)
bn_layer.set_weights(yolo_weight[bn_layer_name])
else:
if exclude:
pass
else:
conv_layer.set_weights(yolo_weight[conv_layer_name])
return yolo_weight['final_layer']
if __name__ == '__main__':
from Config.config import cfg
from model.yolov3 import YOLOv3
from model import ops
NUM_CLASS = len(read_class_names(cfg.YOLO.CLASSES))
STRIDES = np.array(cfg.YOLO.STRIDES)
ANCHORS = get_anchors(cfg.YOLO.ANCHORS)
input_layer = tf.keras.layers.Input([cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3])
feature_maps = YOLOv3(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode_train(fm, NUM_CLASS, STRIDES, ANCHORS, i)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
extract_weights = extract_weights_v3(model, weights_file='D:\\coursera\\YoLoSerirs\\pretrain\\yolov3.weights')
np.save('/pretrain/yolov3.npy', extract_weights)
def evaluate(model_name, weight_path):
assert model_name in ['yolov3_tiny', 'yolov3', 'yolov4']
physical_devices = tf.config.experimental.list_physical_devices('GPU')
if len(physical_devices) > 0:
tf.config.experimental.set_memory_growth(physical_devices[0], True)
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
STRIDES = np.array(cfg.YOLO.STRIDES)
IOU_LOSS_THRESH = cfg.YOLO.IOU_LOSS_THRESH
XYSCALE = cfg.YOLO.XYSCALE
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS)
trainset = Dataset('train')
isfreeze = False
steps_per_epoch = len(trainset)
first_stage_epochs = cfg.TRAIN.FISRT_STAGE_EPOCHS
second_stage_epochs = cfg.TRAIN.SECOND_STAGE_EPOCHS
global_steps = tf.Variable(1, trainable=False, dtype=tf.int64)
warmup_steps = cfg.TRAIN.WARMUP_EPOCHS * steps_per_epoch
total_steps = (first_stage_epochs + second_stage_epochs) * steps_per_epoch
input_layer = tf.keras.layers.Input([cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3])
if model_name=='yolov3_tiny':
feature_maps = YOLOv3_tiny(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode_train(fm, NUM_CLASS, STRIDES, ANCHORS, i)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
elif model_name=='yolov3':
feature_maps = YOLOv3(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode_train(fm, NUM_CLASS, STRIDES, ANCHORS, i)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
elif model_name=='yolov4':
feature_maps = YOLOv4(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode_train(fm, NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
else:
raise ValueError
if weight_path:
if weight_path.split(".")[-1] == "weights":
if model_name == 'yolov3_tiny':
utils.load_weights_tiny(model, weight_path)
elif model_name=='yolov3':
utils.load_weights_v3(model, weight_path)
elif model_name=='yolov4':
utils.load_weights(model, weight_path)
else:
raise ValueError
else:
model.load_weights(weight_path)
print('Restoring weights from: %s ... ' % weight_path)
trainset = Dataset('train')
for image_data, target in trainset:
pred_result = model(image_data, training=True)
giou_loss = conf_loss = prob_loss = 0
for i in range(3):
conv, pred = pred_result[i * 2], pred_result[i * 2 + 1]
loss_items = ops.compute_loss(pred, conv, target[i][0], target[i][1],
STRIDES=STRIDES, NUM_CLASS=NUM_CLASS,
IOU_LOSS_THRESH=IOU_LOSS_THRESH, i=i)
giou_loss += loss_items[0]
conf_loss += loss_items[1]
prob_loss += loss_items[2]
total_loss = giou_loss + conf_loss + prob_loss
tf.print("=> STEP %4d giou_loss: %4.2f conf_loss: %4.2f "
"prob_loss: %4.2f total_loss: %4.2f" % (global_steps, giou_loss,
conf_loss, prob_loss, total_loss))
def __init__(self, config):
self.config = config
# Train on device
target_device = config['train']['device']
if torch.cuda.is_available():
torch.backends.cudnn.benchmark = True
self.device = target_device
else:
self.device = "cpu"
# Load dataset
train_transform = get_yolo_transform(config['dataset']['size'],
mode='train')
valid_transform = get_yolo_transform(config['dataset']['size'],
mode='test')
train_dataset = YOLODataset(
csv_file=config['dataset']['train']['csv'],
img_dir=config['dataset']['train']['img_root'],
label_dir=config['dataset']['train']['label_root'],
anchors=config['dataset']['anchors'],
scales=config['dataset']['scales'],
n_classes=config['dataset']['n_classes'],
transform=train_transform)
valid_dataset = YOLODataset(
csv_file=config['dataset']['valid']['csv'],
img_dir=config['dataset']['valid']['img_root'],
label_dir=config['dataset']['valid']['label_root'],
anchors=config['dataset']['anchors'],
scales=config['dataset']['scales'],
n_classes=config['dataset']['n_classes'],
transform=valid_transform)
# DataLoader
self.train_loader = DataLoader(
dataset=train_dataset,
batch_size=config['dataloader']['batch_size'],
num_workers=config['dataloader']['num_workers'],
pin_memory=True,
shuffle=True,
drop_last=False)
self.valid_loader = DataLoader(
dataset=valid_dataset,
batch_size=config['dataloader']['batch_size'],
num_workers=config['dataloader']['num_workers'],
pin_memory=True,
shuffle=False,
drop_last=False)
# Model
model = YOLOv3(
in_channels=config['model']['in_channels'],
num_classes=config['model']['num_classes'],
)
self.model = model.to(self.device)
# Faciliated Anchor boxes with model
torch_anchors = torch.tensor(config['dataset']['anchors']) # (3, 3, 2)
torch_scales = torch.tensor(config['dataset']['scales']) # (3,)
scaled_anchors = ( # (3, 3, 2)
torch_anchors *
(torch_scales.unsqueeze(1).unsqueeze(1).repeat(1, 3, 2)))
self.scaled_anchors = scaled_anchors.to(self.device)
# Optimizer
self.scaler = torch.cuda.amp.GradScaler()
self.optimizer = optim.Adam(
params=self.model.parameters(),
lr=config['optimizer']['lr'],
weight_decay=config['optimizer']['weight_decay'],
)
# Scheduler
self.scheduler = OneCycleLR(
self.optimizer,
max_lr=config['optimizer']['lr'],
epochs=config['train']['n_epochs'],
steps_per_epoch=len(self.train_loader),
)
# Loss function
self.loss_fn = YOLOLoss()
# Tensorboard
self.logdir = config['train']['logdir']
self.board = SummaryWriter(logdir=config['train']['logdir'])
# Training State
self.current_epoch = 0
self.current_map = 0
def prune_train(model_name, weight_path, logdir, save_path, epoches):
assert model_name in ['yolov3_tiny', 'yolov3', 'yolov4']
physical_devices = tf.config.experimental.list_physical_devices('GPU')
if len(physical_devices) > 0:
tf.config.experimental.set_memory_growth(physical_devices[0], True)
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
STRIDES = np.array(cfg.YOLO.STRIDES)
IOU_LOSS_THRESH = cfg.YOLO.IOU_LOSS_THRESH
XYSCALE = cfg.YOLO.XYSCALE
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS)
trainset = Dataset('train')
isfreeze = False
steps_per_epoch = len(trainset)
first_stage_epochs = cfg.TRAIN.FISRT_STAGE_EPOCHS
second_stage_epochs = cfg.TRAIN.SECOND_STAGE_EPOCHS
global_steps = tf.Variable(1, trainable=False, dtype=tf.int64)
warmup_steps = cfg.TRAIN.WARMUP_EPOCHS * steps_per_epoch
total_steps = (first_stage_epochs + second_stage_epochs) * steps_per_epoch
input_layer = tf.keras.layers.Input(
[cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3])
if model_name == 'yolov3_tiny':
feature_maps = YOLOv3_tiny(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode_train(fm, NUM_CLASS, STRIDES, ANCHORS, i)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
elif model_name == 'yolov3':
feature_maps = YOLOv3(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode_train(fm, NUM_CLASS, STRIDES, ANCHORS, i)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
elif model_name == 'yolov4':
feature_maps = YOLOv4(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode_train(fm, NUM_CLASS, STRIDES, ANCHORS, i,
XYSCALE)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
else:
raise ValueError
if weight_path:
if weight_path.split(".")[-1] == "weights":
if model_name == 'yolov3_tiny':
utils.load_weights_tiny(model, weight_path)
elif model_name == 'yolov3':
utils.load_weights_v3(model, weight_path)
elif model_name == 'yolov4':
utils.load_weights(model, weight_path)
else:
raise ValueError
else:
model.load_weights(weight_path)
print('Restoring weights from: %s ... ' % weight_path)
optimizer = tf.keras.optimizers.Adam(learning_rate=0.0001)
if os.path.exists(logdir):
shutil.rmtree(logdir)
# for layer in model.layers:
# print(layer.name, isinstance(layer, tf.keras.layers.Conv2D))
def apply_pruning_to_dense(layer):
if isinstance(layer, tf.keras.layers.Conv2D):
return tfmot.sparsity.keras.prune_low_magnitude(layer)
return layer
# Use `tf.keras.models.clone_model` to apply `apply_pruning_to_dense`
# to the layers of the model.
model_for_pruning = tf.keras.models.clone_model(
model,
clone_function=apply_pruning_to_dense,
)
# model_for_pruning.summary()
unused_arg = -1
model_for_pruning.optimizer = optimizer
step_callback = tfmot.sparsity.keras.UpdatePruningStep()
step_callback.set_model(model_for_pruning)
log_callback = tfmot.sparsity.keras.PruningSummaries(
log_dir=logdir) # Log sparsity and other metrics in Tensorboard.
log_callback.set_model(model_for_pruning)
step_callback.on_train_begin() # run pruning callback
for epoch in range(epoches):
log_callback.on_epoch_begin(epoch=unused_arg) # run pruning callback
for image_data, target in trainset:
step_callback.on_train_batch_begin(
batch=unused_arg) # run pruning callback
with tf.GradientTape() as tape:
pred_result = model_for_pruning(image_data, training=True)
giou_loss = conf_loss = prob_loss = 0
# optimizing process
for i in range(3):
conv, pred = pred_result[i * 2], pred_result[i * 2 + 1]
loss_items = ops.compute_loss(
pred,
conv,
target[i][0],
target[i][1],
STRIDES=STRIDES,
NUM_CLASS=NUM_CLASS,
IOU_LOSS_THRESH=IOU_LOSS_THRESH,
i=i)
giou_loss += loss_items[0]
conf_loss += loss_items[1]
prob_loss += loss_items[2]
total_loss = giou_loss + conf_loss + prob_loss
gradients = tape.gradient(
total_loss, model_for_pruning.trainable_variables)
optimizer.apply_gradients(
zip(gradients, model_for_pruning.trainable_variables))
tf.print(
"=> STEP %4d lr: %.6f giou_loss: %4.2f conf_loss: %4.2f "
"prob_loss: %4.2f total_loss: %4.2f" %
(global_steps, optimizer.lr.numpy(), giou_loss, conf_loss,
prob_loss, total_loss))
step_callback.on_epoch_end(batch=unused_arg) # run pruning callback
model_for_export = tfmot.sparsity.keras.strip_pruning(model_for_pruning)
return model_for_export
def train(model_name, weight_path, save_path, stage, learn_rate, epochs,
use_self_npy):
assert model_name in ['yolov3']
physical_devices = tf.config.experimental.list_physical_devices('GPU')
if len(physical_devices) > 0:
tf.config.experimental.set_memory_growth(physical_devices[0], True)
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
STRIDES = np.array(cfg.YOLO.STRIDES)
IOU_LOSS_THRESH = cfg.YOLO.IOU_LOSS_THRESH
XYSCALE = cfg.YOLO.XYSCALE
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS)
trainset = General_Dataset('train', cfg=cfg)
isfreeze = False
steps_per_epoch = len(trainset)
global_steps = tf.Variable(1, trainable=False, dtype=tf.int64)
total_steps = epochs * steps_per_epoch
print('steps_per_epoch:', steps_per_epoch)
input_layer = tf.keras.layers.Input(
[cfg.TRAIN.INPUT_SIZE, cfg.TRAIN.INPUT_SIZE, 3])
feature_maps = YOLOv3(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode_train(fm, NUM_CLASS, STRIDES, ANCHORS, i)
bbox_tensors.append(fm)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
if weight_path:
if use_self_npy:
weight = np.load(weight_path, allow_pickle=True)
model.set_weights(weight)
final_layers = []
else:
final_layers = utils.load_weights_v3_npy(model,
weight_path,
exclude=True)
print('Restoring weights from: %s ... ' % weight_path)
else:
final_layers = []
optimizer = tf.keras.optimizers.Adam(learn_rate)
avg_giou_loss = []
avg_conf_loss = []
def train_step(image_data, target):
global avg_giou_loss, avg_conf_loss
with tf.GradientTape() as tape:
pred_result = model(image_data, training=True)
giou_loss = conf_loss = prob_loss = 0
# optimizing process
for i in range(3):
conv, pred = pred_result[i * 2], pred_result[i * 2 + 1]
loss_items = ops.compute_loss(pred,
conv,
target[i][0],
target[i][1],
STRIDES=STRIDES,
NUM_CLASS=NUM_CLASS,
IOU_LOSS_THRESH=IOU_LOSS_THRESH,
i=i)
giou_loss += loss_items[0]
conf_loss += loss_items[1]
prob_loss += loss_items[2]
total_loss = giou_loss + conf_loss + prob_loss
gradients = tape.gradient(total_loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients,
model.trainable_variables))
avg_giou_loss.append(giou_loss)
avg_conf_loss.append(conf_loss)
if global_steps % 10 == 0:
tf.print(
"=> STEP %4d lr: %.6f giou_loss: %4.2f conf_loss: %4.2f "
"prob_loss: %4.2f total_loss: %4.2f" %
(global_steps,
optimizer.lr.numpy(), np.mean(avg_giou_loss),
np.mean(avg_conf_loss), prob_loss, total_loss))
avg_giou_loss = []
avg_conf_loss = []
global_steps.assign_add(1)
if stage == 'last':
for layer in model.layers:
if layer.name not in ['conv2d_74', 'conv2d_66', 'conv2d_58']:
layer.trainable = False
else:
print(layer.name)
for epoch in range(epochs):
for image_data, target in trainset:
train_step(image_data, target)
if save_path:
np.save(save_path, model.get_weights())
def save_tflite(model_name, weight_path, quantize_mode, output, input_size):
assert model_name in ['yolov3_tiny', 'yolov3', 'yolov4']
assert quantize_mode in ['int8', 'float16', 'full_int8']
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
if model_name == 'yolov3_tiny':
feature_maps = YOLOv3_tiny(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode(fm, NUM_CLASS)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
elif model_name == 'yolov3':
feature_maps = YOLOv3(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode(fm, NUM_CLASS)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
elif model_name == 'yolov4':
feature_maps = YOLOv4(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = ops.decode(fm, NUM_CLASS)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
else:
model = None
raise ValueError
if weight_path.split(".")[-1] == "weights":
if model_name == 'yolov3_tiny':
utils.load_weights_tiny(model, weight_path)
elif model_name == ' yolov3':
utils.load_weights_v3(model, weight_path)
elif model_name == 'yolov4':
utils.load_weights(model, weight_path)
else:
raise ValueError
else:
model.load_weights(weight_path).expect_partial()
print('Restoring weights from: %s ... ' % weight_path)
# model.summary()
converter = tf.lite.TFLiteConverter.from_keras_model(model)
if tf.__version__ >= '2.2.0':
converter.experimental_new_converter = False
if quantize_mode == 'int8':
converter.optimizations = [tf.lite.Optimize.DEFAULT]
elif quantize_mode == 'float16':
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [
tf.compat.v1.lite.constants.FLOAT16
]
elif quantize_mode == 'full_int8':
converter.target_spec.supported_ops = [
tf.lite.OpsSet.TFLITE_BUILTINS_INT8
]
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [
tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS
]
converter.allow_custom_ops = True
converter.representative_dataset = representative_data_gen
else:
raise ValueError
tflite_model = converter.convert()
open(output, 'wb').write(tflite_model)
logging.info("model saved to: {}".format(output))