Reformat with ruff

Author: qubvel

docs/conf.py

@@ -14,10 +14,9 @@
 # import sys
 # sys.path.insert(0, os.path.abspath('.'))
 
-import os
-import re
 import sys
 import datetime
+import sphinx_rtd_theme
 
 sys.path.append("..")
 
@@ -68,14 +67,11 @@ def get_version():
 # a list of builtin themes.
 #
 
-import sphinx_rtd_theme
-
 html_theme = "sphinx_rtd_theme"
 html_theme_path = [sphinx_rtd_theme.get_html_theme_path()]
 
 # import karma_sphinx_theme
 # html_theme = "karma_sphinx_theme"
-import faculty_sphinx_theme
 
 html_theme = "faculty_sphinx_theme"
 

misc/generate_table.py

@@ -4,19 +4,17 @@
 
 
 WIDTH = 32
-COLUMNS = [
-    "Encoder",
-    "Weights",
-    "Params, M",
-]
+COLUMNS = ["Encoder", "Weights", "Params, M"]
 
 
 def wrap_row(r):
     return "|{}|".format(r)
 
 
 header = "|".join([column.ljust(WIDTH, " ") for column in COLUMNS])
-separator = "|".join(["-" * WIDTH] + [":" + "-" * (WIDTH - 2) + ":"] * (len(COLUMNS) - 1))
+separator = "|".join(
+    ["-" * WIDTH] + [":" + "-" * (WIDTH - 2) + ":"] * (len(COLUMNS) - 1)
+)
 
 print(wrap_row(header))
 print(wrap_row(separator))

misc/generate_table_timm.py

@@ -24,20 +24,29 @@ def make_table(data):
 
     l1 = "+" + "-" * max_len1 + "+" + "-" * max_len2 + "+\n"
     l2 = "+" + "=" * max_len1 + "+" + "=" * max_len2 + "+\n"
-    top = "| " + "Encoder name".ljust(max_len1 - 2) + " | " + "Support dilation".center(max_len2 - 2) + " |\n"
+    top = (
+        "| "
+        + "Encoder name".ljust(max_len1 - 2)
+        + " | "
+        + "Support dilation".center(max_len2 - 2)
+        + " |\n"
+    )
 
     table = l1 + top + l2
 
     for k in sorted(data.keys()):
-        support = "✅".center(max_len2 - 3) if data[k]["has_dilation"] else " ".center(max_len2 - 2)
+        support = (
+            "✅".center(max_len2 - 3)
+            if data[k]["has_dilation"]
+            else " ".center(max_len2 - 2)
+        )
         table += "| " + k.ljust(max_len1 - 2) + " | " + support + " |\n"
         table += l1
 
     return table
 
 
 if __name__ == "__main__":
-
     supported_models = {}
 
     with tqdm(timm.list_models()) as names:

segmentation_models_pytorch/__init__.py

