|
- import os
- import colorsys
-
- import numpy as np
- from keras import backend as K
- from keras.models import load_model
- from keras.layers import Input
-
- from yolo4.model import yolo_eval, yolo4_body
- from yolo4.utils import letterbox_image
-
- from PIL import Image, ImageFont, ImageDraw
- from timeit import default_timer as timer
- import matplotlib.pyplot as plt
-
- from operator import itemgetter
-
- class Yolo4(object):
- def get_class(self):
- classes_path = os.path.expanduser(self.classes_path)
- with open(classes_path) as f:
- class_names = f.readlines()
- class_names = [c.strip() for c in class_names]
- return class_names
-
- def get_anchors(self):
- anchors_path = os.path.expanduser(self.anchors_path)
- with open(anchors_path) as f:
- anchors = f.readline()
- anchors = [float(x) for x in anchors.split(',')]
- return np.array(anchors).reshape(-1, 2)
-
- def load_yolo(self):
- model_path = os.path.expanduser(self.model_path)
- assert model_path.endswith('.h5'), 'Keras model or weights must be a .h5 file.'
-
- self.class_names = self.get_class()
- self.anchors = self.get_anchors()
-
- num_anchors = len(self.anchors)
- num_classes = len(self.class_names)
-
- # Generate colors for drawing bounding boxes.
- hsv_tuples = [(x / len(self.class_names), 1., 1.)
- for x in range(len(self.class_names))]
- self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
- self.colors = list(
- map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
- self.colors))
-
- self.sess = K.get_session()
-
- # Load model, or construct model and load weights.
- self.yolo4_model = yolo4_body(Input(shape=(608, 608, 3)), num_anchors//3, num_classes)
-
- # Read and convert darknet weight
- print('Loading weights.')
- weights_file = open(self.weights_path, 'rb')
- major, minor, revision = np.ndarray(
- shape=(3, ), dtype='int32', buffer=weights_file.read(12))
- if (major*10+minor)>=2 and major<1000 and minor<1000:
- seen = np.ndarray(shape=(1,), dtype='int64', buffer=weights_file.read(8))
- else:
- seen = np.ndarray(shape=(1,), dtype='int32', buffer=weights_file.read(4))
- print('Weights Header: ', major, minor, revision, seen)
-
- convs_to_load = []
- bns_to_load = []
- for i in range(len(self.yolo4_model.layers)):
- layer_name = self.yolo4_model.layers[i].name
- if layer_name.startswith('conv2d_'):
- convs_to_load.append((int(layer_name[7:]), i))
- if layer_name.startswith('batch_normalization_'):
- bns_to_load.append((int(layer_name[20:]), i))
-
- convs_sorted = sorted(convs_to_load, key=itemgetter(0))
- bns_sorted = sorted(bns_to_load, key=itemgetter(0))
-
- bn_index = 0
- for i in range(len(convs_sorted)):
- print('Converting ', i)
- if i == 93 or i == 101 or i == 109:
- #no bn, with bias
- weights_shape = self.yolo4_model.layers[convs_sorted[i][1]].get_weights()[0].shape
- bias_shape = self.yolo4_model.layers[convs_sorted[i][1]].get_weights()[0].shape[3]
- filters = bias_shape
- size = weights_shape[0]
- darknet_w_shape = (filters, weights_shape[2], size, size)
- weights_size = np.product(weights_shape)
-
- conv_bias = np.ndarray(
- shape=(filters, ),
- dtype='float32',
- buffer=weights_file.read(filters * 4))
- conv_weights = np.ndarray(
- shape=darknet_w_shape,
- dtype='float32',
- buffer=weights_file.read(weights_size * 4))
- conv_weights = np.transpose(conv_weights, [2, 3, 1, 0])
- self.yolo4_model.layers[convs_sorted[i][1]].set_weights([conv_weights, conv_bias])
- else:
- #with bn, no bias
- weights_shape = self.yolo4_model.layers[convs_sorted[i][1]].get_weights()[0].shape
- size = weights_shape[0]
- bn_shape = self.yolo4_model.layers[bns_sorted[bn_index][1]].get_weights()[0].shape
- filters = bn_shape[0]
- darknet_w_shape = (filters, weights_shape[2], size, size)
- weights_size = np.product(weights_shape)
-
- conv_bias = np.ndarray(
- shape=(filters, ),
- dtype='float32',
- buffer=weights_file.read(filters * 4))
- bn_weights = np.ndarray(
- shape=(3, filters),
- dtype='float32',
- buffer=weights_file.read(filters * 12))
-
- bn_weight_list = [
- bn_weights[0], # scale gamma
- conv_bias, # shift beta
- bn_weights[1], # running mean
- bn_weights[2] # running var
- ]
- self.yolo4_model.layers[bns_sorted[bn_index][1]].set_weights(bn_weight_list)
-
- conv_weights = np.ndarray(
- shape=darknet_w_shape,
- dtype='float32',
- buffer=weights_file.read(weights_size * 4))
- conv_weights = np.transpose(conv_weights, [2, 3, 1, 0])
- self.yolo4_model.layers[convs_sorted[i][1]].set_weights([conv_weights])
-
- bn_index += 1
-
- weights_file.close()
-
- self.yolo4_model.save(self.model_path)
-
-
- if self.gpu_num>=2:
- self.yolo4_model = multi_gpu_model(self.yolo4_model, gpus=self.gpu_num)
-
- self.input_image_shape = K.placeholder(shape=(2, ))
- self.boxes, self.scores, self.classes = yolo_eval(self.yolo4_model.output, self.anchors,
- len(self.class_names), self.input_image_shape,
- score_threshold=self.score)
-
- def __init__(self, score, iou, anchors_path, classes_path, model_path, weights_path, gpu_num=1):
- self.score = score
- self.iou = iou
- self.anchors_path = anchors_path
- self.classes_path = classes_path
- self.weights_path = weights_path
- self.model_path = model_path
- self.gpu_num = gpu_num
- self.load_yolo()
-
- def close_session(self):
- self.sess.close()
-
- if __name__ == '__main__':
- model_path = 'yolo4_weight.h5'
- anchors_path = 'model_data/yolo4_anchors.txt'
- classes_path = 'model_data/coco_classes.txt'
- weights_path = 'yolov4.weights'
-
- score = 0.5
- iou = 0.5
-
- model_image_size = (608, 608)
-
- yolo4_model = Yolo4(score, iou, anchors_path, classes_path, model_path, weights_path)
-
- yolo4_model.close_session()
|