def main():
physical_devices = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(physical_devices[0], True)
#get train set
with open(train_path) as f:
_line = f.readlines()
train_set = [i.rstrip('\n') for i in _line]
train_generator = YoloGenerator(train_list=train_set,
anchors=anchors,
num_classes=num_classes,
batch_size=batch_size,
shuffle=False,
multi_scale=True,
visual_effect=VisualEffect(),
misc_effect=MiscEffect(border_value=0),
input_size=416)
#creat model
model, debug_model = yoloNano(anchors, num_classes=num_classes)
#if you want to resume the train,open the code
#model.load_weights('./model_save/multi_coco_car.h5')
model.compile(optimizer=tf.keras.optimizers.Adam(lr=1e-4),
loss={
'yolo_loss': lambda y_true, y_pred: y_pred
})
callbacks = create_callbacks()
# start training
model.fit_generator(generator=train_generator,
epochs=epochs,
callbacks=callbacks)
return 0
def main():
physical_devices = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(physical_devices[0], True)
train_model, test_model = yoloNano(anchors,
input_size=416,
num_classes=num_classes,
include_attention=True)
test_model.summary()
test_model.load_weights('./model_save/multi_coco_car.h5')
cap = cv2.VideoCapture('1.mp4')
while True:
ret, img = cap.read()
org_h = img.shape[0]
org_w = img.shape[1]
max_side = max(org_h, org_w)
if org_h > org_w:
scale = org_w / max_side
pts1 = np.array([[0, 0], [org_w, 0], [0, org_h]], dtype=np.float32)
pts2 = np.array([[img_size * (1 - scale) / 2, 0],
[img_size * (1 + scale) / 2, 0],
[img_size * (1 - scale) / 2, img_size]],
dtype=np.float32)
M = cv2.getAffineTransform(pts1, pts2)
img = cv2.warpAffine(img, M, (img_size, img_size))
else:
scale = org_h / max_side
pts1 = np.array([[0, 0], [org_w, 0], [0, org_h]], dtype=np.float32)
offset1 = img_size * (1 - scale) / 2
offset2 = img_size * (1 + scale) / 2
pts2 = np.array([[0, offset1], [img_size, offset1], [0, offset2]],
dtype=np.float32)
M = cv2.getAffineTransform(pts1, pts2)
img = cv2.warpAffine(img, M, (img_size, img_size))
img = img / 255.0
pred_img = img[np.newaxis, :]
yolo_output = test_model.predict(pred_img)
boxes_, scores_, classes_ = yolo_eval(yolo_outputs=yolo_output,
anchors=anchors,
num_classes=num_classes,
image_shape=np.array(
[img_size, img_size]),
score_threshold=0.7)
for box in boxes_[0]:
ymin = int(box[0])
xmin = int(box[1])
ymax = int(box[2])
xmax = int(box[3])
cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 255, 0))
cv2.imshow('pred', img)
cv2.waitKey(100)
def main():
physical_devices = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(physical_devices[0], True)
train_model, test_model = yoloNano(anchors,
input_size=416,
num_classes=num_classes)
test_model.summary()
test_model.load_weights('./model_save/save_model.h5')
#img = cv2.imread('/home/cvos/Datasets/coco_car/val/COCO_val2014_000000003849.jpg')
img = cv2.imread('./test_img/Untitled Folder/4.jpeg')
org_h = img.shape[0]
org_w = img.shape[1]
max_side = max(org_h, org_w)
if org_h > org_w:
scale = org_w / max_side
pts1 = np.array([[0, 0], [org_w, 0], [0, org_h]], dtype=np.float32)
pts2 = np.array([[img_size *
(1 - scale) / 2, 0], [img_size * (1 + scale) / 2, 0],
[img_size * (1 - scale) / 2, img_size]],
dtype=np.float32)
M = cv2.getAffineTransform(pts1, pts2)
img = cv2.warpAffine(img, M, (img_size, img_size))
else:
scale = org_h / max_side
pts1 = np.array([[0, 0], [org_w, 0], [0, org_h]], dtype=np.float32)
offset1 = img_size * (1 - scale) / 2
offset2 = img_size * (1 + scale) / 2
pts2 = np.array([[0, offset1], [img_size, offset1], [0, offset2]],
dtype=np.float32)
M = cv2.getAffineTransform(pts1, pts2)
img = cv2.warpAffine(img, M, (img_size, img_size))
img = img / 255.0
pred_img = img[np.newaxis, :]
yolo_output = test_model.predict(pred_img)
boxes_, scores_, classes_ = yolo_eval(yolo_output,
anchors,
1,
np.array([416, 416]),
score_threshold=.2)
for box in boxes_[0]:
ymin = int(box[0])
xmin = int(box[1])
ymax = int(box[2])
xmax = int(box[3])
cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 255, 0))
cv2.imshow('pred', img)
cv2.waitKey(0)