-
Notifications
You must be signed in to change notification settings - Fork 4
/
get_inference.py
73 lines (57 loc) · 1.92 KB
/
get_inference.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
#! /usr/bin/env python3
from tqdm import tqdm
from utils import draw_boxes, get_session
from frontend import YOLO
from utils import list_images
import tensorflow as tf
import numpy as np
import keras
import json
import argparse
import os
import cv2
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]="0"
argparser = argparse.ArgumentParser(
description='Train and validate YOLO_v2 model on any dataset')
argparser.add_argument(
'-c',
'--conf',
default='config.json',
help='path to configuration file')
argparser.add_argument(
'-w',
'--weights',
default='',
help='path to pretrained weights')
argparser.add_argument(
'-i',
'--input',
help='path to an image or an video (mp4 format)')
def _main_(args):
config_path = args.conf
weights_path = args.weights
image_path = args.input
keras.backend.tensorflow_backend.set_session(get_session())
with open(config_path) as config_buffer:
config = json.load(config_buffer)
if weights_path == '':
weights_path = config['train']['saved_weights_name']
###############################
# Make the model
###############################
yolo = YOLO(backend = config['model']['backend'],
input_size = (config['model']['input_size_h'],config['model']['input_size_w']),
labels = config['model']['labels'],
max_box_per_image = config['model']['max_box_per_image'],
anchors = config['model']['anchors'],
gray_mode = config['model']['gray_mode'])
###############################
# Load trained weights
###############################
yolo.load_weights(weights_path)
inference_model = yolo.get_inference_model()
inference_model.save("inference.h5")
if __name__ == '__main__':
args = argparser.parse_args()
_main_(args)