def box_gen(folder, images):
model = YOLO()
for path in images:
print(path)
path = os.path.join(data_folder, path)
img = cv2.imread(path)
boxes = model.predict(img)
yield img, boxes
def _main_(args):
config_path = args.conf
weights_path = args.weights
image_path = args.input
with open(config_path) as config_buffer:
config = json.load(config_buffer)
yolo = YOLO(backend=config['model']['architecture'],
input_size=config['model']['input_size'],
labels=config['model']['labels'],
max_box_per_image=config['model']['max_box_per_image'],
anchors=config['model']['anchors'])
yolo.load_weights(weights_path)
if image_path[-4:] == '.mp4':
video_out = image_path[:-4] + '_detected' + image_path[-4:]
video_reader = cv2.VideoCapture(image_path)
nb_frames = int(video_reader.get(cv2.CAP_PROP_FRAME_COUNT))
frame_h = int(video_reader.get(cv2.CAP_PROP_FRAME_HEIGHT))
frame_w = int(video_reader.get(cv2.CAP_PROP_FRAME_WIDTH))
video_writer = cv2.VideoWriter(video_out,
cv2.VideoWriter_fourcc(*'MPEG'), 50.0,
(frame_w, frame_h))
for i in tqdm(range(nb_frames)):
_, image = video_reader.read()
boxes = yolo.predict(image)
image = draw_boxes(image, boxes, config['model']['labels'])
video_writer.write(np.uint8(image))
video_reader.release()
video_writer.release()
else:
image = cv2.imread(image_path)
boxes = yolo.predict(image)
image = draw_boxes(image, boxes, config['model']['labels'])
print(len(boxes), 'boxes are found')
cv2.imwrite(image_path[:-4] + '_result' + image_path[-4:], image)
def find_ball(img):
model = YOLO(w_file=WEIGHTS_FILE, c_file=CONFIG_FILE, n_file=NAMES_FILE)
boxes = model.predict(img, 32)
return boxes
from constants import CONFIG_FILE
import json
from keras.models import Model, load_model
import cv2
import os
from model import YOLO
if __name__ == '__main__':
with open(CONFIG_FILE) as config_buffer:
config = json.loads(config_buffer.read())
image_path = os.path.expanduser(config["predict"]["image_path"])
model_path = os.path.expanduser(config["predict"]["model_path"])
output_file_path = os.path.expanduser(config["predict"]["output"])
print('>>>>> Creating YOLO object')
yolo = YOLO(input_size=tuple(config['model']['input_size']),
grid_size=int(config['model']['grid_size']),
bbox_count=int(config['model']['bboxes_per_grid_cell']),
classes=config['model']['class_names'],
lambda_coord=config['model']['lambda_coord'],
lambda_noobj=config['model']['lambda_noobj'],
bbox_params=config['model']['bbox_params'])
if os.path.isfile(image_path) and os.path.isfile(model_path):
yolo.predict(image_path, model_path, output_file_path)
else:
print('Path to image or model does not exist...')