You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@mxnet.apache.org by GitBox <gi...@apache.org> on 2018/05/22 22:08:46 UTC
[GitHub] zhreshold closed pull request #9808: [MXNET-460] Ssd camera demo
zhreshold closed pull request #9808: [MXNET-460] Ssd camera demo
URL: https://github.com/apache/incubator-mxnet/pull/9808
This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:
As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):
diff --git a/example/ssd/README.md b/example/ssd/README.md
index 0b970923e44..55387c5fd2d 100644
--- a/example/ssd/README.md
+++ b/example/ssd/README.md
@@ -17,6 +17,7 @@ remarkable traits of MXNet.
Due to the permission issue, this example is maintained in this [repository](https://github.com/zhreshold/mxnet-ssd) separately. You can use the link regarding specific per example [issues](https://github.com/zhreshold/mxnet-ssd/issues).
### What's new
+* Added live camera capture and detection display (run with --camera flag)
* Added multiple trained models.
* Added a much simpler way to compose network from mainstream classification networks (resnet, inception...) and [Guide](symbol/README.md).
* Update to the latest version according to caffe version, with 5% mAP increase.
@@ -84,6 +85,13 @@ python demo.py --cpu --network resnet50 --data-shape 512
```
* Check `python demo.py --help` for more options.
+### Live Camera detection
+
+Use `init.sh` to download the trained model.
+You can use `./demo.py --camera` to use a video capture device with opencv such as a webcam. This
+will open a window that will display the camera output together with the detections. You can play
+with the detection threshold to get more or less detections.
+
### Train the model
This example only covers training on Pascal VOC dataset. Other datasets should
be easily supported by adding subclass derived from class `Imdb` in `dataset/imdb.py`.
diff --git a/example/ssd/dataset/cv2Iterator.py b/example/ssd/dataset/cv2Iterator.py
new file mode 100644
index 00000000000..469faeac828
--- /dev/null
+++ b/example/ssd/dataset/cv2Iterator.py
@@ -0,0 +1,56 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import mxnet as mx
+import numpy as np
+import cv2
+
+
+class CameraIterator():
+ """
+ An iterator that captures frames with opencv or the specified capture
+ """
+ def __init__(self, capture=cv2.VideoCapture(0), frame_resize=None):
+ self._capture = capture
+ self._frame_resize = frame_resize
+ if frame_resize:
+ assert isinstance(frame_resize, tuple) and (len(tuple) == 2), "frame_resize should be a tuple of (x,y)"
+ self._frame_shape = (1, 3, frame_resize[0], frame_resize[1])
+ else:
+ self._frame_shape = (1, 3,
+ int(self._capture.get(cv2.CAP_PROP_FRAME_WIDTH)),
+ int(self._capture.get(cv2.CAP_PROP_FRAME_HEIGHT)))
+
+ def __iter__(self):
+ return self
+
+ def __next__(self):
+ ret, frame = self._capture.read()
+ if cv2.waitKey(1) & 0xFF == ord('q') or ret is not True:
+ raise StopIteration
+ if self._frame_resize:
+ frame = cv2.resize(frame, (self._frame_resize[0], self._frame_resize[1]))
+ return frame
+
+ def __enter__(self):
+ pass
+
+ def __exit__(self, exc_type, exc_alue, traceback):
+ self.close()
+
+ def close(self):
+ self._capture.release()
diff --git a/example/ssd/demo.py b/example/ssd/demo.py
index 0480bdd658b..4ae8b350742 100755
--- a/example/ssd/demo.py
+++ b/example/ssd/demo.py
@@ -25,6 +25,9 @@
import sys
from detect.detector import Detector
from symbol.symbol_factory import get_symbol
+from dataset.cv2Iterator import CameraIterator
+import logging
+import cv2
def get_detector(net, prefix, epoch, data_shape, mean_pixels, ctx, num_class,
nms_thresh=0.5, force_nms=True, nms_topk=400):
@@ -72,6 +75,8 @@ def parse_args():
type=str, nargs='?')
parser.add_argument('--epoch', dest='epoch', help='epoch of trained model',
default=0, type=int)
+ parser.add_argument('--batch-size', dest='batch_size', help='batch size',
+ default=1, type=int)
parser.add_argument('--prefix', dest='prefix', help='trained model prefix',
default=os.path.join(os.getcwd(), 'model', 'ssd_'),
type=str)
@@ -102,6 +107,8 @@ def parse_args():
car, cat, chair, cow, diningtable, dog, horse, motorbike, \
person, pottedplant, sheep, sofa, train, tvmonitor',
help='string of comma separated names, or text filename')
+ parser.add_argument('--camera', action='store_true',
+ help="use camera for image capturing")
args = parser.parse_args()
return args
@@ -131,13 +138,59 @@ def parse_data_shape(data_shape_str):
raise ValueError("Unexpected data_shape: %s", data_shape_str)
return data_shape
-if __name__ == '__main__':
- args = parse_args()
- if args.cpu:
- ctx = mx.cpu()
- else:
- ctx = mx.gpu(args.gpu_id)
+def draw_detection(frame, det, class_names):
+ (klass, score, x0, y0, x1, y1) = det
+ klass_name = class_names[int(klass)]
+ h = frame.shape[0]
+ w = frame.shape[1]
+ # denormalize detections from [0,1] to the frame size
+ p0 = tuple(map(int, (x0*w,y0*h)))
+ p1 = tuple(map(int, (x1*w,y1*h)))
+ logging.info("detection: %s %s", klass_name, score)
+ cv2.rectangle(frame, p0, p1, (0,0,255), 2)
+ # Where to draw the text, a few pixels above the top y coordinate
+ tp0 = (p0[0], p0[1]-5)
+ draw_text = "{} {}".format(klass_name, score)
+ cv2.putText(frame, draw_text, tp0, cv2.FONT_HERSHEY_COMPLEX_SMALL, 0.5, (0,0,255))
+
+def network_path(prefix, network, data_shape):
+ return "{}{}_{}".format(prefix, network, data_shape)
+
+def run_camera(args,ctx):
+ assert args.batch_size == 1, "only batch size of 1 is supported"
+ logging.info("Detection threshold is {}".format(args.thresh))
+ iter = CameraIterator()
+ class_names = parse_class_names(args.class_names)
+ mean_pixels = (args.mean_r, args.mean_g, args.mean_b)
+ data_shape = int(args.data_shape)
+ batch_size = int(args.batch_size)
+ detector = Detector(
+ get_symbol(args.network, data_shape, num_classes=len(class_names)),
+ network_path(args.prefix, args.network, data_shape),
+ args.epoch,
+ data_shape,
+ mean_pixels,
+ batch_size,
+ ctx
+ )
+ for frame in iter:
+ logging.info("Frame info: shape %s type %s", frame.shape, frame.dtype)
+ logging.info("Generating batch")
+ data_batch = detector.create_batch(frame)
+ logging.info("Detecting objects")
+ detections_batch = detector.detect_batch(data_batch)
+ #detections = [mx.nd.array((1,1,0.2,0.2,0.4,0.4))]
+ detections = detections_batch[0]
+ logging.info("%d detections", len(detections))
+ for det in detections:
+ obj = det.asnumpy()
+ (klass, score, x0, y0, x1, y1) = obj
+ if score > args.thresh:
+ draw_detection(frame, obj, class_names)
+ cv2.imshow('frame', frame)
+
+def run_images(args,ctx):
# parse image list
image_list = [i.strip() for i in args.images.split(',')]
assert len(image_list) > 0, "No valid image specified to detect"
@@ -156,3 +209,22 @@ def parse_data_shape(data_shape_str):
# run detection
detector.detect_and_visualize(image_list, args.dir, args.extension,
class_names, args.thresh, args.show_timer)
+
+def main():
+ logging.getLogger().setLevel(logging.INFO)
+ logging.basicConfig(format='%(asctime)-15s %(message)s')
+ args = parse_args()
+ if args.cpu:
+ ctx = mx.cpu()
+ else:
+ ctx = mx.gpu(args.gpu_id)
+
+ if args.camera:
+ run_camera(args, ctx)
+ else:
+ run_images(args, ctx)
+ return 0
+
+if __name__ == '__main__':
+ sys.exit(main())
+
diff --git a/example/ssd/detect/detector.py b/example/ssd/detect/detector.py
index 8d72eaa8669..1b5e8cb76ee 100644
--- a/example/ssd/detect/detector.py
+++ b/example/ssd/detect/detector.py
@@ -21,6 +21,9 @@
from dataset.testdb import TestDB
from dataset.iterator import DetIter
import logging
+import cv2
+from mxnet.io import DataBatch, DataDesc
+
class Detector(object):
"""
@@ -58,8 +61,25 @@ def __init__(self, symbol, model_prefix, epoch, data_shape, mean_pixels, \
self.mod.bind(data_shapes=[('data', (batch_size, 3, data_shape[0], data_shape[1]))])
self.mod.set_params(args, auxs)
self.mean_pixels = mean_pixels
+ self.mean_pixels_nd = mx.nd.array(mean_pixels).reshape((3,1,1))
- def detect(self, det_iter, show_timer=False):
+ def create_batch(self, frame):
+ """
+ :param frame: an (w,h,channels) numpy array (image)
+ :return: DataBatch of (1,channels,data_shape,data_shape)
+ """
+ frame_resize = mx.nd.array(cv2.resize(frame, (self.data_shape[0], self.data_shape[1])))
+ #frame_resize = mx.img.imresize(frame, self.data_shape[0], self.data_shape[1], cv2.INTER_LINEAR)
+ # Change dimensions from (w,h,channels) to (channels, w, h)
+ frame_t = mx.nd.transpose(frame_resize, axes=(2,0,1))
+ frame_norm = frame_t - self.mean_pixels_nd
+ # Add dimension for batch, results in (1,channels,w,h)
+ batch_frame = [mx.nd.expand_dims(frame_norm, axis=0)]
+ batch_shape = [DataDesc('data', batch_frame[0].shape)]
+ batch = DataBatch(data=batch_frame, provide_data=batch_shape)
+ return batch
+
+ def detect_iter(self, det_iter, show_timer=False):
"""
detect all images in iterator
@@ -86,6 +106,17 @@ def detect(self, det_iter, show_timer=False):
result = Detector.filter_positive_detections(detections)
return result
+ def detect_batch(self, batch):
+ """
+ Return detections for batch
+ :param batch:
+ :return:
+ """
+ self.mod.forward(batch, is_train=False)
+ detections = self.mod.get_outputs()[0]
+ positive_detections = Detector.filter_positive_detections(detections)
+ return positive_detections
+
def im_detect(self, im_list, root_dir=None, extension=None, show_timer=False):
"""
wrapper for detecting multiple images
@@ -108,7 +139,7 @@ def im_detect(self, im_list, root_dir=None, extension=None, show_timer=False):
test_db = TestDB(im_list, root_dir=root_dir, extension=extension)
test_iter = DetIter(test_db, 1, self.data_shape, self.mean_pixels,
is_train=False)
- return self.detect(test_iter, show_timer)
+ return self.detect_iter(test_iter, show_timer)
def visualize_detection(self, img, dets, classes=[], thresh=0.6):
"""
@@ -197,7 +228,6 @@ def detect_and_visualize(self, im_list, root_dir=None, extension=None,
----------
"""
- import cv2
dets = self.im_detect(im_list, root_dir, extension, show_timer=show_timer)
if not isinstance(im_list, list):
im_list = [im_list]
diff --git a/example/ssd/init.sh b/example/ssd/init.sh
index 53104bb7d78..863ba871da8 100755
--- a/example/ssd/init.sh
+++ b/example/ssd/init.sh
@@ -53,6 +53,7 @@ function download_demo_images() {
popd
}
-download_pascal_voc
+# Uncomment to download training dataset
+#download_pascal_voc
download_model
download_demo_images
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services