doanh25032004

Upload folder using huggingface_hub

872b0a0 verified 8 months ago

3.28 kB

import numpy as np
import torch

# import torch.nn.functional as F


class ObjectCache:
    def __init__(self, cache_size=500):
        self.cache_size = cache_size
        self._obj_mask_cache = None
        self._img_cache = None
        self._motion_cache = None
        self.count = 0

    # We initialize the cache when we receive the first push sample so that it adapts to the current img size automatically
    def init_cache(self, img_size):
        self._obj_mask_cache = torch.zeros(
            (self.cache_size, 1, *img_size), dtype=torch.float32
        )
        self._img_cache = torch.zeros(
            (self.cache_size, 3, *img_size), dtype=torch.float32
        )
        self._motion_cache = torch.zeros((self.cache_size, 2), dtype=torch.float32)
        return

    def pop(self, B=8, with_aug=True):  # we do not remove objects after popping
        if self.count < self.cache_size:  # do not use it before it is full
            return None

        idx = np.random.choice(self.cache_size, B, replace=False)
        obj_mask = self._obj_mask_cache[idx]
        img = self._img_cache[idx]
        motion = self._motion_cache[idx]

        if with_aug:
            rand_scale = (
                torch.rand(B) * 0.7 + 0.8
            )  # randomly rescale motion by 0.8-1.5 times
            rand_scale *= (-1) ** (
                torch.rand(B) > 0.5
            ).float()  # randomly reverse motion
            motion = motion * rand_scale[:, None]

            flip_flag = torch.rand(B) > 0.5  # randomly horitontal-flip obj mask
            img[flip_flag] = img[flip_flag].flip(dims=[3])
            obj_mask[flip_flag] = obj_mask[flip_flag].flip(dims=[3])
            motion[flip_flag, 0] *= -1

        return obj_mask, img, motion

    def push(self, obj_mask, img, motion):
        """
        obj_mask: [B, 1, H, W]
        img: [B, 3, H, W]
        motion: [B, 2]
        """

        if self._obj_mask_cache is None:
            self.init_cache(img_size=img.shape[-2:])

        B = obj_mask.shape[0]

        if self.count <= self.cache_size - B:  # many spaces
            self._obj_mask_cache[self.count : (self.count + B)] = obj_mask
            self._img_cache[self.count : (self.count + B)] = img
            self._motion_cache[self.count : (self.count + B)] = motion
            self.count += B
            return

        elif self.count < self.cache_size:  # partial space
            space = self.cache_size - self.count
            self._obj_mask_cache[self.count :] = obj_mask[:space]
            self._img_cache[self.count :] = img[:space]
            self._motion_cache[self.count :] = motion[:space]

            overwrite_idx = np.random.choice(self.count, B - space, replace=False)
            self._obj_mask_cache[overwrite_idx] = obj_mask[space:]
            self._img_cache[overwrite_idx] = img[space:]
            self._motion_cache[overwrite_idx] = motion[space:]
            self.count += space
            return

        else:  # no spaces; random overwrite
            overwrite_idx = np.random.choice(self.cache_size, B, replace=False)
            self._obj_mask_cache[overwrite_idx] = obj_mask
            self._img_cache[overwrite_idx] = img
            self._motion_cache[overwrite_idx] = motion
            return

	import numpy as np
	import torch

	# import torch.nn.functional as F


	class ObjectCache:
	def __init__(self, cache_size=500):
	self.cache_size = cache_size
	self._obj_mask_cache = None
	self._img_cache = None
	self._motion_cache = None
	self.count = 0

	# We initialize the cache when we receive the first push sample so that it adapts to the current img size automatically
	def init_cache(self, img_size):
	self._obj_mask_cache = torch.zeros(
	(self.cache_size, 1, *img_size), dtype=torch.float32
	)
	self._img_cache = torch.zeros(
	(self.cache_size, 3, *img_size), dtype=torch.float32
	)
	self._motion_cache = torch.zeros((self.cache_size, 2), dtype=torch.float32)
	return

	def pop(self, B=8, with_aug=True): # we do not remove objects after popping
	if self.count < self.cache_size: # do not use it before it is full
	return None

	idx = np.random.choice(self.cache_size, B, replace=False)
	obj_mask = self._obj_mask_cache[idx]
	img = self._img_cache[idx]
	motion = self._motion_cache[idx]

	if with_aug:
	rand_scale = (
	torch.rand(B) * 0.7 + 0.8
	) # randomly rescale motion by 0.8-1.5 times
	rand_scale = (-1) * (
	torch.rand(B) > 0.5
	).float() # randomly reverse motion
	motion = motion * rand_scale[:, None]

	flip_flag = torch.rand(B) > 0.5 # randomly horitontal-flip obj mask
	img[flip_flag] = img[flip_flag].flip(dims=[3])
	obj_mask[flip_flag] = obj_mask[flip_flag].flip(dims=[3])
	motion[flip_flag, 0] *= -1

	return obj_mask, img, motion

	def push(self, obj_mask, img, motion):
	"""
	obj_mask: [B, 1, H, W]
	img: [B, 3, H, W]
	motion: [B, 2]
	"""

	if self._obj_mask_cache is None:
	self.init_cache(img_size=img.shape[-2:])

	B = obj_mask.shape[0]

	if self.count <= self.cache_size - B: # many spaces
	self._obj_mask_cache[self.count : (self.count + B)] = obj_mask
	self._img_cache[self.count : (self.count + B)] = img
	self._motion_cache[self.count : (self.count + B)] = motion
	self.count += B
	return

	elif self.count < self.cache_size: # partial space
	space = self.cache_size - self.count
	self._obj_mask_cache[self.count :] = obj_mask[:space]
	self._img_cache[self.count :] = img[:space]
	self._motion_cache[self.count :] = motion[:space]

	overwrite_idx = np.random.choice(self.count, B - space, replace=False)
	self._obj_mask_cache[overwrite_idx] = obj_mask[space:]
	self._img_cache[overwrite_idx] = img[space:]
	self._motion_cache[overwrite_idx] = motion[space:]
	self.count += space
	return

	else: # no spaces; random overwrite
	overwrite_idx = np.random.choice(self.cache_size, B, replace=False)
	self._obj_mask_cache[overwrite_idx] = obj_mask
	self._img_cache[overwrite_idx] = img
	self._motion_cache[overwrite_idx] = motion
	return