class ObjectDetector(AlgoStep):
def __init__(self, score=0.25, dataset=DataSet.tiny):
self.logger = get_logger(self)
if dataset == DataSet.normal:
self.yolo = YOLO(
score=score
)
elif dataset == DataSet.tiny:
self.yolo = YOLO(
score=score,
model_path='model_data/yolov3-tiny.h5',
anchors_path='model_data/yolov3-tiny_anchors.txt'
)
elif dataset == DataSet.crowdhuman:
self.yolo = YOLO(
score=score,
model_path='model_data/yolov3-tiny-crowdhuman-80000.h5',
anchors_path='model_data/yolov3-tiny-crowdhuman_anchors.txt',
classes_path='model_data/crowdhuman.names'
)
else:
raise ValueError('Dataset {} unknown'.format(dataset))
def process(self, container: AlgoContainer):
image = Image.fromarray(container.frame.image[..., ::-1]) # bgr to rgb
detections = self.yolo.detect_image(image)
if not detections:
detections = []
# TODO Transform detections into ValueObject
return container.extend_with(detections=detections)
class YoloPreditor:
def __init__(self, model_path=None, classes_path=None, classes=None):
model_path = model_path or ''
self.model = YOLO(model_path=model_path, classes_path=classes_path)
def predict(self, img):
objects = self.model.detect_image(img, with_prob_and_class=True)
return objects
def get(self):
self.set_header('Access-Control-Allow-Origin', '*')
self.set_header('Access-Control-Allow-Methods', 'POST, GET, OPTIONS')
modelId = self.get_query_argument('modelId', 'None')
if int(config["server1"]["modelId"]) == int(modelId):
_,B = model.detect_image("/home/nlp/lmy/pdf8-web/01.jpg")
else:
model = YOLO()
_,B = model.detect_image("/home/nlp/lmy/pdf8-web/01.jpg")
respon = {"obj":str(B)}
self.set_header('Content-Type', 'application/json; charset=UTF-8')
self.write(json.dumps(respon))
self.finish()
def test(filename,
model_path,
anchors_path,
classes_path,
*,
score_threshold=0.2,
iou=0.45,
max_boxes=100):
yolo = YOLO(
model_path=model_path,
anchors_path=anchors_path,
classes_path=classes_path,
score_threshold=score_threshold,
iou=iou,
max_boxes=max_boxes,
)
x = imread(filename)
x = Image.fromarray(x)
x = yolo.detect_image(x)
x = np.array(x)
imsave('out.png', x)
class TLClassifier(object):
def __init__(self, is_site, configuration):
#TODO load classifier
self.is_site = is_site
self.site_use_yolo = False
if is_site:
if self.site_use_yolo:
self.model = YOLO(model_path=configuration["model_path"], anchors_path=configuration["model_anchor"],
classes_path=configuration["model_classes"])
self.traffic_light_classes = 9 # TODO: hard code
else:
self.model = Model(model_path=configuration["model_path"])
else:
self.model = YOLO(model_path=configuration["model_path"], anchors_path=configuration["model_anchor"],
classes_path=configuration["model_classes"])
self.traffic_light_classes = 9 # TODO: hard code
def get_classification(self, image):
"""Determines the color of the traffic light in the image
Args:
image (cv::Mat): image containing the traffic light
Returns:
int: ID of traffic l ight color (specified in styx_msgs/TrafficLight)
"""
#TODO implement light color prediction
if image is not None:
image_data = np.asarray(image)
if self.is_site:
if self.site_use_yolo:
pil_image = Image.fromarray(image_data)
out_boxes, out_scores, out_classes = self.model.detect_image(pil_image)
detected_states = []
for i, c in reversed(list(enumerate(out_classes))):
# predicted_class = self.class_names[c]
box = out_boxes[i]
score = out_scores[i]
if c == self.traffic_light_classes:
top, left, bottom, right = box
top = max(0, np.floor(top + 0.5).astype('int32'))
left = max(0, np.floor(left + 0.5).astype('int32'))
bottom = min(pil_image.size[1], np.floor(bottom + 0.5).astype('int32'))
right = min(pil_image.size[0], np.floor(right + 0.5).astype('int32'))
corp_traffic_light = image_data[top:bottom, left:right, :]
detected_states.append(self.detect_traffic_light_state_luminous(corp_traffic_light))
if len(detected_states) > 0:
(_, idx, counts) = np.unique(detected_states, return_index=True, return_counts=True)
index = idx[np.argmax(counts)]
return detected_states[index]
else:
return TrafficLight.UNKNOWN
else:
return self.model.detect_traffic_light(image)
else:
pil_image = Image.fromarray(image_data)
out_boxes, out_scores, out_classes = self.model.detect_image(pil_image)
detected_states = []
for i, c in reversed(list(enumerate(out_classes))):
#predicted_class = self.class_names[c]
box = out_boxes[i]
score = out_scores[i]
if c == self.traffic_light_classes:
top, left, bottom, right = box
top = max(0, np.floor(top + 0.5).astype('int32'))
left = max(0, np.floor(left + 0.5).astype('int32'))
bottom = min(pil_image.size[1], np.floor(bottom + 0.5).astype('int32'))
right = min(pil_image.size[0], np.floor(right + 0.5).astype('int32'))
corp_traffic_light = image_data[top:bottom, left:right, :]
detected_states.append(self.detect_traffic_light_state_color(corp_traffic_light))
if len(detected_states) > 0:
(_, idx, counts) = np.unique(detected_states, return_index=True, return_counts=True)
index = idx[np.argmax(counts)]
return detected_states[index]
else:
return TrafficLight.UNKNOWN
return TrafficLight.UNKNOWN
def detect_traffic_light_state_color(self, image):
def simple_thresh(img, thresh=(0, 255)):
# apply only thresholding on image plane and return binary image
binary = np.zeros_like(img)
binary[(img > thresh[0]) & (img <= thresh[1])] = 1
return binary
rows = image.shape[0]
# thresholding on red channel
red_channel = image[:, :, 0]
# thresholding on green channel
green_channel = image[:, :, 1]
blue_channel = image[:, :, 2]
# patitioning on r_binary
r_binary = simple_thresh(red_channel, (196, 255))
# plt.figure()
# plt.imshow(r_binary)
r_binary_1 = r_binary[0:int(rows / 3) - 1, :]
r_binary_2 = r_binary[int(rows / 3):(2 * int(rows / 3)) - 1, :]
r_binary_3 = r_binary[(2 * int(rows / 3)):rows - 1, :]
is_red = (r_binary_1.sum() > 2 * (r_binary_2.sum())) & (r_binary_1.sum() > 2 * (r_binary_3.sum()))
# (GREEN) patitioning on g_binary
g_binary = simple_thresh(green_channel, (183, 255))
# plt.figure()
# plt.imshow(g_binary)
g_binary_1 = g_binary[0:int(rows / 3) - 1, :]
g_binary_2 = g_binary[int(rows / 3):(2 * int(rows / 3)) - 1, :]
g_binary_3 = g_binary[(2 * int(rows / 3)):rows - 1, :]
is_green = (g_binary_3.sum() > 2 * (g_binary_2.sum())) & (g_binary_3.sum() > 2 * (g_binary_1.sum()))
# (YELLOW) partition on g_binary and b_binary
g_binary = simple_thresh(green_channel, (146, 255))
b_binary = simple_thresh(blue_channel, (143, 255))
combine = g_binary + b_binary
y_binary_1 = combine[0:int(rows / 3) - 1, :]
y_binary_2 = combine[int(rows / 3):(2 * int(rows / 3)) - 1, :]
y_binary_3 = combine[(2 * int(rows / 3)):rows - 1, :]
is_yellow = (y_binary_2.sum() > 2 * (y_binary_1.sum())) & (y_binary_2.sum() > 2 * (y_binary_3.sum()))
if is_red:
return TrafficLight.RED
elif is_yellow:
return TrafficLight.YELLOW
elif is_green:
return TrafficLight.GREEN
else:
return TrafficLight.UNKNOWN
def detect_traffic_light_state_luminous(self, image):
def adjust_histo(image):
img_yuv = cv2.cvtColor(image, cv2.COLOR_BGR2YUV)
if np.max(img_yuv[:, :, 0]) - np.min(img_yuv[:, :, 0]) < 180:
img_yuv[:, :, 0] = cv2.equalizeHist(img_yuv[:, :, 0])
img_output = cv2.cvtColor(img_yuv, cv2.COLOR_YUV2BGR)
return img_output
image = adjust_histo(image)
img_h, img_w, img_ch = image.shape
l_channel = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)[:, :, 2]
top_third_marker = int(img_h / 3)
bottom_third_marker = img_h - top_third_marker
# Magnitude of L is established for each section of the image
top = 0
mid = 0
bottom = 0
count_result = {'RED': 0, 'YELLOW': 0, 'GREEN': 0}
for i in range(top_third_marker):
for j in range(img_w):
top += l_channel[i][j]
count_result['RED'] = top # compensate for R, G, B => L, U, V
for i in range(top_third_marker, bottom_third_marker):
for j in range(img_w):
mid += l_channel[i][j]
count_result['YELLOW'] = mid
for i in range(bottom_third_marker, img_h):
for j in range(img_w):
bottom += l_channel[i][j]
count_result['GREEN'] = bottom
# The result is classified into one of the 3 colors and returned
max_count = max(count_result, key=count_result.get)
std = np.std(list(count_result.values()))
mean = np.mean(list(count_result.values()))
# print('x', std, std/mean)
if (std / mean < 0.1):
return TrafficLight.UNKNOWN
if max_count == 'RED':
return TrafficLight.RED
elif max_count == 'YELLOW':
return TrafficLight.YELLOW
elif max_count == 'GREEN':
return TrafficLight.GREEN
else:
return TrafficLight.UNKNOWN
detector = YOLO()
results = {}
name_list = []
for file in os.listdir(input_path):
if file.endswith(".jpg"):
name_list.append(file)
number_objects = 0
for name in name_list: #test
img_path = os.path.join(input_path, name)
result = {}
image = Image.open(img_path)
image_ = cv2.imread(img_path)
height, width, depth = image_.shape
_, res = detector.detect_image(image)
out_boxes, out_scores, out_classes = res
result['height'] = int(height)
result['width'] = int(width)
result['depth'] = int(depth)
objects = {}
for c, bbox in zip(out_classes, out_boxes):
number_objects += 1
objects[str(number_objects)] = {
'category': labels[c],
'bbox': [int(b) for b in bbox]
}
result['objects'] = objects
results[name] = result
with open(save_json_path, 'w', encoding='utf-8') as f:
json.dump(results, f)
class DroneStabilization:
image_sub = None
def __init__(self):
rospy.Subscriber("/state_machine/state", String, self.state_callback)
self.pub_condition = rospy.Publisher(
"/state_machine/window/find_condition", String, queue_size=10)
rospy.init_node('window', anonymous=True)
self.yolo = YOLO()
self.count_stable = 0
def setup_window(self):
self.image_pub = rospy.Publisher("image_topic_2", Image, queue_size=10)
self.vel_pub = rospy.Publisher('controle/velocity',
Vector3D,
queue_size=10)
self.bridge = CvBridge()
# self.stab = camStab()
self.detect = detect_window()
self.vec_vel = Vector3D()
self.vec_vel.x = 0
self.vec_vel.y = 0
self.vec_vel.z = 0
self.pixel_movement = [0, 0]
self.frame = 0
self.count_callbacks = 0
self.opt_flow = camStab()
self.image_sub = rospy.Subscriber("/bebop/image_raw",
Image,
self.callback,
queue_size=None)
rospy.loginfo("setup ok")
# self.pub_condition.publish("ok")
def state_callback(self, data):
if (data.data == "find_window"):
rospy.loginfo(" WINDOW !!!")
self.setup_window()
# condition!!
elif (self.image_sub != None):
rospy.loginfo(" end! ")
cv2.destroyAllWindows()
self.image_sub.unregister()
self.image_sub = None
def callback(self, data):
# if self.count_callbacks < 10:
# self.count_callbacks += 1
# return
try:
cv_image = self.bridge.imgmsg_to_cv2(data, "bgr8")
except CvBridgeError as e:
print(e)
(rows, cols, channels) = cv_image.shape
if not (cols > 60 and rows > 60): # returns if data have unvalid shape
return
# self.detect.update(cv_image)
if self.frame == 0:
pil_image = PIL.Image.fromarray(cv_image)
r_image, out_boxes, out_score = self.yolo.detect_image(pil_image)
# window detected with min precision of (defined on yolo file)
if out_score.shape[0] > 0:
rospy.loginfo("detect")
m_box = out_boxes.mean(0)
self.mean_box_cord = m_box[:]
self.image_h = r_image.shape[0]
self.image_w = r_image.shape[1]
self.image_center = np.array(
[self.image_w / 2, self.image_h / 2])
box_margin = 30
max_score_index = np.argmax(out_score)
box_lt = (max(
int(out_boxes[max_score_index][1]) - box_margin,
0), max(
int(out_boxes[max_score_index][0]) - box_margin, 0))
box_rb = (min(
int(out_boxes[max_score_index][3]) + box_margin,
self.image_w),
min(
int(out_boxes[max_score_index][2]) + box_margin,
self.image_h))
img = cv2.rectangle(r_image, box_lt, box_rb, (0, 220, 200), 2)
self.opt_flow.resetFeatures(
cv_image,
interval=[box_lt[1], box_rb[1], box_lt[0], box_rb[0]],
get_corners_only=True)
# print(init_feature + np.array(box_lt))
# self.opt_flow.set_p0(init_feature + np.array(box_lt))
# self.opt_flow.set_initial_img(cv_image)
self.frame += 1
cv2.imshow("window", img)
else:
rospy.loginfo("no windows on view, trying again")
# self.vel_pub.publish(self.vec_vel)
cv2.imshow("window", r_image)
else:
self.opt_flow.update(cv_image)
print(self.opt_flow.get_square_shape())
self.adjust_position(self.opt_flow.get_square_center(),
self.opt_flow.get_square_shape())
k = cv2.waitKey(30) & 0xff
if k == 27:
exit()
elif k == 'r' or k == ord('r'):
print("r pressed ")
self.frame = 0
elif k == ord('p') or k == 'p':
print("through_window ")
self.pub_condition.publish("ok")
try:
# self.vel_pub.publish(self.vec_vel)
# print(self.vec_vel.x)
# self.image_pub.publish(self.bridge.cv2_to_imgmsg(cv_image, "bgr8"))
pass
except CvBridgeError as e:
print(e)
self.count_callbacks = 0
def adjust_position(self,
window_center_pos,
window_shape,
ideal_window_image_rate=0.60):
delta = self.image_center - window_center_pos
# rospy.loginfo(delta)
self.vec_vel.x = 0
self.vec_vel.y = 0
self.vec_vel.z = 0
if abs(delta[0]) > 5 or abs(delta[1]) > 5:
self.vec_vel.y = np.sign(delta[0]) * min(
np.abs(delta[0]) / 20, 0.2)
self.vec_vel.z = np.sign(delta[1]) * min(
np.abs(delta[1]) / 20, 0.2)
self.count_stable = 0
rate = window_shape[1] / self.image_h
delta_rate = ideal_window_image_rate - rate
# (x,y) goes back
if abs(delta_rate) > 0.02 and abs(delta[0]) < 20 and abs(
delta[1]) < 20:
self.vec_vel.x = delta_rate * 2
self.count_stable = 0
self.vel_pub.publish(self.vec_vel)
self.count_stable += 1
if self.count_stable > 3:
rospy.loginfo("GOOD!!")
self.pub_condition.publish("ok")