diff --git a/alodataset/transforms.py b/alodataset/transforms.py index 1734f2fc..73da7f6a 100644 --- a/alodataset/transforms.py +++ b/alodataset/transforms.py @@ -76,7 +76,6 @@ def __call__(self, frames: Union[Mapping[str, Frame], List[Frame], Frame], **kwa # Go through each image if isinstance(frames, dict): - n_set = {} if same_on_sequence is None or same_on_frames is None: @@ -86,11 +85,9 @@ def __call__(self, frames: Union[Mapping[str, Frame], List[Frame], Frame], **kwa ) for key in frames: - # Go throguh each element of the sequence # (If needed to apply save the params for each time step if "T" in frames[key].names and same_on_frames and not same_on_sequence: - n_set[key] = [] for t in range(0, frames[key].shape[0]): if t not in seqid2params: @@ -110,7 +107,6 @@ def __call__(self, frames: Union[Mapping[str, Frame], List[Frame], Frame], **kwa # Different for each element of the sequence, but we don't need to save # the params for each image neither elif "T" in frames[key].names and not same_on_frames and not same_on_sequence: - n_set[key] = [] for t in range(0, frames[key].shape[0]): @@ -216,8 +212,7 @@ def __init__(self, transforms1: AloTransform, transforms2: AloTransform, p: floa """ self.transforms1 = transforms1 self.transforms2 = transforms2 - self.p = p - super().__init__(*args, **kwargs) + super().__init__(p=p, *args, **kwargs) def sample_params(self): """Sample a `number` between and 1. The first transformation @@ -256,8 +251,7 @@ def __init__(self, p: float = 0.5, *args, **kwargs): p: float Probability to apply the transformation """ - self.p = p - super().__init__(*args, **kwargs) + super().__init__(p=p, *args, **kwargs) def sample_params(self): """Sample a `number` between and 1. The transformation @@ -376,7 +370,6 @@ def set_params(self, pad_left, pad_right, pad_top, pad_bottom): self._pad_bottom = pad_bottom def __call__(self, frame): - print((self._pad_top, self._pad_bottom), (self._pad_left, self._pad_right)) return frame.pad( @@ -453,53 +446,39 @@ def apply(self, frame: Frame): class RandomResizeWithAspectRatio(AloTransform): - def __init__(self, sizes: list, max_size: int = None, *args, **kwargs): - """Reszie the given given frame to a sampled `size` from the list of - given `sizes` so that the largest side is equal to `size` and always < to - `max_size` (if given). + def __init__(self, sizes: list, *args, **kwargs): + """Resize the given frame to a sampled `size` from the list of + given `sizes` so that the largest side is equal to `size`. Parameters ---------- sizes: list - List of int. Possible size to sample from + List of int. Possible sizes to sample from """ - assert isinstance(sizes, list) or max_size == None + assert isinstance(sizes, list) self.sizes = sizes - self.max_size = max_size super().__init__(*args, **kwargs) @staticmethod - def get_size_with_aspect_ratio(frame: Frame, size: int, max_size: int = None): - """Given a `frame` and a `size` this method compute a new size so that the largest - side is equal to `size` and always < to `max_size` (if given). + def get_size_with_aspect_ratio(frame: Frame, size: int): + """Given a `frame` and a `size` this method compute a new size so that the largest + side is equal to `size`. Parameters ---------- frame : Frame Frame to resize. Used only to get the width and the height of the target frame to resize. size: int - Desired size - max_size: int - Maximum size of the largest side. + Desired max size """ h, w = frame.H, frame.W - if max_size is not None: - min_original_size = float(min((w, h))) - max_original_size = float(max((w, h))) - if max_original_size / min_original_size * size > max_size: - size = int(round(max_size * min_original_size / max_original_size)) - - if (w <= h and w == size) or (h <= w and h == size): - return (h, w) - - if w < h: + if w > h: ow = size - oh = int(size * h / w) + oh = round(size * h / w) else: oh = size - ow = int(size * w / h) - + ow = round(size * w / h) return (oh, ow) def sample_params(self): @@ -521,7 +500,7 @@ def apply(self, frame: Frame): Frame to apply the transformation on """ # Sample one frame size - size = self.get_size_with_aspect_ratio(frame, self._size, self.max_size) + size = self.get_size_with_aspect_ratio(frame, self._size) # Resize frame frame = frame.resize(size) return frame @@ -910,13 +889,14 @@ def apply(self, frame: Frame, center: Union[Tuple[int, int], Tuple[float, float] class RandomFocusBlur(AloTransform): """Randomly introduces motion blur. - + Parameters ---------- max_filter_size : int Max filter size to use, the higher the more blured the image. - + """ + def __init__(self, max_filter_size=10, *args, **kwargs): assert isinstance(max_filter_size, int) self.max_filter_size = max_filter_size @@ -937,7 +917,7 @@ def sample_params(self): def set_params(self, h_size, v_size): self.h_filter_size = h_size self.v_filter_size = v_size - + @torch.no_grad() def apply(self, frame): c, h, w = frame.shape @@ -963,13 +943,14 @@ def apply(self, frame): class RandomFocusBlurV2(AloTransform): """Randomly introduces motion blur. - + Parameters ---------- max_filter_size : int Max filter size to use, the higher the more blured the image. """ + def __init__(self, max_filter_size=10, *args, **kwargs): assert isinstance(max_filter_size, int) self.max_filter_size = max_filter_size @@ -987,25 +968,25 @@ def sample_params(self): def set_params(self, h_size, v_size): self.h_filter_size = h_size self.v_filter_size = v_size - + @staticmethod def h_trans(frame, size): v_left_frames = [frame[:, :, i:] for i in range(1, size // 2 + 1)] - v_left_frames = [torch.nn.functional.pad(x, pad=(0, i + 1), value=0) for i, x in enumerate(v_left_frames)] - + v_left_frames = [torch.nn.functional.pad(x, pad=(0, i + 1), value=0) for i, x in enumerate(v_left_frames)] + v_right_frames = [frame[:, :, :-i] for i in range(1, size // 2 + 1)] - v_right_frames = [torch.nn.functional.pad(x, pad=(i + 1, 0), value=0) for i, x in enumerate(v_right_frames)] + v_right_frames = [torch.nn.functional.pad(x, pad=(i + 1, 0), value=0) for i, x in enumerate(v_right_frames)] v_frames = [*v_left_frames, frame, *v_right_frames] return v_frames - + @staticmethod def v_trans(frame, size): h_top_frames = [frame[:, i:, :] for i in range(1, size // 2 + 1)] - h_top_frames = [torch.nn.functional.pad(x, pad=(0, 0, 0, i + 1), value=0) for i, x in enumerate(h_top_frames)] - + h_top_frames = [torch.nn.functional.pad(x, pad=(0, 0, 0, i + 1), value=0) for i, x in enumerate(h_top_frames)] + h_bot_frames = [frame[:, :-i, :] for i in range(1, size // 2 + 1)] - h_bot_frames = [torch.nn.functional.pad(x, pad=(0, 0, i + 1, 0), value=0) for i, x in enumerate(h_bot_frames)] + h_bot_frames = [torch.nn.functional.pad(x, pad=(0, 0, i + 1, 0), value=0) for i, x in enumerate(h_bot_frames)] h_frames = [*h_top_frames, frame, *h_bot_frames] return h_frames @@ -1020,10 +1001,10 @@ def apply(self, frame): v_frame = sum(v_frames) / self.v_filter_size h_frame = sum(h_frames) / self.h_filter_size - + blured = (h_frame + v_frame) / 2 blured = Frame(blured) - + blured = blured.norm_as(frame) blured.__dict__ = frame.__dict__.copy() return blured @@ -1035,19 +1016,19 @@ def h_trans(frame, size): c, h, _ = frame.shape v_left_frames = [frame[:, :, i:] for i in range(1, size // 2 + 1)] v_left_frames = [torch.cat([f, torch.zeros((c, h, i + 1))], dim=2) for i, f in enumerate(v_left_frames)] - + v_right_frames = [frame[:, :, :-i] for i in range(1, size // 2 + 1)] v_right_frames = [torch.cat([torch.zeros((c, h, i + 1)), f], dim=2) for i, f in enumerate(v_right_frames)] v_frames = [*v_left_frames, frame, *v_right_frames] return v_frames - + @staticmethod def v_trans(frame, size): c, _, w = frame.shape h_top_frames = [frame[:, i:, :] for i in range(1, size // 2 + 1)] h_top_frames = [torch.cat([f, torch.zeros((c, i + 1, w))], dim=1) for i, f in enumerate(h_top_frames)] - + h_bot_frames = [frame[:, :-i, :] for i in range(1, size // 2 + 1)] h_bot_frames = [torch.cat([torch.zeros((c, i + 1, w)), f], dim=1) for i, f in enumerate(h_bot_frames)] @@ -1057,7 +1038,7 @@ def v_trans(frame, size): class RandomFlowMotionBlur(AloTransform): """Introduces motion blur from optical flow. - + Idea : Let OpticalFlow : x, y --> x', y' retrive the indexes betwe x, x' and y, y' i.e x -> x1 ... -> x' , y -> y1 ... -> y' @@ -1076,20 +1057,21 @@ class RandomFlowMotionBlur(AloTransform): Motion blur intensity. If this arg is set, the value will not be random anymore. """ + def __init__( - self, - subframes: int = 10, - flow_model=None, - model_kwargs={}, - intensity=None, - **kwargs, - ): + self, + subframes: int = 10, + flow_model=None, + model_kwargs={}, + intensity=None, + **kwargs, + ): if isinstance(intensity, list): assert all([isinstance(x, float) for x in intensity]) assert intensity[0] < intensity[1] assert len(intensity) == 2 - self.intensity = 1. if intensity is None else intensity + self.intensity = 1.0 if intensity is None else intensity self.model_kwargs = model_kwargs self.flow_model = flow_model self.inter_intensity = None @@ -1116,12 +1098,12 @@ def _get_flow_model_kwargs(self, frame1, frame2): frame2 = Frame(frame2).norm_minmax_sym().batch() return {"frame1": frame1, "frame2": frame2, **self.model_kwargs} - + @staticmethod def _adapt_model_output(output): - """Adapts model output to be an optical flow of size [2, H, W] where the first channel + """Adapts model output to be an optical flow of size [2, H, W] where the first channel is the OF over X axis and the second is over Y axis - + Example with alonet/raft/raft ... -> """ @@ -1152,7 +1134,7 @@ def apply(self, frame, flow=None, p_frame=None): if HW != HW_: flow = torch.nn.functional.interpolate(flow.unsqueeze(0), size=HW_, mode="nearest") flow = flow.squeeze() - + flow = flow * self.inter_intensity # XY Coordinates @@ -1164,18 +1146,20 @@ def apply(self, frame, flow=None, p_frame=None): # Map coridinates of intermediate points X -> X, intemediate X points ..., X + X_displacement (same for Y) subcoords = [ [ - (coords[0] - map_coords[0]) * i / self.subframes + coords[0], # X - (coords[0] - map_coords[0]) * i / self.subframes + coords[1]] # Y - for i in range(self.subframes + 1) - ] + (coords[0] - map_coords[0]) * i / self.subframes + coords[0], # X + (coords[0] - map_coords[0]) * i / self.subframes + coords[1], + ] # Y + for i in range(self.subframes + 1) + ] # Round and clamp indexes (float -> int + Occlusion) subcoords = [ [ torch.round(torch.clamp(s[0], min=0, max=HW_[0] - 1)).long(), - torch.round(torch.clamp(s[1], min=0, max=HW_[1] - 1)).long() + torch.round(torch.clamp(s[1], min=0, max=HW_[1] - 1)).long(), ] - for s in subcoords] + for s in subcoords + ] # Frame to indexed intermediate frames frame_ = [frame_[:, subcoord[0], subcoord[1]] for subcoord in subcoords] @@ -1203,18 +1187,19 @@ class RandomCornersMask(AloTransform): ## p_sides = [top, bottom, right, left] """ + def __init__( - self, - max_mask_size: float = 0.2, - p_sides: List = [0.2, 0.2, 0.2, 0.2], - **kwargs, - ): + self, + max_mask_size: float = 0.2, + p_sides: List = [0.2, 0.2, 0.2, 0.2], + **kwargs, + ): assert len(p_sides) == 4 assert isinstance(p_sides, list) assert isinstance(max_mask_size, float) assert max_mask_size >= 0 and max_mask_size < 1 assert all([isinstance(x, float) for x in p_sides]) - + # Random var param self.max_mask_size = max_mask_size self.p_sides = p_sides