From 8eb0d400e80f3dd135313834ce103bd0f9e62dbc Mon Sep 17 00:00:00 2001
From: Edmond Yoo <hj3yoo@uwaterloo.ca>
Date: Fri, 14 Sep 2018 03:26:38 +0000
Subject: [PATCH] Update for training

---
 darknet.py |   47 ++++++++++++++++++++++++++++++++++++++++++++++-
 1 files changed, 46 insertions(+), 1 deletions(-)

diff --git a/darknet.py b/darknet.py
index 4bca5f5..ef5f98d 100644
--- a/darknet.py
+++ b/darknet.py
@@ -31,6 +31,7 @@
 import math
 import random
 import os
+import cv2
 
 def sample(probs):
     s = sum(probs)
@@ -421,5 +422,49 @@
             print("Unable to show image: "+str(e))
     return detections
 
+
+def capture(thresh=.5, hier_thresh=.5, nms=.45, configPath="./cfg/yolov3.cfg", weightPath="yolov3.weights",
+            metaPath="./data/coco.data", showImage=True, makeImageOnly=False, initOnly=False):
+    global metaMain, netMain, altNames  # pylint: disable=W0603
+    netMain = load_net_custom(configPath.encode("ascii"), weightPath.encode("ascii"), 0, 1)  # batch size = 1
+    metaMain = load_meta(metaPath.encode("ascii"))
+
+    num = c_int(0)
+    pnum = pointer(num)
+    num = pnum[0]
+
+    capture = cv2.VideoCapture('../data/test3.mp4')
+    print(capture.get(cv2.CAP_PROP_FPS))
+
+    capture.set(cv2.CAP_PROP_FRAME_WIDTH, 1024)
+    capture.set(cv2.CAP_PROP_FRAME_HEIGHT, 768)
+
+    while True:
+        ret, frame = capture.read()
+        im, arr = array_to_image(frame)
+        predict_image(netMain, im)
+        dets = get_network_boxes(netMain, im.w, im.h, thresh, hier_thresh, None, 0, pnum, 1)
+        if nms:
+            do_nms_sort(dets, num, metaMain.classes, nms)
+        res = []
+        for j in range(num):
+            for i in range(metaMain.classes):
+                if dets[j].prob[i] > 0:
+                    b = dets[j].bbox
+                    nameTag = metaMain.names[i]
+                    res.append((nameTag, dets[j].prob[i], (b.x, b.y, b.w, b.h)))
+        print(res)
+        cv2.imshow('frame', frame)
+        if cv2.waitKey(1) & 0xFF == ord('q'):
+            break
+
+    capture.release()
+    cv2.destroyAllWindows()
+
+
 if __name__ == "__main__":
-    print(performDetect())
+    performDetect(imagePath="../data/test1.jpg", thresh=0.25, configPath="./cfg/tiny_yolo.cfg",
+                  weightPath="./weights/second_general/tiny_yolo_17000.weights",
+                  metaPath="./data/obj.data", showImage=True, makeImageOnly=False, initOnly=False)
+    #print(performDetect(showImage=False))
+    #capture()

--
Gitblit v1.10.0