From e6598d13c703029b2686bc2eb8d5c09badf42992 Mon Sep 17 00:00:00 2001 From: "Phillip K.S. Chu" Date: Mon, 30 Jul 2018 16:42:51 -0700 Subject: [PATCH] Added command line option parsing Added command option parsing, and --model, --anchors, --classes, and --gpu_num, with default values. --- README.md | 32 ++++++++++++++++---- yolo.py | 55 +++++++++++++++------------------- yolo_video.py | 81 ++++++++++++++++++++++++++++++++++++++++++++------- 3 files changed, 121 insertions(+), 47 deletions(-) diff --git a/README.md b/README.md index 3ae009f02..46f12387b 100644 --- a/README.md +++ b/README.md @@ -18,14 +18,36 @@ A Keras implementation of YOLOv3 (Tensorflow backend) inspired by [allanzelener/ ``` wget https://pjreddie.com/media/files/yolov3.weights python convert.py yolov3.cfg yolov3.weights model_data/yolo.h5 -python yolo.py OR python yolo_video.py [video_path] [output_path(optional)] +python yolo_video.py [OPTIONS...] --image, for image detection mode, OR +python yolo_video.py [video_path] [output_path (optional)] ``` -For Tiny YOLOv3, just do in a similar way. And modify model path and anchor path in yolo.py. +For Tiny YOLOv3, just do in a similar way, just specify model path and anchor path with `--model model_file` and `--anchors anchor_file`. +### Usage +Use --help to see usage of yolo_video.py: +``` +usage: yolo_video.py [-h] [--model MODEL] [--anchors ANCHORS] + [--classes CLASSES] [--gpu_num GPU_NUM] [--image] + [--input] [--output] + +positional arguments: + --input Video input path + --output Video output path + +optional arguments: + -h, --help show this help message and exit + --model MODEL path to model weight file, default model_data/yolo.h5 + --anchors ANCHORS path to anchor definitions, default + model_data/yolo_anchors.txt + --classes CLASSES path to class definitions, default + model_data/coco_classes.txt + --gpu_num GPU_NUM Number of GPU to use, default 1 + --image Image detection mode, will ignore all positional arguments +``` --- -4. MultiGPU usage is an optinal. Change the number of gpu and add gpu device id +4. MultiGPU usage: use `--gpu_num N` to use N GPUs. It is passed to the [Keras multi_gpu_model()](https://keras.io/utils/#multi_gpu_model). ## Training @@ -46,8 +68,8 @@ For Tiny YOLOv3, just do in a similar way. And modify model path and anchor path 3. Modify train.py and start training. `python train.py` - Use your trained weights or checkpoint weights in yolo.py. - Remember to modify class path or anchor path. + Use your trained weights or checkpoint weights with command line option `--model model_file` when using yolo_video.py + Remember to modify class path or anchor path, with `--classes class_file` and `--anchors anchor_file`. If you want to use original pretrained weights for YOLOv3: 1. `wget https://pjreddie.com/media/files/darknet53.conv.74` diff --git a/yolo.py b/yolo.py index d1b6c47f4..4aa348643 100644 --- a/yolo.py +++ b/yolo.py @@ -1,7 +1,6 @@ -#! /usr/bin/env python # -*- coding: utf-8 -*- """ -Run a YOLO_v3 style detection model on test images. +Class definition of YOLO_v3 style detection model on image and video """ import colorsys @@ -17,21 +16,32 @@ from yolo3.model import yolo_eval, yolo_body, tiny_yolo_body from yolo3.utils import letterbox_image import os -os.environ['CUDA_VISIBLE_DEVICES'] = '0' from keras.utils import multi_gpu_model -gpu_num=1 class YOLO(object): - def __init__(self): - self.model_path = 'model_data/yolo.h5' # model path or trained weights path - self.anchors_path = 'model_data/yolo_anchors.txt' - self.classes_path = 'model_data/coco_classes.txt' - self.score = 0.3 - self.iou = 0.45 + _defaults = { + "model_path": 'model_data/yolo.h5', + "anchors_path": 'model_data/yolo_anchors.txt', + "classes_path": 'model_data/coco_classes.txt', + "score" : 0.3, + "iou" : 0.45, + "model_image_size" : (416, 416), + "gpu_num" : 1, + } + + @classmethod + def get_defaults(cls, n): + if n in cls._defaults: + return cls._defaults[n] + else: + return "Unrecognized attribute name '" + n + "'" + + def __init__(self, **kwargs): + self.__dict__.update(self._defaults) # set up default values + self.__dict__.update(kwargs) # and update with user overrides self.class_names = self._get_class() self.anchors = self._get_anchors() self.sess = K.get_session() - self.model_image_size = (416, 416) # fixed size or (None, None), hw self.boxes, self.scores, self.classes = self.generate() def _get_class(self): @@ -82,8 +92,8 @@ def generate(self): # Generate output tensor targets for filtered bounding boxes. self.input_image_shape = K.placeholder(shape=(2, )) - if gpu_num>=2: - self.yolo_model = multi_gpu_model(self.yolo_model, gpus=gpu_num) + if self.gpu_num>=2: + self.yolo_model = multi_gpu_model(self.yolo_model, gpus=self.gpu_num) boxes, scores, classes = yolo_eval(self.yolo_model.output, self.anchors, len(self.class_names), self.input_image_shape, score_threshold=self.score, iou_threshold=self.iou) @@ -159,7 +169,6 @@ def detect_image(self, image): def close_session(self): self.sess.close() - def detect_video(yolo, video_path, output_path=""): import cv2 vid = cv2.VideoCapture(video_path) @@ -201,21 +210,3 @@ def detect_video(yolo, video_path, output_path=""): break yolo.close_session() - -def detect_img(yolo): - while True: - img = input('Input image filename:') - try: - image = Image.open(img) - except: - print('Open Error! Try again!') - continue - else: - r_image = yolo.detect_image(image) - r_image.show() - yolo.close_session() - - - -if __name__ == '__main__': - detect_img(YOLO()) diff --git a/yolo_video.py b/yolo_video.py index daef5769d..7c394617c 100644 --- a/yolo_video.py +++ b/yolo_video.py @@ -1,16 +1,77 @@ import sys +import argparse +from yolo import YOLO, detect_video +from PIL import Image -if len(sys.argv) < 2: - print("Usage: $ python {0} [video_path] [output_path(optional)]", sys.argv[0]) - exit() +def detect_img(yolo): + while True: + img = input('Input image filename:') + try: + image = Image.open(img) + except: + print('Open Error! Try again!') + continue + else: + r_image = yolo.detect_image(image) + r_image.show() + yolo.close_session() -from yolo import YOLO -from yolo import detect_video +FLAGS = None if __name__ == '__main__': - video_path = sys.argv[1] - if len(sys.argv) > 2: - output_path = sys.argv[2] - detect_video(YOLO(), video_path, output_path) + # class YOLO defines the default value, so suppress any default here + parser = argparse.ArgumentParser(argument_default=argparse.SUPPRESS) + ''' + Command line options + ''' + parser.add_argument( + '--model', type=str, + help='path to model weight file, default ' + YOLO.get_defaults("model_path") + ) + + parser.add_argument( + '--anchors', type=str, + help='path to anchor definitions, default ' + YOLO.get_defaults("anchors_path") + ) + + parser.add_argument( + '--classes', type=str, + help='path to class definitions, default ' + YOLO.get_defaults("classes_path") + ) + + parser.add_argument( + '--gpu_num', type=int, + help='Number of GPU to use, default ' + str(YOLO.get_defaults("gpu_num")) + ) + + parser.add_argument( + '--image', default=False, action="store_true", + help='Image detection mode, will ignore all positional arguments' + ) + ''' + Command line positional arguments -- for video detection mode + ''' + parser.add_argument( + "--input", nargs='?', type=str,required=False,default='./path2your_video', + help = "Video input path" + ) + + parser.add_argument( + "--output", nargs='?', type=str, default="", + help = "[Optional] Video output path" + ) + + FLAGS = parser.parse_args() + + if FLAGS.image: + """ + Image detection mode, disregard any remaining command line arguments + """ + print("Image detection mode") + if "input" in FLAGS: + print(" Ignoring remaining command line arguments: " + FLAGS.input + "," + FLAGS.output) + detect_img(YOLO(**vars(FLAGS))) + elif "input" in FLAGS: + detect_video(YOLO(**vars(FLAGS)), FLAGS.input, FLAGS.output) else: - detect_video(YOLO(), video_path) + print("Must specify at least video_input_path. See usage with --help.")