Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add predict directory option #282

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
4 changes: 4 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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).
Expand Down
13 changes: 12 additions & 1 deletion predict.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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)
Expand Down