tensorflow-gpu==2.0.0

malin9402 · Oct 4, 2019 · 4128617 · 4128617
1 parent b353814
commit 4128617
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 13 deletions.
diff --git a/4-Object_Detection/YOLOV3/core/config.py b/4-Object_Detection/YOLOV3/core/config.py
@@ -23,7 +23,7 @@
 __C.YOLO                      = edict()
 
 # Set the class name
-__C.YOLO.CLASSES              = "./data/classes/coco.names"
+__C.YOLO.CLASSES              = "./data/classes/yymnist.names"
 __C.YOLO.ANCHORS              = "./data/anchors/basline_anchors.txt"
 __C.YOLO.STRIDES              = [8, 16, 32]
 __C.YOLO.ANCHOR_PER_SCALE     = 3

diff --git a/4-Object_Detection/YOLOV3/image_demo.py b/4-Object_Detection/YOLOV3/image_demo.py
@@ -17,10 +17,8 @@
 import tensorflow as tf
 from core.yolov3 import YOLOv3, decode
 from PIL import Image
-from core.config import cfg
 
 input_size   = 416
-NUM_CLASS    = len(utils.read_class_names(cfg.YOLO.CLASSES))
 image_path   = "./docs/kite.jpg"
 
 input_layer  = tf.keras.layers.Input([input_size, input_size, 3])
@@ -36,14 +34,15 @@
 bbox_tensors = []
 for i, fm in enumerate(feature_maps):
     bbox_tensor = decode(fm, i)
-    bbox_tensors.append(tf.reshape(bbox_tensor, (-1, 5+NUM_CLASS)))
+    bbox_tensors.append(bbox_tensor)
 
-bbox_tensors = tf.concat(bbox_tensors, axis=0)
 model = tf.keras.Model(input_layer, bbox_tensors)
 utils.load_weights(model, "./yolov3.weights")
 model.summary()
 
 pred_bbox = model.predict(image_data)
+pred_bbox = [tf.reshape(x, (-1, tf.shape(x)[-1])) for x in pred_bbox]
+pred_bbox = tf.concat(pred_bbox, axis=0)
 bboxes = utils.postprocess_boxes(pred_bbox, original_image_size, input_size, 0.3)
 bboxes = utils.nms(bboxes, 0.45, method='nms')
 

diff --git a/4-Object_Detection/YOLOV3/test.py b/4-Object_Detection/YOLOV3/test.py
@@ -42,9 +42,8 @@
 bbox_tensors = []
 for i, fm in enumerate(feature_maps):
     bbox_tensor = decode(fm, i)
-    bbox_tensors.append(tf.reshape(bbox_tensor, (-1, 5+NUM_CLASS)))
+    bbox_tensors.append(bbox_tensor)
 
-bbox_tensors = tf.concat(bbox_tensors, axis=0)
 model = tf.keras.Model(input_layer, bbox_tensors)
 model.load_weights("./yolov3")
 
@@ -81,6 +80,8 @@
         image_data = image_data[np.newaxis, ...].astype(np.float32)
 
         pred_bbox = model.predict(image_data)
+        pred_bbox = [tf.reshape(x, (-1, tf.shape(x)[-1])) for x in pred_bbox]
+        pred_bbox = tf.concat(pred_bbox, axis=0)
         bboxes = utils.postprocess_boxes(pred_bbox, image_size, INPUT_SIZE, cfg.TEST.SCORE_THRESHOLD)
         bboxes = utils.nms(bboxes, cfg.TEST.IOU_THRESHOLD, method='nms')
 

diff --git a/4-Object_Detection/YOLOV3/video_demo.py b/4-Object_Detection/YOLOV3/video_demo.py
@@ -20,7 +20,7 @@
 
 
 video_path      = "./docs/road.mp4"
-# video_path      = 0
+video_path      = 0
 num_classes     = 80
 input_size      = 416
 
@@ -30,9 +30,8 @@
 bbox_tensors = []
 for i, fm in enumerate(feature_maps):
     bbox_tensor = decode(fm, i)
-    bbox_tensors.append(tf.reshape(bbox_tensor, (-1, 5+num_classes)))
+    bbox_tensors.append(bbox_tensor)
 
-bbox_tensors = tf.concat(bbox_tensors, axis=0)
 model = tf.keras.Model(input_layer, bbox_tensors)
 utils.load_weights(model, "./yolov3.weights")
 model.summary()
@@ -46,15 +45,18 @@
     frame_size = frame.shape[:2]
     image_data = utils.image_preporcess(np.copy(frame), [input_size, input_size])
     image_data = image_data[np.newaxis, ...].astype(np.float32)
-    prev_time = time.time()
 
+    prev_time = time.time()
     pred_bbox = model.predict(image_data)
+    curr_time = time.time()
+    exec_time = curr_time - prev_time
+
+    pred_bbox = [tf.reshape(x, (-1, tf.shape(x)[-1])) for x in pred_bbox]
+    pred_bbox = tf.concat(pred_bbox, axis=0)
     bboxes = utils.postprocess_boxes(pred_bbox, frame_size, input_size, 0.3)
     bboxes = utils.nms(bboxes, 0.45, method='nms')
     image = utils.draw_bbox(frame, bboxes)
 
-    curr_time = time.time()
-    exec_time = curr_time - prev_time
     result = np.asarray(image)
     info = "time: %.2f ms" %(1000*exec_time)
     cv2.putText(result, text=info, org=(50, 70), fontFace=cv2.FONT_HERSHEY_SIMPLEX,
@@ -64,3 +66,5 @@
     cv2.imshow("result", result)
     if cv2.waitKey(1) & 0xFF == ord('q'): break
 
+
+