From bf777b33a0b455cbd214682886b4d9bc0aa040a9 Mon Sep 17 00:00:00 2001 From: sipp11 Date: Mon, 21 Oct 2019 14:47:51 +0900 Subject: [PATCH] Initial commit --- README.md | 26 +++ data/mscoco_label_map.pbtxt | 394 ++++++++++++++++++++++++++++++++++++ pip.txt | 8 + run.py | 93 +++++++++ 4 files changed, 521 insertions(+) create mode 100644 README.md create mode 100644 data/mscoco_label_map.pbtxt create mode 100644 pip.txt create mode 100644 run.py diff --git a/README.md b/README.md new file mode 100644 index 0000000..aee5761 --- /dev/null +++ b/README.md @@ -0,0 +1,26 @@ +TO get this code to work + +1. clone models + + git clone https://github.com/tensorflow/models + +2. run [`protoc`](http://google.github.io/proto-lens/installing-protoc.html) in `models/research` + + protoc object_detection/protos/*.proto --python_out=. + +3. Add models/research into PYTHONPATH + + export PYTHONPATH=/home/sipp11/dev/tf1-test/models/research:$PYTHONPATH + +4. Copy `mscoco_label_map.pbtxt` from `/models/research/object_detection/data` to `./data/` However, `.pbtxt` is just an id & label conversion as following: + + item { + id: 1 + display_name: "person" + } + item { + id: 2 + display_name: "bicycle" + } + +5. It should work! diff --git a/data/mscoco_label_map.pbtxt b/data/mscoco_label_map.pbtxt new file mode 100644 index 0000000..6244d2f --- /dev/null +++ b/data/mscoco_label_map.pbtxt @@ -0,0 +1,394 @@ +item { + id: 1 + display_name: "person" +} +item { + id: 2 + display_name: "bicycle" +} +item { + id: 3 + display_name: "car" +} +item { + id: 4 + display_name: "motorcycle" +} +item { + id: 5 + display_name: "airplane" +} +item { + id: 6 + display_name: "bus" +} +item { + name: "/m/07jdr" + id: 7 + display_name: "train" +} +item { + name: "/m/07r04" + id: 8 + display_name: "truck" +} +item { + name: "/m/019jd" + id: 9 + display_name: "boat" +} +item { + name: "/m/015qff" + id: 10 + display_name: "traffic light" +} +item { + name: "/m/01pns0" + id: 11 + display_name: "fire hydrant" +} +item { + name: "/m/02pv19" + id: 13 + display_name: "stop sign" +} +item { + name: "/m/015qbp" + id: 14 + display_name: "parking meter" +} +item { + name: "/m/0cvnqh" + id: 15 + display_name: "bench" +} +item { + name: "/m/015p6" + id: 16 + display_name: "bird" +} +item { + name: "/m/01yrx" + id: 17 + display_name: "cat" +} +item { + name: "/m/0bt9lr" + id: 18 + display_name: "dog" +} +item { + name: "/m/03k3r" + id: 19 + display_name: "horse" +} +item { + name: "/m/07bgp" + id: 20 + display_name: "sheep" +} +item { + name: "/m/01xq0k1" + id: 21 + display_name: "cow" +} +item { + name: "/m/0bwd_0j" + id: 22 + display_name: "elephant" +} +item { + name: "/m/01dws" + id: 23 + display_name: "bear" +} +item { + name: "/m/0898b" + id: 24 + display_name: "zebra" +} +item { + name: "/m/03bk1" + id: 25 + display_name: "giraffe" +} +item { + name: "/m/01940j" + id: 27 + display_name: "backpack" +} +item { + name: "/m/0hnnb" + id: 28 + display_name: "umbrella" +} +item { + name: "/m/080hkjn" + id: 31 + display_name: "handbag" +} +item { + name: "/m/01rkbr" + id: 32 + display_name: "tie" +} +item { + name: "/m/01s55n" + id: 33 + display_name: "suitcase" +} +item { + name: "/m/02wmf" + id: 34 + display_name: "frisbee" +} +item { + name: "/m/071p9" + id: 35 + display_name: "skis" +} +item { + name: "/m/06__v" + id: 36 + display_name: "snowboard" +} +item { + name: "/m/018xm" + id: 37 + display_name: "sports ball" +} +item { + name: "/m/02zt3" + id: 38 + display_name: "kite" +} +item { + name: "/m/03g8mr" + id: 39 + display_name: "baseball bat" +} +item { + name: "/m/03grzl" + id: 40 + display_name: "baseball glove" +} +item { + name: "/m/06_fw" + id: 41 + display_name: "skateboard" +} +item { + name: "/m/019w40" + id: 42 + display_name: "surfboard" +} +item { + name: "/m/0dv9c" + id: 43 + display_name: "tennis racket" +} +item { + name: "/m/04dr76w" + id: 44 + display_name: "bottle" +} +item { + name: "/m/09tvcd" + id: 46 + display_name: "wine glass" +} +item { + name: "/m/08gqpm" + id: 47 + display_name: "cup" +} +item { + name: "/m/0dt3t" + id: 48 + display_name: "fork" +} +item { + name: "/m/04ctx" + id: 49 + display_name: "knife" +} +item { + name: "/m/0cmx8" + id: 50 + display_name: "spoon" +} +item { + name: "/m/04kkgm" + id: 51 + display_name: "bowl" +} +item { + name: "/m/09qck" + id: 52 + display_name: "banana" +} +item { + name: "/m/014j1m" + id: 53 + display_name: "apple" +} +item { + name: "/m/0l515" + id: 54 + display_name: "sandwich" +} +item { + name: "/m/0cyhj_" + id: 55 + display_name: "orange" +} +item { + name: "/m/0hkxq" + id: 56 + display_name: "broccoli" +} +item { + name: "/m/0fj52s" + id: 57 + display_name: "carrot" +} +item { + name: "/m/01b9xk" + id: 58 + display_name: "hot dog" +} +item { + name: "/m/0663v" + id: 59 + display_name: "pizza" +} +item { + name: "/m/0jy4k" + id: 60 + display_name: "donut" +} +item { + name: "/m/0fszt" + id: 61 + display_name: "cake" +} +item { + name: "/m/01mzpv" + id: 62 + display_name: "chair" +} +item { + name: "/m/02crq1" + id: 63 + display_name: "couch" +} +item { + name: "/m/03fp41" + id: 64 + display_name: "potted plant" +} +item { + name: "/m/03ssj5" + id: 65 + display_name: "bed" +} +item { + name: "/m/04bcr3" + id: 67 + display_name: "dining table" +} +item { + name: "/m/09g1w" + id: 70 + display_name: "toilet" +} +item { + name: "/m/07c52" + id: 72 + display_name: "tv" +} +item { + name: "/m/01c648" + id: 73 + display_name: "laptop" +} +item { + name: "/m/020lf" + id: 74 + display_name: "mouse" +} +item { + name: "/m/0qjjc" + id: 75 + display_name: "remote" +} +item { + name: "/m/01m2v" + id: 76 + display_name: "keyboard" +} +item { + name: "/m/050k8" + id: 77 + display_name: "cell phone" +} +item { + name: "/m/0fx9l" + id: 78 + display_name: "microwave" +} +item { + name: "/m/029bxz" + id: 79 + display_name: "oven" +} +item { + name: "/m/01k6s3" + id: 80 + display_name: "toaster" +} +item { + name: "/m/0130jx" + id: 81 + display_name: "sink" +} +item { + name: "/m/040b_t" + id: 82 + display_name: "refrigerator" +} +item { + name: "/m/0bt_c3" + id: 84 + display_name: "book" +} +item { + name: "/m/01x3z" + id: 85 + display_name: "clock" +} +item { + name: "/m/02s195" + id: 86 + display_name: "vase" +} +item { + name: "/m/01lsmm" + id: 87 + display_name: "scissors" +} +item { + name: "/m/0kmg4" + id: 88 + display_name: "teddy bear" +} +item { + name: "/m/03wvsk" + id: 89 + display_name: "hair drier" +} +item { + name: "/m/012xff" + id: 90 + display_name: "toothbrush" +} diff --git a/pip.txt b/pip.txt new file mode 100644 index 0000000..a36847b --- /dev/null +++ b/pip.txt @@ -0,0 +1,8 @@ +Cython +contextlib2 +pillow +lxml +jupyter +matplotlib +protobuf +tensorflow==1.14.0 diff --git a/run.py b/run.py new file mode 100644 index 0000000..4304d97 --- /dev/null +++ b/run.py @@ -0,0 +1,93 @@ +import numpy as np +import os +import six.moves.urllib as urllib +import sys +import tarfile +import tensorflow as tf +import zipfile + +from collections import defaultdict +from io import StringIO +from matplotlib import pyplot as plt +from PIL import Image + +import cv2 + +# cap = cv2.VideoCapture(0) +cap = cv2.VideoCapture("/home/sipp11/Desktop/5min.mp4") + + +sys.path.append("..") + +from utils import label_map_util + +from utils import visualization_utils as vis_util + +MODEL_NAME = "ssd_mobilenet_v1_coco_11_06_2017" +MODEL_FILE = MODEL_NAME + ".tar.gz" +DOWNLOAD_BASE = "http://download.tensorflow.org/models/object_detection/" + +# Path to frozen detection graph. This is the actual model that is used for the object detection. +PATH_TO_CKPT = MODEL_NAME + "/frozen_inference_graph.pb" + +# List of the strings that is used to add correct label for each box. +PATH_TO_LABELS = os.path.join("data", "mscoco_label_map.pbtxt") + +NUM_CLASSES = 90 + +opener = urllib.request.URLopener() +opener.retrieve(DOWNLOAD_BASE + MODEL_FILE, MODEL_FILE) +tar_file = tarfile.open(MODEL_FILE) +for file in tar_file.getmembers(): + file_name = os.path.basename(file.name) + if "frozen_inference_graph.pb" in file_name: + tar_file.extract(file, os.getcwd()) + +detection_graph = tf.Graph() +with detection_graph.as_default(): + od_graph_def = tf.compat.v1.GraphDef() + with tf.io.gfile.GFile(PATH_TO_CKPT, "rb") as fid: + serialized_graph = fid.read() + od_graph_def.ParseFromString(serialized_graph) + tf.import_graph_def(od_graph_def, name="") + +label_map = label_map_util.load_labelmap(PATH_TO_LABELS) +categories = label_map_util.convert_label_map_to_categories( + label_map, max_num_classes=NUM_CLASSES, use_display_name=True +) +category_index = label_map_util.create_category_index(categories) + +with detection_graph.as_default(): + with tf.compat.v1.Session(graph=detection_graph) as sess: + while True: + ret, image_np = cap.read() + # Expand dimensions since the model expects images to have shape: [1, None, None, 3] + image_np_expanded = np.expand_dims(image_np, axis=0) + image_tensor = detection_graph.get_tensor_by_name("image_tensor:0") + # Each box represents a part of the image where a particular object was detected. + boxes = detection_graph.get_tensor_by_name("detection_boxes:0") + # Each score represent how level of confidence for each of the objects. + # Score is shown on the result image, together with the class label. + scores = detection_graph.get_tensor_by_name("detection_scores:0") + classes = detection_graph.get_tensor_by_name("detection_classes:0") + num_detections = detection_graph.get_tensor_by_name("num_detections:0") + # Actual detection. + (boxes, scores, classes, num_detections) = sess.run( + [boxes, scores, classes, num_detections], + feed_dict={image_tensor: image_np_expanded}, + ) + # Visualization of the results of a detection. + vis_util.visualize_boxes_and_labels_on_image_array( + image_np, + np.squeeze(boxes), + np.squeeze(classes).astype(np.int32), + np.squeeze(scores), + category_index, + use_normalized_coordinates=True, + line_thickness=8, + ) + + cv2.imshow("object detection", cv2.resize(image_np, (800, 600))) + if cv2.waitKey(25) & 0xFF == ord("q"): + cv2.destroyAllWindows() + break