Added command line option parsing

Added command option parsing, and --model, --anchors, --classes, and --gpu_num, with default values.

Author: philtrade

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`  

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())

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.")