Add files via upload

Author: cchighman

watch_cameras.py

@@ -13,11 +13,17 @@
 from PIL import Image
 
 from PriusImageCache import ImageDeduplication
-from imageai.Detection import ObjectDetection
+from yolo4 import Yolo4
 from imageai.Prediction.Custom import CustomImagePrediction
 from prius_color import has_prius_color_from_array
 from prius_color import has_prius_contour_from_array
 
+import colorsys
+from keras import backend as K
+from keras.models import load_model
+from keras.layers import Input
+
+
 import warnings
 warnings.filterwarnings('ignore', category=DeprecationWarning)
 warnings.filterwarnings('ignore', category=FutureWarning)
@@ -30,8 +36,6 @@
 tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
 deprecation._PER_MODULE_WARNING_LIMIT = 0
 
-
-
 cams = []
 
 cam_threads = []
@@ -49,27 +53,28 @@
                 help='Output Path')
 ap.add_argument("-a", "--accuracy", default=50,
                 help="predict accuracy")
-ap.add_argument("-y", "--detectspeed", default='normal',
-                help="detection speed")
 ap.add_argument("-z", "--predictspeed", default='normal',
                 help="prediction speed")
 ap.add_argument("-n", "--name", default='predictor1',
                 help="name")
 ap.add_argument("-m", "--models", default='./',
                 help='model path')
