Skip to content

Commit

Permalink
Fix cropping when rendering video (#842)
Browse files Browse the repository at this point in the history
  • Loading branch information
@roomrys
roomrys authored Aug 3, 2022
1 parent 28b395d commit 44e4661
Show file tree
Hide file tree
Showing 2 changed files with 119 additions and 82 deletions.
181 changes: 101 additions & 80 deletions sleap/io/visuals.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,38 +49,41 @@ def reader(out_q: Queue, video: Video, frames: List[int], scale: float = 1.0):

logger.info(f"Chunks: {chunk_count}, chunk size: {chunk_size}")

i = 0
for chunk_i in range(chunk_count):

# Read the next chunk of frames
frame_start = chunk_size * chunk_i
frame_end = min(frame_start + chunk_size, total_count)
frames_idx_chunk = frames[frame_start:frame_end]
try:
i = 0
for chunk_i in range(chunk_count):

t0 = perf_counter()
# Read the next chunk of frames
frame_start = chunk_size * chunk_i
frame_end = min(frame_start + chunk_size, total_count)
frames_idx_chunk = frames[frame_start:frame_end]

# Safely load frames from video, skipping frames we can't load
loaded_chunk_idxs, video_frame_images = video.get_frames_safely(
frames_idx_chunk
)
t0 = perf_counter()

if not loaded_chunk_idxs:
print(f"No frames could be loaded from chunk {chunk_i}")
i += 1
continue
# Safely load frames from video, skipping frames we can't load
loaded_chunk_idxs, video_frame_images = video.get_frames_safely(
frames_idx_chunk
)

if scale != 1.0:
video_frame_images = resize_images(video_frame_images, scale)
if not loaded_chunk_idxs:
print(f"No frames could be loaded from chunk {chunk_i}")
i += 1
continue

elapsed = perf_counter() - t0
fps = len(loaded_chunk_idxs) / elapsed
logger.debug(f"reading chunk {i} in {elapsed} s = {fps} fps")
i += 1
if scale != 1.0:
video_frame_images = resize_images(video_frame_images, scale)

out_q.put((loaded_chunk_idxs, video_frame_images))
elapsed = perf_counter() - t0
fps = len(loaded_chunk_idxs) / elapsed
logger.debug(f"reading chunk {i} in {elapsed} s = {fps} fps")
i += 1

# send _sentinal object into queue to signal that we're done
out_q.put(_sentinel)
out_q.put((loaded_chunk_idxs, video_frame_images))
except Exception as e:
raise e
finally:
# send _sentinal object into queue to signal that we're done
out_q.put(_sentinel)


def writer(
Expand Down Expand Up @@ -112,36 +115,42 @@ def writer(
total_frames_written = 0
start_time = perf_counter()
i = 0
while True:
data = in_q.get()

if data is _sentinel:
# no more data to be received so stop
in_q.put(_sentinel)
break
try:
while True:
data = in_q.get()

if writer_object is None and data:
h, w = data[0].shape[:2]
writer_object = VideoWriter.safe_builder(
filename, height=h, width=w, fps=fps
)
if data is _sentinel:
# no more data to be received so stop
in_q.put(_sentinel)
break

t0 = perf_counter()
for img in data:
writer_object.add_frame(img, bgr=True)
if writer_object is None and data:
h, w = data[0].shape[:2]
writer_object = VideoWriter.safe_builder(
filename, height=h, width=w, fps=fps
)

elapsed = perf_counter() - t0
fps = len(data) / elapsed
logger.debug(f"writing chunk {i} in {elapsed} s = {fps} fps")
i += 1
t0 = perf_counter()
for img in data:
writer_object.add_frame(img, bgr=True)

total_frames_written += len(data)
total_elapsed = perf_counter() - start_time
progress_queue.put((total_frames_written, total_elapsed))
elapsed = perf_counter() - t0
fps = len(data) / elapsed
logger.debug(f"writing chunk {i} in {elapsed} s = {fps} fps")
i += 1

writer_object.close()
# send (-1, time) to signal done
progress_queue.put((-1, total_elapsed))
total_frames_written += len(data)
total_elapsed = perf_counter() - start_time
progress_queue.put((total_frames_written, total_elapsed))
except Exception as e:
# Stop receiving data
in_q.put(_sentinel)
raise e
finally:
if writer_object is not None:
writer_object.close()
# Send (-1, time) to signal done
progress_queue.put((-1, total_elapsed))


class VideoMarkerThread(Thread):
Expand Down Expand Up @@ -217,32 +226,38 @@ def run(self):
def marker(self):
cv2.setNumThreads(usable_cpu_count())

chunk_i = 0
while True:
data = self.in_q.get()
try:
chunk_i = 0
while True:
data = self.in_q.get()

if data is _sentinel:
# no more data to be received so stop
self.in_q.put(_sentinel)
break
if data is _sentinel:
# no more data to be received so stop
self.in_q.put(_sentinel)
break

frames_idx_chunk, video_frame_images = data
frames_idx_chunk, video_frame_images = data

t0 = perf_counter()
t0 = perf_counter()

imgs = self._mark_images(
frame_indices=frames_idx_chunk,
frame_images=video_frame_images,
)
imgs = self._mark_images(
frame_indices=frames_idx_chunk,
frame_images=video_frame_images,
)

elapsed = perf_counter() - t0
fps = len(imgs) / elapsed
logger.debug(f"drawing chunk {chunk_i} in {elapsed} s = {fps} fps")
chunk_i += 1
self.out_q.put(imgs)
elapsed = perf_counter() - t0
fps = len(imgs) / elapsed
logger.debug(f"drawing chunk {chunk_i} in {elapsed} s = {fps} fps")
chunk_i += 1
self.out_q.put(imgs)
except Exception as e:
# Stop receiving data
self.in_q.put(_sentinel)
raise e

# send _sentinal object into queue to signal that we're done
self.out_q.put(_sentinel)
finally:
# Send _sentinal object into queue to signal that we're done
self.out_q.put(_sentinel)

def _mark_images(self, frame_indices, frame_images):
imgs = []
Expand All @@ -255,7 +270,7 @@ def _mark_images(self, frame_indices, frame_images):
return imgs

def _mark_single_frame(self, video_frame: np.ndarray, frame_idx: int) -> np.ndarray:
"""Returns single annotated frame image.
"""Return single annotated frame image.
Args:
video_frame: The ndarray of the frame image.
Expand All @@ -264,21 +279,27 @@ def _mark_single_frame(self, video_frame: np.ndarray, frame_idx: int) -> np.ndar
Returns:
ndarray of frame image with visual annotations added.
"""

# Use OpenCV to convert to BGR color image
video_frame = img_to_cv(video_frame)

# Add the instances to the image
overlay = self._plot_instances_cv(video_frame.copy(), frame_idx)

return cv2.addWeighted(overlay, self.alpha, video_frame, 1 - self.alpha, 0)
# Crop video_frame to same size as overlay
video_frame_cropped = (
self._crop_frame(video_frame.copy())[0] if self.crop else video_frame
)

return cv2.addWeighted(
overlay, self.alpha, video_frame_cropped, 1 - self.alpha, 0
)

def _plot_instances_cv(
self,
img: np.ndarray,
frame_idx: int,
) -> Optional[np.ndarray]:
"""Adds visuals annotations to single frame image.
) -> np.ndarray:
"""Add visual annotations to single frame image.
Args:
img: The ndarray of the frame image.
Expand All @@ -293,7 +314,7 @@ def _plot_instances_cv(
lfs = labels.find(labels.videos[video_idx], frame_idx)

if len(lfs) == 0:
return self._crop_frame(img) if self.crop else img
return self._crop_frame(img)[0] if self.crop else img

instances = lfs[0].instances_to_show

Expand All @@ -311,9 +332,9 @@ def _get_crop_center(
) -> Tuple[int, int]:
if instances:
centroids = np.array([inst.centroid for inst in instances])
center_xy = np.median(centroids, axis=0)

center_xy = np.nanmedian(centroids, axis=0)
self._crop_centers.append(center_xy)

elif not self._crop_centers:
# no crops so far and no instances yet so just use image center
img_w, img_h = img.shape[:2]
Expand All @@ -322,7 +343,7 @@ def _get_crop_center(
self._crop_centers.append(center_xy)

# use a running average of the last N centers to smooth movement
center_xy = tuple(np.mean(np.stack(self._crop_centers), axis=0))
center_xy = tuple(np.nanmean(np.stack(self._crop_centers), axis=0))

return center_xy

Expand Down Expand Up @@ -520,7 +541,7 @@ def save_labeled_video(
progress_win.setValue(frames_complete)
else:
print(
f"Finished {frames_complete} frames in {elapsed} s, fps = {fps}, approx {remaining_time} s remaining"
f"Finished {frames_complete} frames in {elapsed:.1f} s, fps = {round(fps)}, approx {remaining_time:.1f} s remaining"
)

elapsed = perf_counter() - t0
Expand Down
20 changes: 18 additions & 2 deletions tests/io/test_visuals.py
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
import numpy as np
import os
import pytest
from sleap.io.dataset import Labels
from sleap.io.visuals import (
save_labeled_video,
resize_images,
Expand Down Expand Up @@ -67,13 +69,27 @@ def test_sleap_render(centered_pair_predictions):
assert os.path.exists("testvis.avi")


def test_write_visuals(tmpdir, centered_pair_predictions):
@pytest.mark.parametrize("crop", ["Half", "Quarter", None])
def test_write_visuals(tmpdir, centered_pair_predictions: Labels, crop: str):
labels = centered_pair_predictions
video = centered_pair_predictions.videos[0]

# Determine crop size relative to original size and scale
crop_size_xy = None
w = int(video.backend.width)
h = int(video.backend.height)
if crop == "Half":
crop_size_xy = (w // 2, h // 2)
elif crop == "Quarter":
crop_size_xy = (w // 4, h // 4)

path = os.path.join(tmpdir, "clip.avi")
save_labeled_video(
filename=path,
labels=centered_pair_predictions,
video=centered_pair_predictions.videos[0],
video=video,
frames=(0, 1, 2),
fps=15,
crop_size_xy=crop_size_xy,
)
assert os.path.exists(path)

0 comments on commit 44e4661

Please sign in to comment.