Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
141 changes: 63 additions & 78 deletions alodataset/transforms.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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:
Expand All @@ -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]):
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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):
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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)]

Expand All @@ -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'
Expand All @@ -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
Expand All @@ -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 ... ->

"""
Expand Down Expand Up @@ -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
Expand All @@ -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]
Expand Down Expand Up @@ -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
Expand Down