Skip to content

Fix: Albumentations imports, Remove imgaug and Update torch model loading #351

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Fix: Albumentations imports, Remove imgaug and Update torch model loading #351

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
32d7fd7
Fix albumentations imports, remove imgaug and update torch load
dineshnagumothusr May 31, 2025
aa8f3fd
Fix perspective transform
dineshnagumothusr Jun 2, 2025
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
12 changes: 6 additions & 6 deletions saicinpainting/training/data/aug.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from albumentations import DualIAATransform, to_tuple
import imgaug.augmenters as iaa
from albumentations import DualTransform, Affine, Perspective
from albumentations.core.utils import to_tuple

class IAAAffine2(DualIAATransform):
class IAAAffine2(DualTransform):
"""Place a regular grid of points on the input and randomly move the neighbourhood of these point around
via affine transformations.

Expand Down Expand Up @@ -39,7 +39,7 @@ def __init__(

@property
def processor(self):
return iaa.Affine(
return Affine(
self.scale,
self.translate_percent,
self.translate_px,
Expand All @@ -54,7 +54,7 @@ def get_transform_init_args_names(self):
return ("scale", "translate_percent", "translate_px", "rotate", "shear", "order", "cval", "mode")


class IAAPerspective2(DualIAATransform):
class IAAPerspective2(DualTransform):
"""Perform a random four point perspective transform of the input.

Note: This class introduce interpolation artifacts to mask if it has values other than {0;1}
Expand All @@ -78,7 +78,7 @@ def __init__(self, scale=(0.05, 0.1), keep_size=True, always_apply=False, p=0.5,

@property
def processor(self):
return iaa.PerspectiveTransform(self.scale, keep_size=self.keep_size, mode=self.mode, cval=self.cval)
return Perspective(self.scale, keep_size=self.keep_size, mode=self.mode, cval=self.cval)

def get_transform_init_args_names(self):
return ("scale", "keep_size")
2 changes: 1 addition & 1 deletion saicinpainting/training/trainers/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ def make_training_model(config):

def load_checkpoint(train_config, path, map_location='cuda', strict=True):
model: torch.nn.Module = make_training_model(train_config)
state = torch.load(path, map_location=map_location)
state = torch.load(path, map_location=map_location, weights_only=False)
model.load_state_dict(state['state_dict'], strict=strict)
model.on_load_checkpoint(state)
return model