Add files via upload

coco.names

@@ -0,0 +1,80 @@
+person
+bicycle
+car
+motorbike
+aeroplane
+bus
+train
+truck
+boat
+traffic light
+fire hydrant
+stop sign
+parking meter
+bench
+bird
+cat
+dog
+horse
+sheep
+cow
+elephant
+bear
+zebra
+giraffe
+backpack
+umbrella
+handbag
+tie
+suitcase
+frisbee
+skis
+snowboard
+sports ball
+kite
+baseball bat
+baseball glove
+skateboard
+surfboard
+tennis racket
+bottle
+wine glass
+cup
+fork
+knife
+spoon
+bowl
+banana
+apple
+sandwich
+orange
+broccoli
+carrot
+hot dog
+pizza
+donut
+cake
+chair
+sofa
+pottedplant
+bed
+diningtable
+toilet
+tvmonitor
+laptop
+mouse
+remote
+keyboard
+cell phone
+microwave
+oven
+toaster
+sink
+refrigerator
+book
+clock
+vase
+scissors
+teddy bear
+hair drier
+toothbrush

detect.py

@@ -0,0 +1,109 @@
+"""Yolo v3 detection script.
+
+Saves the detections in the `detection` folder.
+
+Usage:
+    python detect.py <images/video> <iou threshold> <confidence threshold> <filenames>
+
+Example:
+    python detect.py images 0.5 0.5 data/images/dog.jpg data/images/office.jpg
+    python detect.py video 0.5 0.5 data/video/shinjuku.mp4
+
+Note that only one video can be processed at one run.
+"""
+
+import os
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+
+import tensorflow as tf
+import sys
+import cv2
+
+from yolo_v3 import Yolo_v3
+from utils import load_images, load_class_names, draw_boxes, draw_frame
+
+_MODEL_SIZE = (416, 416)
+_CLASS_NAMES_FILE = './data/labels/coco.names'
+_MAX_OUTPUT_SIZE = 20
+
+detection_result = {}
+
+
+def main(type, iou_threshold, confidence_threshold, input_names):
+    global detection_result
+    class_names = load_class_names(_CLASS_NAMES_FILE)
+    n_classes = len(class_names)
+
+    model = Yolo_v3(n_classes=n_classes, model_size=_MODEL_SIZE,
+                    max_output_size=_MAX_OUTPUT_SIZE,
+                    iou_threshold=iou_threshold,
+                    confidence_threshold=confidence_threshold)
+
+    if type == 'images':
+        #batch_size = len(input_names)
+        batch = load_images(input_names, model_size=_MODEL_SIZE)
+        inputs = tf.placeholder(tf.float32, [1, *_MODEL_SIZE, 3])
+        detections = model(inputs, training=False)
+        saver = tf.train.Saver(tf.global_variables(scope='yolo_v3_model'))
+
+        with tf.Session() as sess:
+            saver.restore(sess, './weights/model.ckpt')
+            detection_result = sess.run(detections, feed_dict={inputs: batch})
+
+        #detection_result = detection_result[0]
+        print("detection_result", detection_result)
+        draw_boxes(input_names, detection_result, class_names, _MODEL_SIZE)
+
+        print('Detections have been saved successfully.')
+
+    elif type == 'video':
+        inputs = tf.placeholder(tf.float32, [1, *_MODEL_SIZE, 3])
+        detections = model(inputs, training=False)
+        saver = tf.train.Saver(tf.global_variables(scope='yolo_v3_model'))
+
+        with tf.Session() as sess:
+            saver.restore(sess, './weights/model.ckpt')
+
+            win_name = 'Video detection'
+            cv2.namedWindow(win_name)
+            cap = cv2.VideoCapture(input_names[0])
+            frame_size = (cap.get(cv2.CAP_PROP_FRAME_WIDTH),
+                          cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+            fourcc = cv2.VideoWriter_fourcc(*'X264')
+            fps = cap.get(cv2.CAP_PROP_FPS)
+            out = cv2.VideoWriter('./detections/detections.mp4', fourcc, fps,
+                                  (int(frame_size[0]), int(frame_size[1])))
+
+            try:
+                while True:
+                    ret, frame = cap.read()
+                    if not ret:
+                        break
+                    resized_frame = cv2.resize(frame, dsize=_MODEL_SIZE[::-1],
+                                               interpolation=cv2.INTER_NEAREST)
+                    detection_result = sess.run(detections,
+                                                feed_dict={inputs: [resized_frame]})
+
+                    draw_frame(frame, frame_size, detection_result,
+                               class_names, _MODEL_SIZE)
+
+                    cv2.imshow(win_name, frame)
+
+                    key = cv2.waitKey(1) & 0xFF
+
+                    if key == ord('q'):
+                        break
+
+                    out.write(frame)
+            finally:
+                cv2.destroyAllWindows()
+                cap.release()
+                print('Detections have been saved successfully.')
+
+    else:
+        raise ValueError("Inappropriate data type. Please choose either 'video' or 'images'.")
+
+
+if __name__ == '__main__':
+    #main(sys.argv[1], float(sys.argv[2]), float(sys.argv[3]), sys.argv[4:])
+    main("images", 0.5, 0.5, "road.jpg")

detect_image.py

@@ -0,0 +1,44 @@
+# Yolo v3 image detection
+import os
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+import tensorflow as tf
+import sys
+import cv2
+
+from yolo_v3 import Yolo_v3
+from utils import load_images, load_class_names, draw_boxes
+
+_MODEL_SIZE = (416, 416)
+_CLASS_NAMES_FILE = 'coco.names'
+_MAX_OUTPUT_SIZE = 50
+
+detection_result = {}
+
+
+def main(iou_threshold, confidence_threshold, input_names):
+    global detection_result
+    class_names = load_class_names(_CLASS_NAMES_FILE)
+    n_classes = len(class_names)
+
+    model = Yolo_v3(n_classes=n_classes, model_size=_MODEL_SIZE,
+                    max_output_size=_MAX_OUTPUT_SIZE,
+                    iou_threshold=iou_threshold,
+                    confidence_threshold=confidence_threshold)
+
+
+    batch = load_images(input_names, model_size=_MODEL_SIZE)
+    inputs = tf.placeholder(tf.float32, [1, *_MODEL_SIZE, 3])
+    detections = model(inputs, training=False)
+    saver = tf.train.Saver(tf.global_variables(scope='yolo_v3_model'))
+
+    with tf.Session() as sess:
+        saver.restore(sess, './weights/model.ckpt')
+        detection_result = sess.run(detections, feed_dict={inputs: batch})
+
+    draw_boxes(input_names, detection_result, class_names, _MODEL_SIZE)
+
+    print('Detections have been saved successfully.')
+
+
+if __name__ == '__main__':
+    main(0.5, 0.5, "images/office.jpg")

detect_video.py

@@ -0,0 +1,71 @@
+# Yolo v3 video detection
+
+import os
+os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+import tensorflow as tf
+import sys
+import cv2
+
+from yolo_v3 import Yolo_v3
+from utils import load_images, load_class_names, draw_boxes, draw_frame
+
+_MODEL_SIZE = (416, 416)
+_CLASS_NAMES_FILE = 'coco.names'
+_MAX_OUTPUT_SIZE = 50
+
+detection_result = {}
+
+
+def main(iou_threshold, confidence_threshold, input_names):
+    global detection_result
+    class_names = load_class_names(_CLASS_NAMES_FILE)
+    n_classes = len(class_names)
+
+    model = Yolo_v3(n_classes=n_classes, model_size=_MODEL_SIZE,
+                    max_output_size=_MAX_OUTPUT_SIZE,
+                    iou_threshold=iou_threshold,
+                    confidence_threshold=confidence_threshold)
+
+    inputs = tf.placeholder(tf.float32, [1, *_MODEL_SIZE, 3])
+    detections = model(inputs, training=False)
+    saver = tf.train.Saver(tf.global_variables(scope='yolo_v3_model'))
+
+    with tf.Session() as sess:
+        saver.restore(sess, './weights/model.ckpt')
+
+        win_name = 'Video detection'
+        cv2.namedWindow(win_name)
+        cap = cv2.VideoCapture(input_names)
+        frame_size = (cap.get(cv2.CAP_PROP_FRAME_WIDTH), cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        fourcc = int(cap.get(cv2.CAP_PROP_FOURCC))
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        if not os.path.exists('detections'):
+            os.mkdir('detections')
+        head, tail = os.path.split(input_names)
+        name = './detections/'+tail[:-4]+'_yolo.mp4'
+        out = cv2.VideoWriter(name, fourcc, fps, (int(frame_size[0]), int(frame_size[1])))
+
+        try:
+            print("Show video")
+            while(cap.isOpened()):
+                ret, frame = cap.read()
+                if not ret:
+                    break
+                resized_frame = cv2.resize(frame, dsize=_MODEL_SIZE[::-1], interpolation=cv2.INTER_NEAREST)
+                detection_result = sess.run(detections, feed_dict={inputs: [resized_frame]})
+                draw_frame(frame, frame_size, detection_result, class_names, _MODEL_SIZE)
+                if ret == True:
+                    cv2.imshow(win_name, frame)
+                    out.write(frame)
+
+                if cv2.waitKey(1) & 0xFF == ord('q'):
+                    break
+
+        finally:
+            cv2.destroyAllWindows()
+            cap.release()
+            print('Detections have been saved successfully.')
+
+
+if __name__ == '__main__':
+    main(0.5, 0.5, "input/driving.mp4")