def net_segment_image_batch(image, segment_engine, latent_engine, scale = 1,occlude = False, mask_image = False, compute_latent=False):
start = time.time()
size = (int(image.shape[1]*scale), int(image.shape[0]*scale))
if occlude:
image = oclude_gripper(image)
resized_image = cv2.resize(image, size)
resize_mask, resized_prob = segment_engine.segment(resized_image, (size[1],size[0]))
mask = cv2.resize(resize_mask, (image.shape[1], image.shape[0]))
ret, seperate_masks = cv2.connectedComponents(mask)
latents = []
coords = []
small_ims = []
boxes = []
for ii in range(1, seperate_masks.max().max() + 1):
bool_mask = (seperate_masks == ii)
curr_mask = mask * bool_mask
x, y = center(curr_mask, size)
curr_mask = np.stack((curr_mask,) * 3, 2)
if mask_image:
bool_mask = curr_mask > curr_mask.mean()
masked_im = image * bool_mask
else:
masked_im = image
slice_x, slice_y = get_bounding_box(curr_mask)
if slice_x is None:
continue
small_im = masked_im[slice_x, slice_y]
small_im = cv2.resize(small_im, (64, 64))
small_ims.append(small_im)
boxes.append((slice_y.start, slice_x.start, slice_y.stop-slice_y.start, slice_x.stop-slice_x.start))
coords.append((x,y))
if compute_latent:
latents_temp = latent_engine.get_latent_batch(small_ims)
for jj in range(latents_temp.shape[0]):
a = np.expand_dims(latents_temp[jj], 0)
if a is np.ndarray:
latents.append(utils.encode_numpy_array(a))
end = time.time()
elapsed = end - start
return latents, coords, utils.encode_numpy_array(mask), boxes
def detect_brick(engine, frame):
start = time.time()
size = (240, 320)
mask, prob = segment_image(engine, frame, size=size)
logger.debug("Segmentation %d", time.time() - start)
x, y = center(mask, size)
if x and y:
is_tracked = True
x = utils.scale_to_range(x, 0, size[1], 0, frame.shape[1])
y = utils.scale_to_range(y, 0, size[0], 0, frame.shape[0])
else:
is_tracked = False
return is_tracked, x, y, mask
def image_to_latents_batch(image, segment_engine, latent_engine, scale = 1,occlude = False, mask_image = False):
start = time.time()
size = (int(image.shape[1]*scale), int(image.shape[0]*scale))
if occlude:
image = oclude_gripper(image)
resized_image = cv2.resize(image, size)
resize_mask, resized_prob = segment_engine.segment(resized_image, (size[1],size[0]))
mask = cv2.resize(resize_mask, (image.shape[1], image.shape[0]))
ret, seperate_masks = cv2.connectedComponents(mask)
latents = []
coords = []
small_ims = []
for ii in range(1, seperate_masks.max().max() + 1):
bool_mask = (seperate_masks == ii)
curr_mask = mask * bool_mask
x, y = center(curr_mask, size)
curr_mask = np.stack((curr_mask,) * 3, 2)
if mask_image:
bool_mask = curr_mask > curr_mask.mean()
masked_im = image * bool_mask
else:
masked_im = image
slice_x, slice_y = get_bounding_box(curr_mask)
if slice_x is None:
continue
small_im = masked_im[slice_x, slice_y]
small_im = cv2.resize(small_im, (64, 64))
# latents.append(latent_engine.get_latent(small_im))
coords.append((x,y))
small_ims.append(small_im)
latents = latent_engine.get_latent_batch(small_ims)
# latents = np.zeros([seperate_masks.max().max() + 1,32])
end = time.time()
elapsed = end - start
return latents,coords,small_ims,mask,elapsed
def search_and_reach_brick(robot,
engine,
config,
visualize=False,
record=False):
# def reset():
# search_and_reach_brick.track_count = 0
# search_and_reach_brick.prev_x = None
# search_and_reach_brick.prev_y = None
try:
if (search_and_reach_brick.rec_count + 1 > 200):
return False, False
except:
search_and_reach_brick.rec_count = 0
for _ in range(15):
frame = robot.cam["ps3_one"].read()
is_tracked, x, y, mask = detect_brick(engine, frame)
# is_tracked = False
# Only if we got a tracked point, we continue. Else: Stop any
# movement
if is_tracked:
# Improve tracking by considering the momentum of the tracked point
try:
search_and_reach_brick.nothing_in_view = 0
except AttributeError:
search_and_reach_brick.nothing_in_view = 0
try:
x = momentum_scaling(search_and_reach_brick.prev_x, x)
search_and_reach_brick.prev_x = x
except AttributeError:
search_and_reach_brick.prev_x = x
try:
y = momentum_scaling(search_and_reach_brick.prev_y, y)
search_and_reach_brick.prev_y = y
except AttributeError:
search_and_reach_brick.prev_x = x
logger.debug("x=%d, y=%d", x, y)
# Calculate error between current location and desired target
# location
# /(config.XY_TRACK_POINT[0])
errx = (x - config.XY_TRACK_POINT[0])
# /(config.XY_TRACK_POINT[1])
erry = (y - config.XY_TRACK_POINT[1])
logger.debug("%d, %d", errx, erry)
# Reached Center Zone, drop out of this behavior part
try:
if search_and_reach_brick.track_count > 5:
return True, True
except AttributeError:
search_and_reach_brick.track_count = 0
if abs(errx) < 50 and abs(erry) < 30:
search_and_reach_brick.track_count = search_and_reach_brick.track_count + 1
logger.debug("New track count: %d",
search_and_reach_brick.track_count)
errx = utils.scale_to_range(errx, -config.XY_TRACK_POINT[1] / 2.0,
config.XY_TRACK_POINT[1] / 2.0, -100, 100)
erry = utils.scale_to_range(erry, -config.XY_TRACK_POINT[0] / 2.0,
config.XY_TRACK_POINT[0] / 2.0, -100, 100)
# Update the PID controller
robot.x_PID.update(erry)
robot.w_PID.update(errx)
# Generate Motor Signals from the differential drive controller
logger.debug("Control Outputs: vl=%d vr=%d", robot.x_PID.output,
robot.w_PID.output)
v_l, v_r = robot.controller.twist_to_vel(robot.x_PID.output,
robot.w_PID.output)
# Send signal to the robot/slave
robot.controller.move(v_l, v_r)
else:
robot.controller.move(0, 0)
try:
search_and_reach_brick.nothing_in_view += 1
except:
search_and_reach_brick.nothing_in_view = 0
if search_and_reach_brick.nothing_in_view > 5:
robot.controller.move(90, -80)
# try:
# search_and_reach_brick.nothing_in_view += 1
# logger.debug("Nothing seen %d", search_and_reach_brick.nothing_in_view)
# if search_and_reach_brick.nothing_in_view>5:
# # robot.cam["ps3_one"].close()
# robot.cam["ps3_two"].open()
# frame = robot.cam["ps3_two"].read()
# is_tracked, x, y, mask = detect_brick(engine, frame)
# if is_tracked:
# errx = (x - config.XY_TRACK_POINT[0])
# erry = (y - config.XY_TRACK_POINT[1])
# logger.debug("%d, %d", errx, erry)
# errx = utils.scale_to_range(errx, -config.XY_TRACK_POINT[1] / 2.0,
# config.XY_TRACK_POINT[1] / 2.0, -100, 100)
# erry = utils.scale_to_range(erry, -config.XY_TRACK_POINT[0] / 2.0,
# config.XY_TRACK_POINT[0] / 2.0, -100, 100)
# robot.x_PID.update(erry)
# robot.w_PID.update(errx)
# # Generate Motor Signals from the differential drive controller
# v_l, v_r = robot.controller.twist_to_vel(
# robot.x_PID.output, robot.w_PID.output)
# logger.debug("Control Outputs: vl=%d vr=%d",
# robot.x_PID.output, robot.w_PID.output)
# # Send signal to the robot/slave
# robot.controller.move(v_l, v_r)
# else:
# # Halt robot
# robot.controller.move(0, 0)
# robot.cam["ps3_two"].close()
# # robot.cam["ps3_one"].open()
# if search_and_reach_brick.nothing_in_view > 25:
# robot.controller.move(-45, 45)
# except AttributeError:
# search_and_reach_brick.nothing_in_view = 0
# Visualize the tracking and target
if visualize:
if is_tracked:
# cv2.circle(mask, (xold, yold), 10, (0, 0, 255), 1)
# visualize_point_tracking(frame, (int(x), int(y)), config.XY_TRACK_POINT, "tracked_point")
cv2.imshow('mask', mask)
else:
cv2.imshow('mask', mask)
cv2.imshow('tracked_point', frame)
cv2.waitKey(1)
if record:
size = (240, 320)
image = cv2.resize(frame, (320, 240))
x, y = center(mask, size)
contours = get_objects_seg(mask)
if not x:
print('no brick')
# cv2.imshow('frame', image)
# if cv2.waitKey(1) & 0xFF == ord('q'):
# break
return True, True
ret, seperate_masks = cv2.connectedComponents(mask)
print(seperate_masks.max().max())
orig_mask = mask
for ii in range(1, len(contours) + 1):
bool_mask = (seperate_masks == ii)
mask = orig_mask * bool_mask
mask = np.stack((mask, ) * 3, 2)
bool_mask = mask > mask.mean()
masked_im = image * bool_mask
slice_x, slice_y = get_bounding_box(mask)
if slice_x is None:
continue
small_im = masked_im[slice_x, slice_y]
small_im = cv2.resize(small_im, (64, 64))
# cv2.imshow('frame', image)
# cv2.imshow('mask' + str(ii), small_im)
# cv2.imshow('hm', small_im)
# latent = get_latent(small_im, sess, loc)
# print(latent)
# try:
search_and_reach_brick.rec_count += 1
# except AttributeError:
# search_and_reach_brick.rec_count = 0
name = str(1).zfill(5) + '_' + str(
search_and_reach_brick.rec_count).zfill(4)
cv2.imwrite(os.path.join(record, name + '.png'), small_im)
# try:
# search_and_reach_brick.rec_count += 1
# # cv2.imwrite('/home/dlrc/mask2.jpg', mask)
# except AttributeError:
# search_and_reach_brick.rec_count = 0
#
# cv2.imwrite(os.path.join(record, 'frame' + str(search_and_reach_brick.rec_count) + '.jpg'), frame)
# Brick not reached yet, continue loop
return True, True
def image_to_latents_small(image, segment_engine, latent_engine, scale = 1,occlude = False, mask_image = False):
# occlude = True
start = time.time()
size = (int(image.shape[1]*scale), int(image.shape[0]*scale))
if occlude:
image = oclude_gripper(image)
resized_image = cv2.resize(image, size)
resize_mask, resized_prob = segment_engine.segment(resized_image, (size[1],size[0]))
mask = cv2.resize(resize_mask, (image.shape[1], image.shape[0]))
# cv2.imshow("1", mask)
#
kernel = np.ones((19, 19), np.uint8)
mask = cv2.morphologyEx(mask,cv2.MORPH_OPEN,kernel)
#
# cv2.imshow("2", mask)
new_mask = np.zeros((image.shape[0],image.shape[1]))
ret, seperate_masks = cv2.connectedComponents(mask)
for ii in range(1, seperate_masks.max().max() + 1):
bool_mask = (seperate_masks == ii)
curr_mask = mask * bool_mask
nnz = np.count_nonzero(curr_mask)
if(nnz<200):
continue
new_mask += curr_mask
mask = new_mask.astype("uint8")
mask[mask<200] = 0
mask[mask>=200] = 255
# cv2.imshow("3",mask)
#
# # kernel = np.ones((19, 19), np.uint8)
# # mask = cv2.morphologyEx(mask,cv2.MORPH_CLOSE,kernel)
#
# cv2.imshow("4", mask)
# cv2.waitKey(0)
# kernel = np.ones((5, 5), np.uint8)
# mask = cv2.morphologyEx(mask,cv2.MORPH_CLOSE,kernel)
ret, seperate_masks = cv2.connectedComponents(mask)
latents = []
coords = []
small_ims = []
for ii in range(1, seperate_masks.max().max() + 1):
bool_mask = (seperate_masks == ii)
curr_mask = mask * bool_mask
x, y = center(curr_mask, size)
curr_mask = np.stack((curr_mask,) * 3, 2)
if mask_image:
bool_mask = curr_mask > curr_mask.mean()
masked_im = image * bool_mask
else:
masked_im = image
slice_x, slice_y = get_bounding_box(curr_mask)
if slice_x is None:
continue
# New code: not yet tested - not sure if works - don't use
# slice_x = slice(int(min(slice_x.start -5,0)),int(max(slice_x.stop + 5,image.shape[0])),None)
# slice_y = slice(int(min(slice_y.start - 5, 0)), int(max(slice_y.stop + 5, image.shape[0])), None)
# slice_x.stop = max(slice_x.stop + 5,image.shape[0])
# slice_y.start = min(slice_y.start - 5, 0)
# slice_y.stop = max(slice_y.stop + 5, image.shape[1])
small_im = masked_im[slice_x, slice_y]
small_im = cv2.resize(small_im, (64, 64))
latents.append(latent_engine.get_latent(small_im))
coords.append((x,y))
small_ims.append(small_im)
end = time.time()
elapsed = end - start
return latents,coords,small_ims,mask,elapsed
def net_segment_image(image, segment_engine, latent_engine, scale = 1,occlude = False, mask_image = False, compute_latent=False, final_latent = False):
start = time.time()
size = (int(image.shape[1]*scale), int(image.shape[0]*scale))
latents = []
coords = []
boxes = []
small_ims = []
if final_latent:
small_im = image
small_im = cv2.resize(small_im, (64, 64))
small_ims.append(small_im)
latents.append(utils.encode_numpy_array(latent_engine.get_latent(small_im)))
mask = np.empty([size[0],size[1],3])
import os
fol = "/home/dlrc/aaaa/"
num_fol = len(os.listdir(fol))-1
if not os.path.exists(fol + str(num_fol).zfill(4)):
os.makedirs(fol + str(num_fol).zfill(4))
num_im = len(os.listdir(fol + str(num_fol).zfill(4)))
cv2.imwrite(fol + str(num_fol).zfill(4) + "/" + str(num_im).zfill(4) + "_test.jpg", small_im)
else:
# if occlude:
# image = oclude_gripper(image)
resized_image = cv2.resize(image, size)
resize_mask, resized_prob = segment_engine.segment(resized_image, (size[1],size[0]))
mask = cv2.resize(resize_mask, (image.shape[1], image.shape[0]))
if occlude:
mask = oclude_gripper(mask)
ret, seperate_masks = cv2.connectedComponents(mask)
for ii in range(1, seperate_masks.max().max() + 1):
bool_mask = (seperate_masks == ii)
curr_mask = mask * bool_mask
x, y = center(curr_mask, size)
curr_mask = np.stack((curr_mask,) * 3, 2)
if mask_image:
bool_mask = curr_mask > curr_mask.mean()
masked_im = image * bool_mask
else:
masked_im = image
slice_x, slice_y = get_bounding_box(curr_mask)
if slice_x is None:
continue
small_im = masked_im[slice_x, slice_y]
small_im = cv2.resize(small_im, (64, 64))
boxes.append((slice_y.start, slice_x.start, slice_y.stop-slice_y.start, slice_x.stop-slice_x.start))
if compute_latent:
latents.append(utils.encode_numpy_array(latent_engine.get_latent(small_im)))
small_ims.append(small_im)
coords.append((x,y))
print(compute_latent)
if compute_latent:
import os
fol = "/home/dlrc/aaaa/"
num_fol = len(os.listdir(fol))
if not os.path.exists(fol + str(num_fol).zfill(4)):
os.makedirs(fol + str(num_fol).zfill(4))
cv2.imwrite(fol + str(num_fol).zfill(4) + "/orig.jpg", image)
cv2.imwrite(fol + str(num_fol).zfill(4) + "/mask.jpg", mask)
print(fol + str(num_fol).zfill(4) + "/orig.jpg")
print(fol + str(num_fol).zfill(4) + "/mask.jpg")
for i, small_im in enumerate(small_ims):
cv2.imwrite(fol + str(num_fol).zfill(4) + "/" + str(i).zfill(4) + "_small.jpg", small_im)
print(fol + str(num_fol).zfill(4) + "/" + str(i).zfill(4) + "_small.jpg")
end = time.time()
elapsed = end - start
return latents, coords, utils.encode_numpy_array(mask), boxes, small_ims