Esempio n. 1
0
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
Esempio n. 2
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)
Esempio n. 3
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)
    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)