-
Notifications
You must be signed in to change notification settings - Fork 4
/
Copy pathobj_det_MobileNetSSD.py
78 lines (59 loc) · 2.78 KB
/
obj_det_MobileNetSSD.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
import time
import rospy
import importlib
baxter=importlib.import_module("baxter-python3.baxter")
import cv2
import numpy as np
PI = 3.141592
WIDTH = 640
HEIGHT = 400
print("[INFO] loading model...")
PROTO = "./models/MobileNetSSD_deploy.prototxt"
MODEL = "./models/MobileNetSSD_deploy.caffemodel"
net = cv2.dnn.readNetFromCaffe(PROTO, MODEL)
RESIZED_DIMENSIONS = (300, 300) # Dimensions net was trained on.
IMG_NORM_RATIO = 0.007843 # In grayscale a pixel can range between 0 and 255
#pascal voc classes
classes = ["background", "aeroplane", "bicycle", "bird", "boat", "bottle",
"bus", "car", "cat", "chair", "cow",
"diningtable", "dog", "horse", "motorbike", "person",
"pottedplant", "sheep", "sofa", "train", "tvmonitor"]
print("[INFO] starting robot...")
rospy.init_node("testing")
rospy.sleep(2.0)
robot = baxter.BaxterRobot(rate=100, arm="left")
rospy.sleep(2.0)
robot._set_camera(camera_name="left_hand_camera", state=True, width=WIDTH, height=HEIGHT, fps=30)
robot.set_robot_state(True)
print("[INFO] moving in position...")
print(robot.move_to_neutral())
print(robot.move_to_zero())
print(robot.move_to_joint_position({"left_s0": -PI/4}, timeout=10))
data = np.array(list(robot._cam_image.data), dtype=np.uint8)
middle_point = np.array([WIDTH/2, HEIGHT/2])
print("[INFO] getting image stream and passing to DNN...")
while not rospy.is_shutdown():
img = np.array(list(robot._cam_image.data), dtype=np.uint8)
img = img.reshape(int(HEIGHT), int(WIDTH), 4)
img = img[:, :, :3].copy()
#Passing image to DNN
#gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
blob = cv2.dnn.blobFromImage(cv2.resize(img, RESIZED_DIMENSIONS), IMG_NORM_RATIO, RESIZED_DIMENSIONS, 127.5)
net.setInput(blob)
neural_network_output = net.forward()
print("Detections: " + str(neural_network_output.shape[2])) if len(neural_network_output)!=0 else print("No Detections")
for i in np.arange(0, neural_network_output.shape[2]):
confidence = neural_network_output[0, 0, i, 2]
# Confidence must be at least x%
if confidence > 0.40:
idx = int(neural_network_output[0, 0, i, 1])
bounding_box = neural_network_output[0, 0, i, 3:7] * np.array([WIDTH, HEIGHT, WIDTH, HEIGHT])
(startX, startY, endX, endY) = bounding_box.astype("int")
label = "{}: {:.2f}%".format(classes[idx], confidence * 100)
cv2.rectangle(img, (startX, startY), (endX, endY), (255,0,0), 2)
y = startY - 15 if startY - 15 > 15 else startY + 15
cv2.putText(img, label, (startX, y),cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255,0,0), 2)
robot._set_display_data(cv2.resize(img, (1024,600)))
robot.rate.sleep()
print(robot.move_to_neutral())
robot.set_robot_state(False)