Skip to content

Commit

Permalink
[pre-commit.ci] auto fixes from pre-commit.com hooks
Browse files Browse the repository at this point in the history
for more information, see https://pre-commit.ci
  • Loading branch information
pre-commit-ci[bot] committed Sep 14, 2024
1 parent 898364c commit 35af5f1
Showing 1 changed file with 6 additions and 5 deletions.
11 changes: 6 additions & 5 deletions kornia/models/segmentor/segmentation_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,19 +21,20 @@ class SegmentationModels(Module):
classes: Number of classes to predict.
**kwargs: Additional arguments to pass to the model. Detailed arguments can be found at:
https://github.com/qubvel-org/segmentation_models.pytorch/tree/main/segmentation_models_pytorch/decoders
Note:
Only encoder weights are available.
Pretrained weights for the whole model are not available.
"""

def __init__(
self,
model_name: str = "Unet",
encoder_name: str = "resnet34",
encoder_weights: Optional[str] = "imagenet",
in_channels: int = 3,
classes: int = 1,
**kwargs
**kwargs,
) -> None:
super().__init__()
self.preproc_params = smp.encoders.get_preprocessing_params(encoder_name) # type: ignore
Expand All @@ -42,7 +43,7 @@ def __init__(
encoder_weights=encoder_weights,
in_channels=in_channels,
classes=classes,
**kwargs
**kwargs,
)

def preprocessing(self, input: Tensor) -> Tensor:
Expand All @@ -52,7 +53,7 @@ def preprocessing(self, input: Tensor) -> Tensor:
input = kornia.color.rgb_to_bgr(input)
else:
raise ValueError(f"Unsupported input space: {self.preproc_params['input_space']}")

if self.preproc_params["input_range"] is not None:
if input.max() > 1 and self.preproc_params["input_range"][1] == 1:
input = input / 255.0
Expand All @@ -61,7 +62,7 @@ def preprocessing(self, input: Tensor) -> Tensor:
mean = tensor(self.preproc_params["mean"]).to(input.device)
else:
mean = tensor(self.preproc_params["mean"]).to(input.device)

if self.preproc_params["std"] is None:
std = tensor(self.preproc_params["std"]).to(input.device)
else:
Expand Down

0 comments on commit 35af5f1

Please sign in to comment.