Memory Leak in sam2.sam2.utils.misc

When you use sam2 video prediction with asyncronous frame loading, the call starts a frame loader with a thread.
This shuts when the python script finishes as a whole, but it sits munching RAM until the end of time if the python thread doesn't get terminated (or if, like me, you don't want it to end just yet).

It would be a decent idea to make a context manager type thing so that it doesn't need to store all the images within memory after the full prediction set is made.

this is the offending class:

class AsyncVideoFrameLoader:
"""
A list of video frames to be load asynchronously without blocking session start.
"""

def __init__(
self,
img_paths,
image_size,
offload_video_to_cpu,
img_mean,
img_std,
compute_device,
):
self.img_paths = img_paths
self.image_size = image_size
self.offload_video_to_cpu = offload_video_to_cpu
self.img_mean = img_mean
self.img_std = img_std
# items in `self.images` will be loaded asynchronously
self.images = [None] * len(img_paths)
# catch and raise any exceptions in the async loading thread
self.exception = None
# video_height and video_width be filled when loading the first image
self.video_height = None
self.video_width = None
self.compute_device = compute_device

# load the first frame to fill video_height and video_width and also
# to cache it (since it's most likely where the user will click)
self.__getitem__(0)

# load the rest of frames asynchronously without blocking the session start
def _load_frames():
try:
for n in tqdm(range(len(self.images)), desc="frame loading (JPEG)"):
self.__getitem__(n)
except Exception as e:
self.exception = e

self.thread = Thread(target=_load_frames, daemon=True)
self.thread.start()

def __getitem__(self, index):
if self.exception is not None:
raise RuntimeError("Failure in frame loading thread") from self.exception

img = self.images[index]
if img is not None:
return img

img, video_height, video_width = _load_img_as_tensor(
self.img_paths[index], self.image_size
)
self.video_height = video_height
self.video_width = video_width
# normalize by mean and std
img -= self.img_mean
img /= self.img_std
if not self.offload_video_to_cpu:
img = img.to(self.compute_device, non_blocking=True)
self.images[index] = img
return img

def __len__(self):
return len(self.images)