def __getitem__(self, idx):
return_dict = {}
scene_num, center, label = self.pair_dict[idx]
cloud_0, cloud_1 = [
x[extract_area(
x, center, self.context_voxel_size[0].item(
), shape='square'), :] for x in self.loaded_clouds[scene_num]
]
z_max = max(cloud_0[:, 2].max(), cloud_1[:, 2].max())
z_min = min(cloud_0[:, 2].min(), cloud_1[:, 2].min())
z_voxel_centers = self.voxel_center_heights(z_min, z_max)
return_dict['voxels'] = {}
for index, z_voxel_center in enumerate(z_voxel_centers):
vox_center = torch.cat(
(center, z_voxel_center.unsqueeze(0).to(center.device)),
dim=-1)
context_for_1, voxel_1, context_0_0 = self.get_voxels(
cloud_1, cloud_0, vox_center)
context_for_0, voxel_0, context_1_1 = self.get_voxels(
cloud_0, cloud_1, vox_center)
return_dict['voxels'][index] = [
context_for_1, voxel_1, context_0_0, context_for_0, voxel_0,
context_1_1, z_voxel_center
]
return_dict['cloud_0'] = cloud_0
return_dict['cloud_1'] = cloud_1
return return_dict, label
def show_processing(self, frame, name="Press"):
"""
Displays the inner area (bounding box of the tracking object).
Parameters
----------
frame : ndarray
3-channel image.
name : str
Name of the window to display the frame.
"""
main_area = extract_area(frame, self.x_start, self.y_start,
self.x_end - self.x_start,
self.y_end - self.y_start)
x, y = int(self.bbox[0]), int(self.bbox[1])
w, h = int(self.bbox[2]), int(self.bbox[3])
H = main_area.shape[0]
W = main_area.shape[1]
p1 = (x, y)
p2 = (x + w, y + h)
cv.rectangle(main_area, p1, p2, (255, 0, 0), 1, 1)
# Resize for better visualization
main_area = cv.resize(main_area, (W * 3, H * 3))
cv.imshow(name, main_area)
def process_frame(self, frame):
"""
Processes a frame to add its average lightness to history.
Parameters
----------
frame : ndarray
3-channel image.
"""
if self.light_history is None:
raise ValueError('Counter is not initialized')
area = extract_area(frame, self.x_start, self.y_start, self.width,
self.height)
area = cv.cvtColor(area, cv.COLOR_BGR2HLS)
light_avg = np.mean(area[:, :, 1])
self.light_history = np.hstack((self.light_history, light_avg))
def init(self, frame):
"""
Initializes the average lightness array.
Parameters
----------
frame : ndarray
3-channel image
"""
# Extract the shoot area.
area = extract_area(frame, self.x_start, self.y_start, self.width,
self.height)
# Change to color space to get the lightness.
area = cv.cvtColor(area, cv.COLOR_BGR2HLS)
# Get the average lightness and start the history array.
light_avg = np.mean(area[:, :, 1])
self.light_history = np.array(light_avg)
def init(self, frame):
"""
Initializes the tracker. The object to be tracked is defined by the
content of the inner area in the frame
Parameters
----------
frame : ndarray
3-channel image.
"""
# Extract the main area.
main_area = extract_area(frame, self.x_start, self.y_start,
self.x_end - self.x_start,
self.y_end - self.y_start)
# Initialize the tracker with the main and inner areas.
self.tracker.init(main_area, self.bbox)
# Init y value of top-left corner of inner area.
self.y_pos_history = np.array(int(self.bbox[1]))
def show_processing(self, frame, name="Plates"):
"""
Displays the lightness of the shoot.
Parameters
----------
frame : ndarray
3-channel image.
name : str
Name of the window to display the frame.
"""
shoot = extract_area(frame, self.x_start, self.y_start, self.width,
self.height)
H = shoot.shape[0]
W = shoot.shape[1]
shoot = cv.cvtColor(shoot, cv.COLOR_BGR2HLS)
# Resize for better visualization
shoot = cv.resize(shoot, (W * 3, H * 3))
cv.imshow(name, shoot[:, :, 1])
def process_frame(self, frame):
"""
Processes the frame to update the inner area, that is, finding the
position of the object defined in the initial frame.
Parameters
----------
frame : ndarray
3-channel image.
"""
if self.y_pos_history is None:
raise ValueError('Tracker is not initialized')
main_area = extract_area(frame, self.x_start, self.y_start,
self.x_end - self.x_start,
self.y_end - self.y_start)
ok, bbox = self.tracker.update(main_area)
if ok:
self.bbox = bbox
# Vertical position of the tracking object in the current frame.
self.y_pos_history = np.hstack(
(self.y_pos_history, int(self.bbox[1])))
def __init__(self, directory_path_train,directory_path_test, out_path, clearance=10, preload=False,
height_min_dif=0.5, max_height=15.0, device="cpu",n_samples=2048,final_voxel_size=[3., 3., 4.],
rotation_augment = True,n_samples_context=2048, context_voxel_size = [3., 3., 4.],
mode='train',verbose=False,voxel_size_final_downsample=0.07,include_all=False,self_pairs_train=True):
print(f'Dataset mode: {mode}')
self.mode = mode
if self.mode =='train':
directory_path = directory_path_train
elif self.mode == 'test':
directory_path = directory_path_test
else:
raise Exception('Invalid mode')
self.verbose = verbose
self.include_all = include_all
self.voxel_size_final_downsample = voxel_size_final_downsample
self.n_samples_context = n_samples_context
self.context_voxel_size = torch.tensor(context_voxel_size)
self.directory_path = directory_path
self.clearance = clearance
self.self_pairs_train = self_pairs_train
self.out_path = out_path
self.height_min_dif = height_min_dif
self.max_height = max_height
self.rotation_augment = rotation_augment
self.save_name = f'ams_{mode}_save_dict_{clearance}.pt'
name_insert = self.save_name.split('.')[0]
self.filtered_scan_path = os.path.join (
out_path, f'{name_insert}_filtered_scans.pt')
if self.mode == 'train':
self.all_valid_combs_path = os.path.join (
out_path, f'{name_insert}_all_valid_combs_{self_pairs_train}.pt')
else:
self.all_valid_combs_path = os.path.join (
out_path, f'{name_insert}_all_valid_combs.pt')
self.years = [2019, 2020]
self.n_samples = n_samples
self.final_voxel_size = torch.tensor(final_voxel_size)
save_path = os.path.join(self.out_path, self.save_name)
voxel_size_icp = 0.05
if not preload:
with open(os.path.join(directory_path, 'args.json')) as f:
self.args = json.load(f)
with open(os.path.join(directory_path, 'response.json')) as f:
self.response = json.load(f)
print(f"Recreating dataset, saving to: {self.out_path}")
self.scans = [Scan(x, self.directory_path)
for x in self.response['RecordingProperties']]
self.scans = [
x for x in self.scans if x.datetime.year in self.years]
if os.path.isfile(self.filtered_scan_path):
with open(self.filtered_scan_path, "rb") as fp:
self.filtered_scans = pickle.load(fp)
else:
self.filtered_scans = filter_scans(self.scans, 3)
with open(self.filtered_scan_path, "wb") as fp:
pickle.dump(self.filtered_scans, fp)
self.save_dict = {}
save_id = -1
for scene_number, scan in enumerate(tqdm(self.filtered_scans)):
# Gather scans within certain distance of scan center
relevant_scans = [x for x in self.scans if np.linalg.norm(
x.center-scan.center) < 7]
relevant_times = set([x.datetime for x in relevant_scans])
# Group by dates
time_partitions = {time: [
x for x in relevant_scans if x.datetime == time] for time in relevant_times}
# Load and combine clouds from same date
clouds_per_time = [torch.from_numpy(np.concatenate([load_las(
x.path) for x in val])).double().to(device) for key, val in time_partitions.items()]
# Make xy 0 at center to avoid large values
center_trans = torch.cat(
(torch.from_numpy(scan.center), torch.tensor([0, 0, 0, 0]))).double().to(device)
clouds_per_time = [x-center_trans for x in clouds_per_time]
# Extract square at center since only those will be used for grid
clouds_per_time = [x[extract_area(x, center=np.array(
[0, 0]), clearance=self.clearance, shape='square'), :] for x in clouds_per_time]
# Apply registration between each cloud and first in list, store transforms
# First cloud does not need to be transformed
clouds_per_time = registration_pipeline(
clouds_per_time, voxel_size_icp, self.voxel_size_final_downsample)
# Remove below ground and above cutoff
# Cut off slightly under ground height
ground_cutoff = scan.ground_height - 0.05
height_cutoff = ground_cutoff+max_height
clouds_per_time = [x[torch.logical_and(
x[:, 2] > ground_cutoff, x[:, 2] < height_cutoff), ...] for x in clouds_per_time]
clouds_per_time = [x.float().cpu() for x in clouds_per_time]
save_id += 1
save_entry = {'clouds': clouds_per_time,
'ground_height': scan.ground_height}
self.save_dict[save_id] = save_entry
if scene_number % 100 == 0 and scene_number != 0:
print(f"Progressbackup: {scene_number}!")
torch.save(self.save_dict, save_path)
print(f"Saving to {save_path}!")
torch.save(self.save_dict, save_path)
else:
self.save_dict = torch.load(save_path)
if os.path.isfile(self.all_valid_combs_path):
self.all_valid_combs = torch.load(self.all_valid_combs_path)
else:
self.all_valid_combs = []
for idx, (save_id,save_entry) in enumerate(tqdm(self.save_dict.items())):
clouds = save_entry['clouds']
clouds = {index:x for index,x in enumerate(clouds) if x.shape[0] > 5000}
if len(clouds) < 2:
if self.verbose:
print(f'Not enough clouds {idx}, skipping ')
continue
cluster_min = torch.stack([torch.min(x,dim=0)[0][:3] for x in clouds.values()]).min(dim=0)[0]
cluster_max = torch.stack([torch.max(x,dim=0)[0][:3] for x in clouds.values()]).max(dim=0)[0]
clusters = {}
for index,x in clouds.items():
labels,voxel_centers = voxelize(x[:, :3],start= cluster_min,end=cluster_max,size= self.final_voxel_size)
clusters[index] = labels
valid_voxels = {}
for ind,cluster in clusters.items():
cluster_indices, counts = cluster.unique(return_counts=True)
valid_indices = cluster_indices[counts > self.n_samples_context]
valid_voxels[ind] = valid_indices
common_voxels = []
for ind_0, ind_1 in combinations(valid_voxels.keys(), 2):
if ind_0 == ind_1:
continue
common_voxels.append([ind_0, ind_1,[x.item() for x in valid_voxels[ind_0] if x in valid_voxels[ind_1]] ])
valid_combs = []
for val in common_voxels:
valid_combs.extend([(val[0], val[1], x) for x in val[2]])
# Self predict (only on index,since clouds shuffled and 1:1 other to same)
if self.mode == 'train' and self.self_pairs_train:
valid_combs.extend([(val[0], val[0], x) for x in val[2]])
if len(valid_combs) < 1:
# If not enough recursively give other index from dataset
continue
for draw_ind,draw in enumerate(valid_combs):
voxel_center = voxel_centers[draw[2]]
cloud_ind_0 = draw[0]
voxel_0 = get_voxel(clouds[cloud_ind_0],voxel_center,self.context_voxel_size)
if voxel_0.shape[0]>=self.n_samples_context:
final_comb = (save_id,draw[0],draw[1],draw[2])
self.all_valid_combs.append({'combination':final_comb,'voxel_center':voxel_center})
else:
print('Invalid')
continue
n_same =0
n_dif = 0
for comb_dict in self.all_valid_combs:
if comb_dict['combination'][1] == comb_dict['combination'][2]:
n_same+=1
else:
n_dif +=1
print(f"n_same/n_dif: {n_same/n_dif}")
torch.save(self.all_valid_combs,self.all_valid_combs_path)
print('Loaded dataset!')
def __process_frame(self, frame):
"""
Process a frame to apply thresholding and add the limits to the history
elements. A limit is valid as long as the white section reaches de end
of the band.
Parameters
----------
frame : ndarray
3-channel image.
Returns
-------
(ndarray, int, int) : Thresheld warped image for visualization, limit of
the left threshold, limit of the right threshold.
"""
# Frame in zenital view
warp = self.warp(frame)
# Limits to threshold the image in the Light channel
o_l, o_h = 0, 20
# Limits to threshold the image in RGB channels to detect orange color
r_l, r_h, g_l, g_h, b_l, b_h = 230, 255, 115, 170, 50, 85
# Arrays for thresheld images
th = []
gr = []
# Merge the upper and lower rectangles
box_left = (self.upper_left_rect[0], self.lower_left_rect[1])
box_right = (self.upper_right_rect[0], self.lower_right_rect[1])
# In every rectangle...
for idx, box in enumerate([box_left, box_right]):
# Box of interest
boi = extract_area(warp, box[0][0], box[0][1],
box[1][0] - box[0][0], box[1][1] - box[0][1])
# Thresholding based in color
b = boi[:, :, 0]
g = boi[:, :, 1]
r = boi[:, :, 2]
bools_color = (r > r_l) & (r < r_h) & (g > g_l) & (g < g_h) & (
b > b_l) & (b < b_h)
# binary = np.zeros_like(boi[:,:,0])
# Thresholding based on Lightness
boi = cv.cvtColor(boi, cv.COLOR_BGR2HLS)
# Lightness channel
boi = boi[:, :, 1]
# Pixels within the threshold limits
bools_light = (boi > o_l) & (boi < o_h)
# Thresheld box
binary = np.zeros_like(boi)
bools = bools_color | bools_light
# Valid pixels as white
binary[bools == True] = 1
# Find the white pixels
whites = np.argwhere(binary == 1)
left_nones = []
right_nones = []
if whites.shape[0] > 0:
# fill all space between first and last white pixel
rows = whites[:, 0]
rows = np.unique(rows)
low = rows[0]
high = rows[-1] + 1
binary[low:high, :] = 1
# Find the first white pixel
# If the box of whites extends to the end of the band, add the
# first white pixel
if boi.shape[0] - high < self.offset:
if idx == 0:
max_left = low
elif idx == 1:
max_right = low
# If the box of whites does not extend to the end of the band
else:
if idx == 0 and self.history_left.shape[0] > 0:
max_left = self.history_left[-1]
elif idx == 0 and self.history_left.shape[0] == 0:
max_left = None
if idx == 1 and self.history_right.shape[0] > 0:
max_right = self.history_right[-1]
elif idx == 1 and self.history_right.shape[0] == 0:
max_right = None
else:
if idx == 0 and self.history_left.shape[0] > 0:
max_left = self.history_left[-1]
elif idx == 0 and self.history_left.shape[0] == 0:
max_left = None
if idx == 1 and self.history_right.shape[0] > 0:
max_right = self.history_right[-1]
elif idx == 1 and self.history_right.shape[0] == 0:
max_right = None
th.append(binary)
# Grayscale image
box = extract_area(warp, box[0][0], box[0][1],
box[1][0] - box[0][0], box[1][1] - box[0][1])
box = cv.cvtColor(box, cv.COLOR_BGR2RGB)
gr.append(box)
# Left and right thresheld bands
th_image = np.hstack((th[0], th[1])) * 255
th_image = cv.resize(th_image,
(th_image.shape[1] * 3, th_image.shape[0] * 3))
# left and right grayscale bands
gr_image = np.hstack((gr[0], gr[1])) * 255
# Limit lines drawn in the grayscale bands
gr_image = cv.cvtColor(gr_image, cv.COLOR_BGR2GRAY)
upper = self.upper_left_rect[1][1] - self.offset
lower = self.lower_left_rect[1][1] - self.offset
gr_image[upper, :] = 0
gr_image[lower, :] = 0
gr_image = cv.resize(gr_image,
(gr_image.shape[1] * 3, gr_image.shape[0] * 3))
# Merge thresheld and grayscale images
gr_image = np.hstack((th_image, gr_image))
return gr_image.T, max_left, max_right