def save_track(dataset, track, folder, labels_dir, ext="png"):
track_data = dataset.fetch_track(track, original=True)
clip_meta = dataset.db.get_clip_meta(track.clip_id)
clip_dir = os.path.join(folder, str(track.clip_id))
if not os.path.isdir(clip_dir):
logging.debug("Creating %s", clip_dir)
os.mkdir(clip_dir)
background = dataset.db.get_clip_background(track.clip_id)
if background is not None:
background, _ = normalize(background, new_max=255)
img = Image.fromarray(np.uint8(background))
filename = "{}-background.{}".format(track.clip_id, ext)
logging.debug("Saving %s", os.path.join(clip_dir, filename))
img.save(os.path.join(clip_dir, filename))
for i, frame in enumerate(track_data):
thermal = frame.thermal
normed, _ = normalize(thermal, new_max=255)
img = Image.fromarray(np.uint8(normed))
filename = "{}-{}.{}".format(track.clip_id, i + track.start_frame, ext)
logging.debug("Saving %s", os.path.join(clip_dir, filename))
img.save(os.path.join(clip_dir, filename))
save_metadata(clip_meta, track, labels_dir)
def preprocess_frame(
frame,
frame_size,
augment,
thermal_median,
velocity,
output_dim,
preprocess_fn=None,
sample=None,
):
processed_frame, flipped = preprocess_segment(
[frame],
frame_size,
reference_level=[thermal_median],
augment=augment,
default_inset=0,
)
if len(processed_frame) == 0:
return
processed_frame = processed_frame[0]
thermal = processed_frame.get_channel(TrackChannels.thermal)
filtered = processed_frame.get_channel(TrackChannels.filtered)
thermal, stats = imageprocessing.normalize(thermal, min=0)
if not stats[0]:
return None
filtered, stats = imageprocessing.normalize(filtered, min=0)
if not stats[0]:
return None
data = np.empty((*thermal.shape, 3))
data[:, :, 0] = thermal
data[:, :, 1] = filtered
data[:, :, 2] = filtered
# for testing
# tools.saveclassify_image(
# data,
# f"samples/{sample.label}-{sample.clip_id}-{sample.track_id}",
# )
# preprocess expects values in range 0-255
if preprocess_fn:
data = data * 255
data = preprocess_fn(data)
return data
def preprocess_movement(
data,
segment,
frames_per_row,
regions,
channel,
preprocess_fn=None,
augment=False,
):
segment = preprocess_segment(segment,
augment=augment,
filter_to_delta=False,
default_inset=0)
segment = segment[:, channel]
# as long as one frame it's fine
square, success = imageprocessing.square_clip(segment, frames_per_row,
(FRAME_SIZE, FRAME_SIZE),
type)
if not success:
return None
dots, overlay = imageprocessing.movement_images(
data,
regions,
dim=square.shape,
channel=channel,
require_movement=True,
)
dots = dots / 255
overlay, success = imageprocessing.normalize(overlay, min=0)
if not success:
return None
data = np.empty((square.shape[0], square.shape[1], 3))
data[:, :, 0] = square
data[:, :, 1] = dots # dots
data[:, :, 2] = overlay # overlay
if preprocess_fn:
for i, frame in enumerate(data):
frame = frame * 255
data[i] = preprocess_fn(frame)
return data
def _get_filtered_frame(self, clip, thermal):
"""
Calculates filtered frame from thermal
:param thermal: the thermal frame
:param background: (optional) used for background subtraction
:return: uint8 filtered frame and adjusted clip threshold for normalized frame
"""
filtered = np.float32(thermal.copy())
avg_change = int(round(np.average(thermal) - clip.stats.mean_background_value))
np.clip(filtered - clip.background - avg_change, 0, None, out=filtered)
filtered, stats = normalize(filtered, new_max=255)
filtered = cv2.fastNlMeansDenoising(np.uint8(filtered), None)
if stats[1] == stats[2]:
mapped_thresh = clip.background_thresh
else:
mapped_thresh = clip.background_thresh / (stats[1] - stats[2]) * 255
return filtered, mapped_thresh
def create_four_tracking_image(self, frame, min_temp, max_temp):
thermal = frame.thermal
filtered = frame.filtered + min_temp
filtered = tools.convert_heat_to_img(filtered, self.colourmap,
min_temp, max_temp)
thermal = tools.convert_heat_to_img(thermal, self.colourmap, min_temp,
max_temp)
if self.debug:
tools.add_heat_number(thermal, frame.thermal, 1)
if frame.mask is None:
mask = np.zeros((np.array(thermal).shape), dtype=np.uint8)
else:
mask, _ = normalize(frame.mask, new_max=255)
mask = np.uint8(mask)
mask = mask[..., np.newaxis]
mask = np.repeat(mask, 3, axis=2)
mask = Image.fromarray(mask)
flow_h, flow_v = frame.get_flow_split(clip_flow=True)
if flow_h is None and flow_v is None:
flow_magnitude = Image.fromarray(
np.zeros((np.array(thermal).shape), dtype=np.uint8))
else:
flow_magnitude = (
np.linalg.norm(np.float32([flow_h, flow_v]), ord=2, axis=0) /
4.0 + min_temp)
flow_magnitude = tools.convert_heat_to_img(flow_magnitude,
self.colourmap,
min_temp, max_temp)
image = np.hstack((np.vstack(
(thermal, mask)), np.vstack((filtered, flow_magnitude))))
image = Image.fromarray(image)
image = image.resize(
(
int(image.width * 4),
int(image.height * 4),
),
Image.BILINEAR,
)
return image
def preprocess_frame(
data, output_dim, use_thermal=True, augment=False, preprocess_fn=None
):
if use_thermal:
channel = TrackChannels.thermal
else:
channel = TrackChannels.filtered
data = data[channel]
data, stats = imageprocessing.normalize(data)
if not stats[0]:
return None
data = data[np.newaxis, :]
data = np.transpose(data, (1, 2, 0))
data = np.repeat(data, output_dim[2], axis=2)
data = imageprocessing.resize_cv(data, output_dim, channel)
# preprocess expects values in range 0-255
if preprocess_fn:
data = data * 255
data = preprocess_fn(data)
return data
def generate_optical_flow(self, opt_flow, prev_frame, flow_threshold=40):
"""
Generate optical flow from thermal frames
:param opt_flow: An optical flow algorithm
"""
height, width = self.thermal.shape
flow = np.zeros([height, width, 2], dtype=np.float32)
scaled_thermal = self.thermal.copy()
scaled_thermal[self.mask == 0] = 0
scaled_thermal, _ = normalize(scaled_thermal, new_max=255)
scaled_thermal = np.float32(scaled_thermal)
# threshold = np.median(self.thermal) + flow_threshold
# scaled_thermal = np.uint8(np.clip(self.thermal - threshold, 0, 255))
if prev_frame is not None:
# for some reason openCV spins up lots of threads for this which really slows things down, so we
# cap the threads to 2
cv2.setNumThreads(2)
flow = opt_flow.calc(prev_frame.scaled_thermal, scaled_thermal,
flow)
self.scaled_thermal = scaled_thermal
self.flow = flow
if prev_frame:
prev_frame.scaled_thermal = None
def preprocess_movement(
data,
segment,
frames_per_row,
frame_size,
regions,
red_channel,
preprocess_fn=None,
augment=False,
green_type=None,
keep_aspect=False,
reference_level=None,
):
segment = preprocess_segment(
segment,
reference_level=reference_level,
augment=augment,
filter_to_delta=False,
default_inset=0,
keep_aspect=keep_aspect,
frame_size=frame_size,
)
red_segment = segment[:, red_channel]
# as long as one frame it's fine
red_square, success = imageprocessing.square_clip(
red_segment, frames_per_row, (frame_size, frame_size), type
)
if not success:
return None
_, overlay = imageprocessing.movement_images(
data,
regions,
dim=red_square.shape,
require_movement=True,
)
overlay, stats = imageprocessing.normalize(overlay, min=0)
if not stats[0]:
return None
data = np.empty((red_square.shape[0], red_square.shape[1], 3))
data[:, :, 0] = red_square
if green_type == FrameTypes.filtered_square:
green_segment = segment[:, TrackChannels.filtered]
green_square, success = imageprocessing.square_clip(
green_segment, frames_per_row, (frame_size, frame_size), type
)
if not success:
return None
elif green_type == FrameTypes.thermal_square:
green_segment = segment[:, TrackChannels.thermal]
green_square, success = imageprocessing.square_clip(
green_segment, frames_per_row, (frame_size, frame_size), type
)
if not success:
return None
elif green_type == FrameTypes.overlay:
green_square = overlay
else:
green_square = np.zeros(overlay.shape)
data[:, :, 1] = green_square
data[:, :, 2] = overlay # overlay
if preprocess_fn:
data = data * 255
data = preprocess_fn(data)
return data
def normalize(self):
if self.thermal is not None:
self.thermal, _ = normalize(self.thermal, new_max=255)
if self.filtered is not None:
self.filtered, _ = normalize(self.filtered, new_max=255)
def preprocess_movement(
segment,
frames_per_row,
frame_size,
red_type,
green_type,
blue_type,
preprocess_fn=None,
augment=False,
reference_level=None,
sample=None,
keep_edge=False,
):
segment, flipped = preprocess_segment(
segment,
frame_size,
reference_level=reference_level,
augment=augment,
default_inset=0,
keep_edge=keep_edge,
)
frame_types = {}
channel_types = set([green_type, blue_type, red_type])
for type in channel_types:
if type == FrameTypes.overlay:
if overlay is None:
overlay = imageprocessing.overlay_image(
data,
regions,
dim=(frames_per_row * frame_size,
frames_per_row * frame_size),
require_movement=True,
)
channel_data, stats = imageprocessing.normalize(overlay, min=0)
if not stats[0]:
return None
else:
channel_data = np.zeros((square.shape[0], square.shape[1]))
channel_data[:overlay.shape[0], :overlay.shape[1]] = overlay
if flipped:
channel_data = np.flip(channel_data, axis=1)
elif type == FrameTypes.flow_tiled:
channel_segment = [
frame.get_channel(TrackChannels.flow) for frame in segment
]
channel_data, success = imageprocessing.square_clip_flow(
channel_segment, frames_per_row, (frame_size, frame_size))
if not success:
return None
else:
if type == FrameTypes.thermal_tiled:
channel = TrackChannels.thermal
else:
channel = TrackChannels.filtered
channel_segment = [frame.get_channel(channel) for frame in segment]
channel_data, success = imageprocessing.square_clip(
channel_segment, frames_per_row, (frame_size, frame_size))
if not success:
return None
frame_types[type] = channel_data
data = np.stack((frame_types[red_type], frame_types[green_type],
frame_types[blue_type]),
axis=2)
# for testing
# tools.saveclassify_image(
# data,
# f"samples/{sample}",
# )
if preprocess_fn:
data = data * 255
data = preprocess_fn(data)
return data
def remove_background_animals(self, initial_frame, initial_diff):
"""
Try and remove animals that are already in the initial frames, by
checking for connected components in the intital_diff frame
(this is the maximum change between first frame and all other frames in the clip)
"""
# remove some noise
initial_diff[initial_diff < self.background_thresh] = 0
initial_diff[initial_diff > 255] = 255
initial_diff = np.uint8(initial_diff)
initial_diff = cv2.fastNlMeansDenoising(initial_diff, None)
_, lower_mask, lower_objects = detect_objects(initial_diff, otsus=True)
max_region = Region(0, 0, self.res_x, self.res_y)
for component in lower_objects[1:]:
region = Region(component[0], component[1], component[2],
component[3])
region.enlarge(2, max=max_region)
if region.width >= self.res_x or region.height >= self.res_y:
logging.info(
"Background animal bigger than max, probably false positive %s %s",
region,
component[4],
)
continue
background_region = region.subimage(initial_frame)
norm_back = background_region.copy()
norm_back, _ = normalize(norm_back, new_max=255)
sub_components, sub_connected, sub_stats = detect_objects(
norm_back, otsus=True)
if sub_components <= 1:
continue
overlap_image = region.subimage(lower_mask) * 255
overlap_pixels = np.sum(sub_connected[overlap_image > 0])
overlap_pixels = overlap_pixels / float(component[4])
# filter out components which are too big, or dont match original causes
# for filtering
if (overlap_pixels < Clip.MIN_ORIGIN_OVERLAP
or sub_stats[1][4] == 0 or sub_stats[1][4] == region.area):
logging.info(
"Invalid components mass: %s, components: %s region area %s overlap %s",
sub_stats[1][4],
sub_components,
region.area,
overlap_pixels,
)
continue
sub_connected[sub_connected > 0] = 1
# remove this component from the background by painting with
# colours of neighbouring pixels
background_region[:] = cv2.inpaint(
np.float32(background_region),
np.uint8(sub_connected),
3,
cv2.INPAINT_TELEA,
)
return initial_frame