-ap.add_argument("-q", "--model", default='model_ex-011_acc-0.869792.h5',
+ap.add_argument("-q", "--model", default='model_ex-006_acc-0.994420.h5',
                 help='model')
 
 args = vars(ap.parse_args())
 
 if os.path.isdir(args['output']) is False:
 	os.mkdir(args['output'])
 
-detector = ObjectDetection()
-detector.setModelTypeAsYOLOv3()
-detector.setModelPath(args['models'] + "yolo.h5")
-detector.loadModel(detection_speed=args["detectspeed"])
-custom_objects = detector.CustomObjects(car=True)
+model_path = args['models'] + 'yolo4_weight.h5'
+anchors_path = args['models'] + 'yolo4_anchors.txt'
+classes_path = args['models'] + 'coco_classes.txt'
+
+score = 0.5
+iou = 0.5
+model_image_size = (608, 608)
+yolo4_model = Yolo4(score, iou, anchors_path, classes_path, model_path)
 
 prediction = CustomImagePrediction()
 prediction.setModelTypeAsResNet()
@@ -123,18 +128,14 @@ def watch_camera(cam):
 			decoded = cv2.imdecode(np.frombuffer(img_data, np.uint8), -1)
 
 			#start1 = time.time()
-			result = detector.detectCustomObjectsFromImage(custom_objects=custom_objects,
-			                                               input_type="array",
-			                                               extract_detected_objects=True,
-			                                               input_image=decoded,
-			                                               output_type="array",
-			                                               minimum_percentage_probability=50)
+			result = yolo4_model.detect_image(decoded, model_image_size=model_image_size)
+
 			#start2 = time.time()
 
 			#print("Detection Time: " + str(start2 - start1))
 
-			for arr in result[1]:
-				(x1, y1, x2, y2) = arr["box_points"]
+			for car in result:
+				(x1, y1, x2, y2) = car["box"]
 				img = decoded[y1:y2, x1:x2]
 
 				#start2 = time.time()

yolo4.py

@@ -0,0 +1,134 @@
+import os
+import colorsys
+
+import numpy as np
+from keras import backend as K
+from keras.models import load_model
+from keras.layers import Input
+
+from yolo4_model.model import yolo_eval, yolo4_body
+from yolo4_model.utils import letterbox_image
+
+class Yolo4(object):
+    def get_class(self):
+        classes_path = os.path.expanduser(self.classes_path)
+        with open(classes_path) as f:
+            class_names = f.readlines()
+        class_names = [c.strip() for c in class_names]
+        return class_names
+
+    def get_anchors(self):
+        anchors_path = os.path.expanduser(self.anchors_path)
+        with open(anchors_path) as f:
+            anchors = f.readline()
+        anchors = [float(x) for x in anchors.split(',')]
+        return np.array(anchors).reshape(-1, 2)
+
+    def load_yolo(self):
+        model_path = os.path.expanduser(self.model_path)
+        assert model_path.endswith('.h5'), 'Keras model or weights must be a .h5 file.'
+
+        self.class_names = self.get_class()
+        self.anchors = self.get_anchors()
+
+        num_anchors = len(self.anchors)
+        num_classes = len(self.class_names)
+
+        # Generate colors for drawing bounding boxes.
+        hsv_tuples = [(x / len(self.class_names), 1., 1.)
+                      for x in range(len(self.class_names))]
+        self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
+        self.colors = list(
+            map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
+                self.colors))
+
+        self.sess = K.get_session()
+
+        # Load model, or construct model and load weights.
+        self.yolo4_model = yolo4_body(Input(shape=(608, 608, 3)), num_anchors//3, num_classes)
+        self.yolo4_model.load_weights(model_path)
+
+        print('{} model, anchors, and classes loaded.'.format(model_path))
+
+        if self.gpu_num>=2:
+            self.yolo4_model = multi_gpu_model(self.yolo4_model, gpus=self.gpu_num)
+
+        self.input_image_shape = K.placeholder(shape=(2, ))
+        self.boxes, self.scores, self.classes = yolo_eval(self.yolo4_model.output, self.anchors,
+                len(self.class_names), self.input_image_shape,
+                score_threshold=self.score)
+
+    def __init__(self, score, iou, anchors_path, classes_path, model_path, gpu_num=1):
+        self.score = score
+        self.iou = iou
+        self.anchors_path = anchors_path
+        self.classes_path = classes_path
+        self.model_path = model_path
+        self.gpu_num = gpu_num
+        self.load_yolo()
+
+    def close_session(self):
+        self.sess.close()
+
+    def detect_image(self, image, model_image_size=(608, 608)):
+        start = timer()
+
+        boxed_image = letterbox_image(image, tuple(reversed(model_image_size)))
+        image_data = np.array(boxed_image, dtype='float32')
+
+
+        image_data /= 255.
+        image_data = np.expand_dims(image_data, 0)  # Add batch dimension.
+
+        out_boxes, out_scores, out_classes = self.sess.run(
+            [self.boxes, self.scores, self.classes],
+            feed_dict={
+                self.yolo4_model.input: image_data,
+                self.input_image_shape: [image.size[1], image.size[0]],
+                K.learning_phase(): 0
+            })
+
+        print('Found {} boxes for {}'.format(len(out_boxes), 'img'))
+
+        #font = ImageFont.truetype(font='font/FiraMono-Medium.otf',
+         #           size=np.floor(3e-2 * image.size[1] + 0.5).astype('int32'))
+        #thickness = (image.size[0] + image.size[1]) // 300
+        cars = []
+        for i, c in reversed(list(enumerate(out_classes))):
+            predicted_class = self.class_names[c]
+
+            if "car" in predicted_class:
+                car =  {
+                "box": out_boxes[i],
+                "score": out_scores[i]
+                }
+                cars.append(car)
+        return cars
+    '''     
+                label = '{} {:.2f}'.format(predicted_class, score)
+                draw = ImageDraw.Draw(image)
+                label_size = draw.textsize(label, font)
+    
+                top, left, bottom, right = box
+                top = max(0, np.floor(top + 0.5).astype('int32'))
+                left = max(0, np.floor(left + 0.5).astype('int32'))
+                bottom = min(image.size[1], np.floor(bottom + 0.5).astype('int32'))
+                right = min(image.size[0], np.floor(right + 0.5).astype('int32'))
+                print(label, (left, top), (right, bottom))
+    
+                if top - label_size[1] >= 0:
+                    text_origin = np.array([left, top - label_size[1]])
+                else:
+                    text_origin = np.array([left, top + 1])
+    
+                # My kingdom for a good redistributable image drawing library.
+                for i in range(thickness):
+                    draw.rectangle(
+                        [left + i, top + i, right - i, bottom - i],
+                        outline=self.colors[c])
+                draw.rectangle(
+                    [tuple(text_origin), tuple(text_origin + label_size)],
+                    fill=self.colors[c])
+                draw.text(text_origin, label, fill=(0, 0, 0), font=font)
+                del draw
+    '''