def train():
model = yolo_body()
if os.path.exists(weights_path):
print('+++++++++++++++++++++++++++++++++++++++++++++++++')
model.load_weights(weights_path, by_name=True, skip_mismatch=True)
model.compile(optimizer=Adam(lr=1e-3), loss=my_loss)
model.fit_generator(
generator=data_generator(data_train, batch_size=batch_size),
steps_per_epoch=max(1, num_train // batch_size),
validation_data=data_generator(data_test, batch_size=batch_size),
validation_steps=max(1, num_test // batch_size),
epochs=2,
initial_epoch=0,
callbacks=[logging, checkpoint, reduce_lr, early_stopping])
model.save_weights(weights_path)
def predict():
model = yolo_body()
model.load_weights(weights_path)
imgs = list()
for i in range(1):
img_path = r'D:\project\data-set\test\挨络替糙骡_1547711365.jpg'
img = resize_img(img_path)
imgs.append(img)
# img_expand = np.expand_dims(img, axis=0)
imgs = np.array(imgs)
y_pred = model.predict(x=imgs)
print(y_pred.shape)
# 给图片画box
img_box = draw_box(y_pred, img)
img_name = os.path.split(img_path)[-1]
print(os.path.join(result_dir, img_name), img_name)
img_box.save(os.path.join(result_dir, img_name))
def predict():
model = yolo_body()
model.load_weights(weights_path)
imgs = list()
for i in range(1):
img_path = os.path.join(r'D:\project\data-set\test', img_[0])
img = resize_img(img_path)
imgs.append(img)
# img_expand = np.expand_dims(img, axis=0)
imgs = np.array(imgs)
y_pred = model.predict(x=imgs)
print(y_pred[0].shape)
# 给图片画box
img_box = draw_box(y_pred, img)
img_name = os.path.split(img_path)[-1].split('_')[0] + '_' + str(
int(time.time())) + '.jpg'
print(os.path.join(result_dir, img_name), img_name)
img_box.save(os.path.join(result_dir, img_name))
def generate(self):
num_anchors = len(self.anchors)
num_classes = len(self.class_names)
try:
self.yolo_model = load_model(self.model_path, compile=False)
except:
self.yolo_model = yolo_body(Input(shape=(None, None, 3)),
num_anchors // 3, num_classes)
self.yolo_model.load_weights(
self.model_path) # make sure model, anchors and classes match
print('{} model, anchors, and classes loaded.'.format(self.model_path))
# 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))
np.random.seed(10101) # Fixed seed for consistent colors across runs.
np.random.shuffle(
self.colors) # Shuffle colors to decorrelate adjacent classes.
np.random.seed(None) # Reset seed to default.
# Generate output tensor targets for filtered bounding boxes.
self.input_image_shape = K.placeholder(shape=(2, ))
if self.gpu_num >= 2:
self.yolo_model = multi_gpu_model(self.yolo_model,
gpus=self.gpu_num)
boxes, scores, classes = yolo_eval(self.yolo_model.output,
self.anchors,
len(self.class_names),
self.input_image_shape,
score_threshold=self.score,
iou_threshold=self.iou)
return boxes, scores, classes