diff --git a/README.md b/README.md index d06b7f498..1b301060a 100644 --- a/README.md +++ b/README.md @@ -136,6 +136,10 @@ By the end of this process, the code will write the weights of the best model to It carries out detection on the image and write the image with detected bounding boxes to the same folder. +or you can also detect the whole images in a directory by specifying the directory path (the path should end with "/") + +`python predict.py -c config.json -w /path/to/best_weights.h5 -i /path/to/image/folder/` + ## Usage for jupyter notebook Refer to the notebook (https://github.com/experiencor/basic-yolo-keras/blob/master/Yolo%20Step-by-Step.ipynb) for a complete walk-through implementation of YOLOv2 from scratch (training, testing, and scoring). diff --git a/predict.py b/predict.py index 544e29ea6..3d4c4d43c 100644 --- a/predict.py +++ b/predict.py @@ -9,6 +9,7 @@ from utils import draw_boxes from frontend import YOLO import json +import glob os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"]="0" @@ -81,7 +82,17 @@ def _main_(args): video_writer.write(np.uint8(image)) video_reader.release() - video_writer.release() + video_writer.release() + + elif image_path[-1] == '/': + for img in glob.glob(image_path+"*.jpg"): + image = cv2.imread(img) + boxes = yolo.predict(image) + image = draw_boxes(image, boxes, config['model']['labels']) + + print(len(boxes), 'boxes are found') + cv2.imwrite(img[:-4] + '_detected' + img[-4:], image) + else: image = cv2.imread(image_path) boxes = yolo.predict(image)