@@ -50,8 +50,7 @@ def create_model(
     except KeyError:
         raise KeyError(
             "Wrong architecture type `{}`. Available options are: {}".format(
-                arch,
-                list(archs_dict.keys()),
+                arch, list(archs_dict.keys())
             )
         )
     return model_class(
@@ -61,3 +60,24 @@ def create_model(
         classes=classes,
         **kwargs,
     )
+
+
+__all__ = [
+    "datasets",
+    "encoders",
+    "decoders",
+    "losses",
+    "metrics",
+    "Unet",
+    "UnetPlusPlus",
+    "MAnet",
+    "Linknet",
+    "FPN",
+    "PSPNet",
+    "DeepLabV3",
+    "DeepLabV3Plus",
+    "PAN",
+    "from_pretrained",
+    "create_model",
+    "__version__",
+]

segmentation_models_pytorch/base/__init__.py

@@ -1,11 +1,13 @@
 from .model import SegmentationModel
 
-from .modules import (
-    Conv2dReLU,
-    Attention,
-)
+from .modules import Conv2dReLU, Attention
 
-from .heads import (
-    SegmentationHead,
-    ClassificationHead,
-)
+from .heads import SegmentationHead, ClassificationHead
+
+__all__ = [
+    "SegmentationModel",
+    "Conv2dReLU",
+    "Attention",
+    "SegmentationHead",
+    "ClassificationHead",
+]

segmentation_models_pytorch/base/heads.py

@@ -3,17 +3,29 @@
 
 
 class SegmentationHead(nn.Sequential):
-    def __init__(self, in_channels, out_channels, kernel_size=3, activation=None, upsampling=1):
-        conv2d = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, padding=kernel_size // 2)
-        upsampling = nn.UpsamplingBilinear2d(scale_factor=upsampling) if upsampling > 1 else nn.Identity()
+    def __init__(
+        self, in_channels, out_channels, kernel_size=3, activation=None, upsampling=1
+    ):
+        conv2d = nn.Conv2d(
+            in_channels, out_channels, kernel_size=kernel_size, padding=kernel_size // 2
+        )
+        upsampling = (
+            nn.UpsamplingBilinear2d(scale_factor=upsampling)
+            if upsampling > 1
+            else nn.Identity()
+        )
         activation = Activation(activation)
         super().__init__(conv2d, upsampling, activation)
 
 
 class ClassificationHead(nn.Sequential):
-    def __init__(self, in_channels, classes, pooling="avg", dropout=0.2, activation=None):
+    def __init__(
+        self, in_channels, classes, pooling="avg", dropout=0.2, activation=None
+    ):
         if pooling not in ("max", "avg"):
-            raise ValueError("Pooling should be one of ('max', 'avg'), got {}.".format(pooling))
+            raise ValueError(
+                "Pooling should be one of ('max', 'avg'), got {}.".format(pooling)
+            )
         pool = nn.AdaptiveAvgPool2d(1) if pooling == "avg" else nn.AdaptiveMaxPool2d(1)
         flatten = nn.Flatten()
         dropout = nn.Dropout(p=dropout, inplace=True) if dropout else nn.Identity()

segmentation_models_pytorch/base/hub_mixin.py

@@ -2,7 +2,12 @@
 from pathlib import Path
 from typing import Optional, Union
 from functools import wraps
-from huggingface_hub import PyTorchModelHubMixin, ModelCard, ModelCardData, hf_hub_download
+from huggingface_hub import (
+    PyTorchModelHubMixin,
+    ModelCard,
+    ModelCardData,
+    hf_hub_download,
+)
 
 
 MODEL_CARD = """
@@ -45,15 +50,17 @@
 
 def _format_parameters(parameters: dict):
     params = {k: v for k, v in parameters.items() if not k.startswith("_")}
-    params = [f'"{k}": {v}' if not isinstance(v, str) else f'"{k}": "{v}"' for k, v in params.items()]
+    params = [
+        f'"{k}": {v}' if not isinstance(v, str) else f'"{k}": "{v}"'
+        for k, v in params.items()
+    ]
     params = ",\n".join([f"    {param}" for param in params])
     params = "{\n" + f"{params}" + "\n}"
     return params
 
 
 class SMPHubMixin(PyTorchModelHubMixin):
     def generate_model_card(self, *args, **kwargs) -> ModelCard:
-
         model_parameters_json = _format_parameters(self._hub_mixin_config)
         directory = self._save_directory if hasattr(self, "_save_directory") else None
         repo_id = self._repo_id if hasattr(self, "_repo_id") else None
@@ -97,8 +104,9 @@ def _del_attrs(self, attrs):
                 delattr(self, f"_{attr}")
 
     @wraps(PyTorchModelHubMixin.save_pretrained)
-    def save_pretrained(self, save_directory: Union[str, Path], *args, **kwargs) -> Optional[str]:
-
+    def save_pretrained(
+        self, save_directory: Union[str, Path], *args, **kwargs
+    ) -> Optional[str]:
         # set additional attributes to be used in generate_model_card
         self._save_directory = save_directory
         self._set_attrs_from_kwargs(["metrics", "dataset"], kwargs)
@@ -132,7 +140,9 @@ def config(self):
 @wraps(PyTorchModelHubMixin.from_pretrained)
 def from_pretrained(pretrained_model_name_or_path: str, *args, **kwargs):
     config_path = hf_hub_download(
-        pretrained_model_name_or_path, filename="config.json", revision=kwargs.get("revision", None)
+        pretrained_model_name_or_path,
+        filename="config.json",
+        revision=kwargs.get("revision", None),
     )
     with open(config_path, "r") as f:
         config = json.load(f)

segmentation_models_pytorch/base/initialization.py

@@ -3,7 +3,6 @@
 
 def initialize_decoder(module):
     for m in module.modules():
-
         if isinstance(m, nn.Conv2d):
             nn.init.kaiming_uniform_(m.weight, mode="fan_in", nonlinearity="relu")
             if m.bias is not None:

segmentation_models_pytorch/base/model.py

@@ -4,23 +4,27 @@
 from .hub_mixin import SMPHubMixin
 
 
-class SegmentationModel(
-    torch.nn.Module,
-    SMPHubMixin,
-):
+class SegmentationModel(torch.nn.Module, SMPHubMixin):
     def initialize(self):
         init.initialize_decoder(self.decoder)
         init.initialize_head(self.segmentation_head)
         if self.classification_head is not None:
             init.initialize_head(self.classification_head)
 
     def check_input_shape(self, x):
-
         h, w = x.shape[-2:]
         output_stride = self.encoder.output_stride
         if h % output_stride != 0 or w % output_stride != 0:
-            new_h = (h // output_stride + 1) * output_stride if h % output_stride != 0 else h
-            new_w = (w // output_stride + 1) * output_stride if w % output_stride != 0 else w
+            new_h = (
+                (h // output_stride + 1) * output_stride
+                if h % output_stride != 0
+                else h
+            )
+            new_w = (
+                (w // output_stride + 1) * output_stride
+                if w % output_stride != 0
+                else w
+            )
             raise RuntimeError(
                 f"Wrong input shape height={h}, width={w}. Expected image height and width "
                 f"divisible by {output_stride}. Consider pad your images to shape ({new_h}, {new_w})."

segmentation_models_pytorch/base/modules.py

@@ -17,7 +17,6 @@ def __init__(
         stride=1,
         use_batchnorm=True,
     ):
-
         if use_batchnorm == "inplace" and InPlaceABN is None:
             raise RuntimeError(
                 "In order to use `use_batchnorm='inplace'` inplace_abn package must be installed. "
@@ -83,7 +82,6 @@ def forward(self, x):
 
 class Activation(nn.Module):
     def __init__(self, name, **params):
-
         super().__init__()
 
         if name is None or name == "identity":

segmentation_models_pytorch/datasets/__init__.py

@@ -1 +1,3 @@
 from .oxford_pet import OxfordPetDataset, SimpleOxfordPetDataset
+
+__all__ = ["OxfordPetDataset", "SimpleOxfordPetDataset"]

segmentation_models_pytorch/datasets/oxford_pet.py

@@ -10,7 +10,6 @@
 
 class OxfordPetDataset(torch.utils.data.Dataset):
     def __init__(self, root, mode="train", transform=None):
-
         assert mode in {"train", "valid", "test"}
 
         self.root = root
@@ -26,7 +25,6 @@ def __len__(self):
         return len(self.filenames)
 
     def __getitem__(self, idx):
-
         filename = self.filenames[idx]
         image_path = os.path.join(self.images_directory, filename + ".jpg")
         mask_path = os.path.join(self.masks_directory, filename + ".png")
@@ -63,7 +61,6 @@ def _read_split(self):
 
     @staticmethod
     def download(root):
-
         # load images
         filepath = os.path.join(root, "images.tar.gz")
         download_url(
@@ -83,13 +80,18 @@ def download(root):
 
 class SimpleOxfordPetDataset(OxfordPetDataset):
     def __getitem__(self, *args, **kwargs):
-
         sample = super().__getitem__(*args, **kwargs)
 
         # resize images
-        image = np.array(Image.fromarray(sample["image"]).resize((256, 256), Image.BILINEAR))
-        mask = np.array(Image.fromarray(sample["mask"]).resize((256, 256), Image.NEAREST))
-        trimap = np.array(Image.fromarray(sample["trimap"]).resize((256, 256), Image.NEAREST))
+        image = np.array(
+            Image.fromarray(sample["image"]).resize((256, 256), Image.BILINEAR)
+        )
+        mask = np.array(
+            Image.fromarray(sample["mask"]).resize((256, 256), Image.NEAREST)
+        )
+        trimap = np.array(
+            Image.fromarray(sample["trimap"]).resize((256, 256), Image.NEAREST)
+        )
 
         # convert to other format HWC -> CHW
         sample["image"] = np.moveaxis(image, -1, 0)

segmentation_models_pytorch/decoders/deeplabv3/__init__.py

@@ -1 +1,3 @@
 from .model import DeepLabV3, DeepLabV3Plus
+
+__all__ = ["DeepLabV3", "DeepLabV3Plus"]

segmentation_models_pytorch/decoders/deeplabv3/decoder.py

@@ -61,14 +61,18 @@ def __init__(
     ):
         super().__init__()
         if output_stride not in {8, 16}:
-            raise ValueError("Output stride should be 8 or 16, got {}.".format(output_stride))
+            raise ValueError(
+                "Output stride should be 8 or 16, got {}.".format(output_stride)
+            )
 
         self.out_channels = out_channels
         self.output_stride = output_stride
 
         self.aspp = nn.Sequential(
             ASPP(encoder_channels[-1], out_channels, atrous_rates, separable=True),
-            SeparableConv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False),
+            SeparableConv2d(
+                out_channels, out_channels, kernel_size=3, padding=1, bias=False
+            ),
             nn.BatchNorm2d(out_channels),
             nn.ReLU(),
         )
@@ -79,7 +83,9 @@ def __init__(
         highres_in_channels = encoder_channels[-4]
         highres_out_channels = 48  # proposed by authors of paper
         self.block1 = nn.Sequential(
-            nn.Conv2d(highres_in_channels, highres_out_channels, kernel_size=1, bias=False),
+            nn.Conv2d(
+                highres_in_channels, highres_out_channels, kernel_size=1, bias=False
+            ),
             nn.BatchNorm2d(highres_out_channels),
             nn.ReLU(),
         )
@@ -210,10 +216,5 @@ def __init__(
             groups=in_channels,
             bias=False,
         )
-        pointwise_conv = nn.Conv2d(
-            in_channels,
-            out_channels,
-            kernel_size=1,
-            bias=bias,
-        )
+        pointwise_conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=bias)
         super().__init__(dephtwise_conv, pointwise_conv)