We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Add support for weights paths to the SEMANTIC_SEGMENTATION_BACKBONES registry


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `flash/image/segmentation/heads.py`
Content:
```
1 # Copyright The PyTorch Lightning team.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 from functools import partial
15 from typing import Callable
16 
17 from flash.core.registry import FlashRegistry
18 from flash.core.utilities.imports import _SEGMENTATION_MODELS_AVAILABLE
19 
20 if _SEGMENTATION_MODELS_AVAILABLE:
21     import segmentation_models_pytorch as smp
22 
23     SMP_MODEL_CLASS = [
24         smp.Unet, smp.UnetPlusPlus, smp.MAnet, smp.Linknet, smp.FPN, smp.PSPNet, smp.DeepLabV3, smp.DeepLabV3Plus,
25         smp.PAN
26     ]
27     SMP_MODELS = {a.__name__.lower(): a for a in SMP_MODEL_CLASS}
28 
29 SEMANTIC_SEGMENTATION_HEADS = FlashRegistry("backbones")
30 
31 if _SEGMENTATION_MODELS_AVAILABLE:
32 
33     def _load_smp_head(
34         head: str,
35         backbone: str,
36         pretrained: bool = True,
37         num_classes: int = 1,
38         in_channels: int = 3,
39         **kwargs,
40     ) -> Callable:
41 
42         if head not in SMP_MODELS:
43             raise NotImplementedError(f"{head} is not implemented! Supported heads -> {SMP_MODELS.keys()}")
44 
45         encoder_weights = None
46         if pretrained:
47             encoder_weights = "imagenet"
48 
49         return smp.create_model(
50             arch=head,
51             encoder_name=backbone,
52             encoder_weights=encoder_weights,
53             classes=num_classes,
54             in_channels=in_channels,
55             **kwargs,
56         )
57 
58     for model_name in SMP_MODELS:
59         SEMANTIC_SEGMENTATION_HEADS(
60             partial(_load_smp_head, head=model_name),
61             name=model_name,
62             namespace="image/segmentation",
63             package="segmentation_models.pytorch"
64         )
65
```

Path: `flash/image/segmentation/model.py`
Content:
```
1 # Copyright The PyTorch Lightning team.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Type, Union
15 
16 import torch
17 from torch import nn
18 from torch.nn import functional as F
19 from torchmetrics import IoU, Metric
20 
21 from flash.core.classification import ClassificationTask
22 from flash.core.data.data_source import DefaultDataKeys
23 from flash.core.data.process import Postprocess, Serializer
24 from flash.core.registry import FlashRegistry
25 from flash.core.utilities.imports import _KORNIA_AVAILABLE
26 from flash.image.segmentation.backbones import SEMANTIC_SEGMENTATION_BACKBONES
27 from flash.image.segmentation.heads import SEMANTIC_SEGMENTATION_HEADS
28 from flash.image.segmentation.serialization import SegmentationLabels
29 
30 if _KORNIA_AVAILABLE:
31     import kornia as K
32 
33 
34 class SemanticSegmentationPostprocess(Postprocess):
35 
36     def per_sample_transform(self, sample: Any) -> Any:
37         resize = K.geometry.Resize(sample[DefaultDataKeys.METADATA]["size"][-2:], interpolation='bilinear')
38         sample[DefaultDataKeys.PREDS] = resize(torch.stack(sample[DefaultDataKeys.PREDS]))
39         sample[DefaultDataKeys.INPUT] = resize(torch.stack(sample[DefaultDataKeys.INPUT]))
40         return super().per_sample_transform(sample)
41 
42 
43 class SemanticSegmentation(ClassificationTask):
44     """``SemanticSegmentation`` is a :class:`~flash.Task` for semantic segmentation of images. For more details, see
45     :ref:`semantic_segmentation`.
46 
47     Args:
48         num_classes: Number of classes to classify.
49         backbone: A string or model to use to compute image features.
50         backbone_kwargs: Additional arguments for the backbone configuration.
51         head: A string or (model, num_features) tuple to use to compute image features.
52         head_kwargs: Additional arguments for the head configuration.
53         pretrained: Use a pretrained backbone.
54         loss_fn: Loss function for training.
55         optimizer: Optimizer to use for training.
56         metrics: Metrics to compute for training and evaluation. Can either be an metric from the `torchmetrics`
57             package, a custom metric inherenting from `torchmetrics.Metric`, a callable function or a list/dict
58             containing a combination of the aforementioned. In all cases, each metric needs to have the signature
59             `metric(preds,target)` and return a single scalar tensor. Defaults to :class:`torchmetrics.IOU`.
60         learning_rate: Learning rate to use for training.
61         multi_label: Whether the targets are multi-label or not.
62         serializer: The :class:`~flash.core.data.process.Serializer` to use when serializing prediction outputs.
63     """
64 
65     postprocess_cls = SemanticSegmentationPostprocess
66 
67     backbones: FlashRegistry = SEMANTIC_SEGMENTATION_BACKBONES
68 
69     heads: FlashRegistry = SEMANTIC_SEGMENTATION_HEADS
70 
71     required_extras: str = "image"
72 
73     def __init__(
74         self,
75         num_classes: int,
76         backbone: Union[str, nn.Module] = "resnet50",
77         backbone_kwargs: Optional[Dict] = None,
78         head: str = "fpn",
79         head_kwargs: Optional[Dict] = None,
80         pretrained: bool = True,
81         loss_fn: Optional[Callable] = None,
82         optimizer: Type[torch.optim.Optimizer] = torch.optim.AdamW,
83         metrics: Union[Metric, Callable, Mapping, Sequence, None] = None,
84         learning_rate: float = 1e-3,
85         multi_label: bool = False,
86         serializer: Optional[Union[Serializer, Mapping[str, Serializer]]] = None,
87         postprocess: Optional[Postprocess] = None,
88     ) -> None:
89         if metrics is None:
90             metrics = IoU(num_classes=num_classes)
91 
92         if loss_fn is None:
93             loss_fn = F.cross_entropy
94 
95         # TODO: need to check for multi_label
96         if multi_label:
97             raise NotImplementedError("Multi-label not supported yet.")
98 
99         super().__init__(
100             model=None,
101             loss_fn=loss_fn,
102             optimizer=optimizer,
103             metrics=metrics,
104             learning_rate=learning_rate,
105             serializer=serializer or SegmentationLabels(),
106             postprocess=postprocess or self.postprocess_cls()
107         )
108 
109         self.save_hyperparameters()
110 
111         if not backbone_kwargs:
112             backbone_kwargs = {}
113 
114         if not head_kwargs:
115             head_kwargs = {}
116 
117         if isinstance(backbone, nn.Module):
118             self.backbone = backbone
119         else:
120             self.backbone = self.backbones.get(backbone)(**backbone_kwargs)
121 
122         self.head: nn.Module = self.heads.get(head)(
123             backbone=self.backbone, num_classes=num_classes, pretrained=pretrained, **head_kwargs
124         )
125         self.backbone = self.head.encoder
126 
127     def training_step(self, batch: Any, batch_idx: int) -> Any:
128         batch = (batch[DefaultDataKeys.INPUT], batch[DefaultDataKeys.TARGET])
129         return super().training_step(batch, batch_idx)
130 
131     def validation_step(self, batch: Any, batch_idx: int) -> Any:
132         batch = (batch[DefaultDataKeys.INPUT], batch[DefaultDataKeys.TARGET])
133         return super().validation_step(batch, batch_idx)
134 
135     def test_step(self, batch: Any, batch_idx: int) -> Any:
136         batch = (batch[DefaultDataKeys.INPUT], batch[DefaultDataKeys.TARGET])
137         return super().test_step(batch, batch_idx)
138 
139     def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> Any:
140         batch_input = (batch[DefaultDataKeys.INPUT])
141         batch[DefaultDataKeys.PREDS] = super().predict_step(batch_input, batch_idx, dataloader_idx=dataloader_idx)
142         return batch
143 
144     def forward(self, x) -> torch.Tensor:
145         res = self.head(x)
146 
147         # some frameworks like torchvision return a dict.
148         # In particular, torchvision segmentation models return the output logits
149         # in the key `out`.
150         if torch.jit.isinstance(res, Dict[str, torch.Tensor]):
151             out = res['out']
152         elif torch.is_tensor(res):
153             out = res
154         else:
155             raise NotImplementedError(f"Unsupported output type: {type(res)}")
156 
157         return out
158 
159     @staticmethod
160     def _ci_benchmark_fn(history: List[Dict[str, Any]]):
161         """
162         This function is used only for debugging usage with CI
163         """
164         assert history[-1]["val_iou"] > 0.2
165
```

Path: `flash/image/segmentation/backbones.py`
Content:
```
1 # Copyright The PyTorch Lightning team.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 from functools import partial
15 
16 from flash.core.registry import FlashRegistry
17 from flash.core.utilities.imports import _SEGMENTATION_MODELS_AVAILABLE
18 
19 if _SEGMENTATION_MODELS_AVAILABLE:
20     import segmentation_models_pytorch as smp
21 
22 SEMANTIC_SEGMENTATION_BACKBONES = FlashRegistry("backbones")
23 
24 if _SEGMENTATION_MODELS_AVAILABLE:
25 
26     ENCODERS = smp.encoders.get_encoder_names()
27 
28     def _load_smp_backbone(backbone: str, **_) -> str:
29         return backbone
30 
31     for encoder_name in ENCODERS:
32         short_name = encoder_name
33         if short_name.startswith("timm-"):
34             short_name = encoder_name[5:]
35         SEMANTIC_SEGMENTATION_BACKBONES(
36             partial(_load_smp_backbone, backbone=encoder_name), name=short_name, namespace="image/segmentation"
37         )
38
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/flash/image/segmentation/backbones.py b/flash/image/segmentation/backbones.py
--- a/flash/image/segmentation/backbones.py
+++ b/flash/image/segmentation/backbones.py
@@ -32,6 +32,11 @@
         short_name = encoder_name
         if short_name.startswith("timm-"):
             short_name = encoder_name[5:]
+
+        available_weights = smp.encoders.encoders[encoder_name]["pretrained_settings"].keys()
         SEMANTIC_SEGMENTATION_BACKBONES(
-            partial(_load_smp_backbone, backbone=encoder_name), name=short_name, namespace="image/segmentation"
+            partial(_load_smp_backbone, backbone=encoder_name),
+            name=short_name,
+            namespace="image/segmentation",
+            weights_paths=available_weights,
         )
diff --git a/flash/image/segmentation/heads.py b/flash/image/segmentation/heads.py
--- a/flash/image/segmentation/heads.py
+++ b/flash/image/segmentation/heads.py
@@ -12,7 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from functools import partial
-from typing import Callable
+from typing import Union
+
+from torch import nn
 
 from flash.core.registry import FlashRegistry
 from flash.core.utilities.imports import _SEGMENTATION_MODELS_AVAILABLE
@@ -33,17 +35,19 @@
     def _load_smp_head(
         head: str,
         backbone: str,
-        pretrained: bool = True,
+        pretrained: Union[bool, str] = True,
         num_classes: int = 1,
         in_channels: int = 3,
         **kwargs,
-    ) -> Callable:
+    ) -> nn.Module:
 
         if head not in SMP_MODELS:
             raise NotImplementedError(f"{head} is not implemented! Supported heads -> {SMP_MODELS.keys()}")
 
         encoder_weights = None
-        if pretrained:
+        if isinstance(pretrained, str):
+            encoder_weights = pretrained
+        elif pretrained:
             encoder_weights = "imagenet"
 
         return smp.create_model(
diff --git a/flash/image/segmentation/model.py b/flash/image/segmentation/model.py
--- a/flash/image/segmentation/model.py
+++ b/flash/image/segmentation/model.py
@@ -77,7 +77,7 @@
         backbone_kwargs: Optional[Dict] = None,
         head: str = "fpn",
         head_kwargs: Optional[Dict] = None,
-        pretrained: bool = True,
+        pretrained: Union[bool, str] = True,
         loss_fn: Optional[Callable] = None,
         optimizer: Type[torch.optim.Optimizer] = torch.optim.AdamW,
         metrics: Union[Metric, Callable, Mapping, Sequence, None] = None,
@@ -156,6 +156,16 @@
 
         return out
 
+    @classmethod
+    def available_pretrained_weights(cls, backbone: str):
+        result = cls.backbones.get(backbone, with_metadata=True)
+        pretrained_weights = None
+
+        if "weights_paths" in result["metadata"]:
+            pretrained_weights = list(result["metadata"]["weights_paths"])
+
+        return pretrained_weights
+
     @staticmethod
     def _ci_benchmark_fn(history: List[Dict[str, Any]]):
         """

{"golden_diff": "diff --git a/flash/image/segmentation/backbones.py b/flash/image/segmentation/backbones.py\n--- a/flash/image/segmentation/backbones.py\n+++ b/flash/image/segmentation/backbones.py\n@@ -32,6 +32,11 @@\n         short_name = encoder_name\n         if short_name.startswith(\"timm-\"):\n             short_name = encoder_name[5:]\n+\n+        available_weights = smp.encoders.encoders[encoder_name][\"pretrained_settings\"].keys()\n         SEMANTIC_SEGMENTATION_BACKBONES(\n-            partial(_load_smp_backbone, backbone=encoder_name), name=short_name, namespace=\"image/segmentation\"\n+            partial(_load_smp_backbone, backbone=encoder_name),\n+            name=short_name,\n+            namespace=\"image/segmentation\",\n+            weights_paths=available_weights,\n         )\ndiff --git a/flash/image/segmentation/heads.py b/flash/image/segmentation/heads.py\n--- a/flash/image/segmentation/heads.py\n+++ b/flash/image/segmentation/heads.py\n@@ -12,7 +12,9 @@\n # See the License for the specific language governing permissions and\n # limitations under the License.\n from functools import partial\n-from typing import Callable\n+from typing import Union\n+\n+from torch import nn\n \n from flash.core.registry import FlashRegistry\n from flash.core.utilities.imports import _SEGMENTATION_MODELS_AVAILABLE\n@@ -33,17 +35,19 @@\n     def _load_smp_head(\n         head: str,\n         backbone: str,\n-        pretrained: bool = True,\n+        pretrained: Union[bool, str] = True,\n         num_classes: int = 1,\n         in_channels: int = 3,\n         **kwargs,\n-    ) -> Callable:\n+    ) -> nn.Module:\n \n         if head not in SMP_MODELS:\n             raise NotImplementedError(f\"{head} is not implemented! Supported heads -> {SMP_MODELS.keys()}\")\n \n         encoder_weights = None\n-        if pretrained:\n+        if isinstance(pretrained, str):\n+            encoder_weights = pretrained\n+        elif pretrained:\n             encoder_weights = \"imagenet\"\n \n         return smp.create_model(\ndiff --git a/flash/image/segmentation/model.py b/flash/image/segmentation/model.py\n--- a/flash/image/segmentation/model.py\n+++ b/flash/image/segmentation/model.py\n@@ -77,7 +77,7 @@\n         backbone_kwargs: Optional[Dict] = None,\n         head: str = \"fpn\",\n         head_kwargs: Optional[Dict] = None,\n-        pretrained: bool = True,\n+        pretrained: Union[bool, str] = True,\n         loss_fn: Optional[Callable] = None,\n         optimizer: Type[torch.optim.Optimizer] = torch.optim.AdamW,\n         metrics: Union[Metric, Callable, Mapping, Sequence, None] = None,\n@@ -156,6 +156,16 @@\n \n         return out\n \n+    @classmethod\n+    def available_pretrained_weights(cls, backbone: str):\n+        result = cls.backbones.get(backbone, with_metadata=True)\n+        pretrained_weights = None\n+\n+        if \"weights_paths\" in result[\"metadata\"]:\n+            pretrained_weights = list(result[\"metadata\"][\"weights_paths\"])\n+\n+        return pretrained_weights\n+\n     @staticmethod\n     def _ci_benchmark_fn(history: List[Dict[str, Any]]):\n         \"\"\"\n", "issue": "Add support for weights paths to the SEMANTIC_SEGMENTATION_BACKBONES registry\n\n", "before_files": [{"content": "# Copyright The PyTorch Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom functools import partial\nfrom typing import Callable\n\nfrom flash.core.registry import FlashRegistry\nfrom flash.core.utilities.imports import _SEGMENTATION_MODELS_AVAILABLE\n\nif _SEGMENTATION_MODELS_AVAILABLE:\n    import segmentation_models_pytorch as smp\n\n    SMP_MODEL_CLASS = [\n        smp.Unet, smp.UnetPlusPlus, smp.MAnet, smp.Linknet, smp.FPN, smp.PSPNet, smp.DeepLabV3, smp.DeepLabV3Plus,\n        smp.PAN\n    ]\n    SMP_MODELS = {a.__name__.lower(): a for a in SMP_MODEL_CLASS}\n\nSEMANTIC_SEGMENTATION_HEADS = FlashRegistry(\"backbones\")\n\nif _SEGMENTATION_MODELS_AVAILABLE:\n\n    def _load_smp_head(\n        head: str,\n        backbone: str,\n        pretrained: bool = True,\n        num_classes: int = 1,\n        in_channels: int = 3,\n        **kwargs,\n    ) -> Callable:\n\n        if head not in SMP_MODELS:\n            raise NotImplementedError(f\"{head} is not implemented! Supported heads -> {SMP_MODELS.keys()}\")\n\n        encoder_weights = None\n        if pretrained:\n            encoder_weights = \"imagenet\"\n\n        return smp.create_model(\n            arch=head,\n            encoder_name=backbone,\n            encoder_weights=encoder_weights,\n            classes=num_classes,\n            in_channels=in_channels,\n            **kwargs,\n        )\n\n    for model_name in SMP_MODELS:\n        SEMANTIC_SEGMENTATION_HEADS(\n            partial(_load_smp_head, head=model_name),\n            name=model_name,\n            namespace=\"image/segmentation\",\n            package=\"segmentation_models.pytorch\"\n        )\n", "path": "flash/image/segmentation/heads.py"}, {"content": "# Copyright The PyTorch Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Type, Union\n\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\nfrom torchmetrics import IoU, Metric\n\nfrom flash.core.classification import ClassificationTask\nfrom flash.core.data.data_source import DefaultDataKeys\nfrom flash.core.data.process import Postprocess, Serializer\nfrom flash.core.registry import FlashRegistry\nfrom flash.core.utilities.imports import _KORNIA_AVAILABLE\nfrom flash.image.segmentation.backbones import SEMANTIC_SEGMENTATION_BACKBONES\nfrom flash.image.segmentation.heads import SEMANTIC_SEGMENTATION_HEADS\nfrom flash.image.segmentation.serialization import SegmentationLabels\n\nif _KORNIA_AVAILABLE:\n    import kornia as K\n\n\nclass SemanticSegmentationPostprocess(Postprocess):\n\n    def per_sample_transform(self, sample: Any) -> Any:\n        resize = K.geometry.Resize(sample[DefaultDataKeys.METADATA][\"size\"][-2:], interpolation='bilinear')\n        sample[DefaultDataKeys.PREDS] = resize(torch.stack(sample[DefaultDataKeys.PREDS]))\n        sample[DefaultDataKeys.INPUT] = resize(torch.stack(sample[DefaultDataKeys.INPUT]))\n        return super().per_sample_transform(sample)\n\n\nclass SemanticSegmentation(ClassificationTask):\n    \"\"\"``SemanticSegmentation`` is a :class:`~flash.Task` for semantic segmentation of images. For more details, see\n    :ref:`semantic_segmentation`.\n\n    Args:\n        num_classes: Number of classes to classify.\n        backbone: A string or model to use to compute image features.\n        backbone_kwargs: Additional arguments for the backbone configuration.\n        head: A string or (model, num_features) tuple to use to compute image features.\n        head_kwargs: Additional arguments for the head configuration.\n        pretrained: Use a pretrained backbone.\n        loss_fn: Loss function for training.\n        optimizer: Optimizer to use for training.\n        metrics: Metrics to compute for training and evaluation. Can either be an metric from the `torchmetrics`\n            package, a custom metric inherenting from `torchmetrics.Metric`, a callable function or a list/dict\n            containing a combination of the aforementioned. In all cases, each metric needs to have the signature\n            `metric(preds,target)` and return a single scalar tensor. Defaults to :class:`torchmetrics.IOU`.\n        learning_rate: Learning rate to use for training.\n        multi_label: Whether the targets are multi-label or not.\n        serializer: The :class:`~flash.core.data.process.Serializer` to use when serializing prediction outputs.\n    \"\"\"\n\n    postprocess_cls = SemanticSegmentationPostprocess\n\n    backbones: FlashRegistry = SEMANTIC_SEGMENTATION_BACKBONES\n\n    heads: FlashRegistry = SEMANTIC_SEGMENTATION_HEADS\n\n    required_extras: str = \"image\"\n\n    def __init__(\n        self,\n        num_classes: int,\n        backbone: Union[str, nn.Module] = \"resnet50\",\n        backbone_kwargs: Optional[Dict] = None,\n        head: str = \"fpn\",\n        head_kwargs: Optional[Dict] = None,\n        pretrained: bool = True,\n        loss_fn: Optional[Callable] = None,\n        optimizer: Type[torch.optim.Optimizer] = torch.optim.AdamW,\n        metrics: Union[Metric, Callable, Mapping, Sequence, None] = None,\n        learning_rate: float = 1e-3,\n        multi_label: bool = False,\n        serializer: Optional[Union[Serializer, Mapping[str, Serializer]]] = None,\n        postprocess: Optional[Postprocess] = None,\n    ) -> None:\n        if metrics is None:\n            metrics = IoU(num_classes=num_classes)\n\n        if loss_fn is None:\n            loss_fn = F.cross_entropy\n\n        # TODO: need to check for multi_label\n        if multi_label:\n            raise NotImplementedError(\"Multi-label not supported yet.\")\n\n        super().__init__(\n            model=None,\n            loss_fn=loss_fn,\n            optimizer=optimizer,\n            metrics=metrics,\n            learning_rate=learning_rate,\n            serializer=serializer or SegmentationLabels(),\n            postprocess=postprocess or self.postprocess_cls()\n        )\n\n        self.save_hyperparameters()\n\n        if not backbone_kwargs:\n            backbone_kwargs = {}\n\n        if not head_kwargs:\n            head_kwargs = {}\n\n        if isinstance(backbone, nn.Module):\n            self.backbone = backbone\n        else:\n            self.backbone = self.backbones.get(backbone)(**backbone_kwargs)\n\n        self.head: nn.Module = self.heads.get(head)(\n            backbone=self.backbone, num_classes=num_classes, pretrained=pretrained, **head_kwargs\n        )\n        self.backbone = self.head.encoder\n\n    def training_step(self, batch: Any, batch_idx: int) -> Any:\n        batch = (batch[DefaultDataKeys.INPUT], batch[DefaultDataKeys.TARGET])\n        return super().training_step(batch, batch_idx)\n\n    def validation_step(self, batch: Any, batch_idx: int) -> Any:\n        batch = (batch[DefaultDataKeys.INPUT], batch[DefaultDataKeys.TARGET])\n        return super().validation_step(batch, batch_idx)\n\n    def test_step(self, batch: Any, batch_idx: int) -> Any:\n        batch = (batch[DefaultDataKeys.INPUT], batch[DefaultDataKeys.TARGET])\n        return super().test_step(batch, batch_idx)\n\n    def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> Any:\n        batch_input = (batch[DefaultDataKeys.INPUT])\n        batch[DefaultDataKeys.PREDS] = super().predict_step(batch_input, batch_idx, dataloader_idx=dataloader_idx)\n        return batch\n\n    def forward(self, x) -> torch.Tensor:\n        res = self.head(x)\n\n        # some frameworks like torchvision return a dict.\n        # In particular, torchvision segmentation models return the output logits\n        # in the key `out`.\n        if torch.jit.isinstance(res, Dict[str, torch.Tensor]):\n            out = res['out']\n        elif torch.is_tensor(res):\n            out = res\n        else:\n            raise NotImplementedError(f\"Unsupported output type: {type(res)}\")\n\n        return out\n\n    @staticmethod\n    def _ci_benchmark_fn(history: List[Dict[str, Any]]):\n        \"\"\"\n        This function is used only for debugging usage with CI\n        \"\"\"\n        assert history[-1][\"val_iou\"] > 0.2\n", "path": "flash/image/segmentation/model.py"}, {"content": "# Copyright The PyTorch Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom functools import partial\n\nfrom flash.core.registry import FlashRegistry\nfrom flash.core.utilities.imports import _SEGMENTATION_MODELS_AVAILABLE\n\nif _SEGMENTATION_MODELS_AVAILABLE:\n    import segmentation_models_pytorch as smp\n\nSEMANTIC_SEGMENTATION_BACKBONES = FlashRegistry(\"backbones\")\n\nif _SEGMENTATION_MODELS_AVAILABLE:\n\n    ENCODERS = smp.encoders.get_encoder_names()\n\n    def _load_smp_backbone(backbone: str, **_) -> str:\n        return backbone\n\n    for encoder_name in ENCODERS:\n        short_name = encoder_name\n        if short_name.startswith(\"timm-\"):\n            short_name = encoder_name[5:]\n        SEMANTIC_SEGMENTATION_BACKBONES(\n            partial(_load_smp_backbone, backbone=encoder_name), name=short_name, namespace=\"image/segmentation\"\n        )\n", "path": "flash/image/segmentation/backbones.py"}], "after_files": [{"content": "# Copyright The PyTorch Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom functools import partial\nfrom typing import Union\n\nfrom torch import nn\n\nfrom flash.core.registry import FlashRegistry\nfrom flash.core.utilities.imports import _SEGMENTATION_MODELS_AVAILABLE\n\nif _SEGMENTATION_MODELS_AVAILABLE:\n    import segmentation_models_pytorch as smp\n\n    SMP_MODEL_CLASS = [\n        smp.Unet, smp.UnetPlusPlus, smp.MAnet, smp.Linknet, smp.FPN, smp.PSPNet, smp.DeepLabV3, smp.DeepLabV3Plus,\n        smp.PAN\n    ]\n    SMP_MODELS = {a.__name__.lower(): a for a in SMP_MODEL_CLASS}\n\nSEMANTIC_SEGMENTATION_HEADS = FlashRegistry(\"backbones\")\n\nif _SEGMENTATION_MODELS_AVAILABLE:\n\n    def _load_smp_head(\n        head: str,\n        backbone: str,\n        pretrained: Union[bool, str] = True,\n        num_classes: int = 1,\n        in_channels: int = 3,\n        **kwargs,\n    ) -> nn.Module:\n\n        if head not in SMP_MODELS:\n            raise NotImplementedError(f\"{head} is not implemented! Supported heads -> {SMP_MODELS.keys()}\")\n\n        encoder_weights = None\n        if isinstance(pretrained, str):\n            encoder_weights = pretrained\n        elif pretrained:\n            encoder_weights = \"imagenet\"\n\n        return smp.create_model(\n            arch=head,\n            encoder_name=backbone,\n            encoder_weights=encoder_weights,\n            classes=num_classes,\n            in_channels=in_channels,\n            **kwargs,\n        )\n\n    for model_name in SMP_MODELS:\n        SEMANTIC_SEGMENTATION_HEADS(\n            partial(_load_smp_head, head=model_name),\n            name=model_name,\n            namespace=\"image/segmentation\",\n            package=\"segmentation_models.pytorch\"\n        )\n", "path": "flash/image/segmentation/heads.py"}, {"content": "# Copyright The PyTorch Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Type, Union\n\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\nfrom torchmetrics import IoU, Metric\n\nfrom flash.core.classification import ClassificationTask\nfrom flash.core.data.data_source import DefaultDataKeys\nfrom flash.core.data.process import Postprocess, Serializer\nfrom flash.core.registry import FlashRegistry\nfrom flash.core.utilities.imports import _KORNIA_AVAILABLE\nfrom flash.image.segmentation.backbones import SEMANTIC_SEGMENTATION_BACKBONES\nfrom flash.image.segmentation.heads import SEMANTIC_SEGMENTATION_HEADS\nfrom flash.image.segmentation.serialization import SegmentationLabels\n\nif _KORNIA_AVAILABLE:\n    import kornia as K\n\n\nclass SemanticSegmentationPostprocess(Postprocess):\n\n    def per_sample_transform(self, sample: Any) -> Any:\n        resize = K.geometry.Resize(sample[DefaultDataKeys.METADATA][\"size\"][-2:], interpolation='bilinear')\n        sample[DefaultDataKeys.PREDS] = resize(torch.stack(sample[DefaultDataKeys.PREDS]))\n        sample[DefaultDataKeys.INPUT] = resize(torch.stack(sample[DefaultDataKeys.INPUT]))\n        return super().per_sample_transform(sample)\n\n\nclass SemanticSegmentation(ClassificationTask):\n    \"\"\"``SemanticSegmentation`` is a :class:`~flash.Task` for semantic segmentation of images. For more details, see\n    :ref:`semantic_segmentation`.\n\n    Args:\n        num_classes: Number of classes to classify.\n        backbone: A string or model to use to compute image features.\n        backbone_kwargs: Additional arguments for the backbone configuration.\n        head: A string or (model, num_features) tuple to use to compute image features.\n        head_kwargs: Additional arguments for the head configuration.\n        pretrained: Use a pretrained backbone.\n        loss_fn: Loss function for training.\n        optimizer: Optimizer to use for training.\n        metrics: Metrics to compute for training and evaluation. Can either be an metric from the `torchmetrics`\n            package, a custom metric inherenting from `torchmetrics.Metric`, a callable function or a list/dict\n            containing a combination of the aforementioned. In all cases, each metric needs to have the signature\n            `metric(preds,target)` and return a single scalar tensor. Defaults to :class:`torchmetrics.IOU`.\n        learning_rate: Learning rate to use for training.\n        multi_label: Whether the targets are multi-label or not.\n        serializer: The :class:`~flash.core.data.process.Serializer` to use when serializing prediction outputs.\n    \"\"\"\n\n    postprocess_cls = SemanticSegmentationPostprocess\n\n    backbones: FlashRegistry = SEMANTIC_SEGMENTATION_BACKBONES\n\n    heads: FlashRegistry = SEMANTIC_SEGMENTATION_HEADS\n\n    required_extras: str = \"image\"\n\n    def __init__(\n        self,\n        num_classes: int,\n        backbone: Union[str, nn.Module] = \"resnet50\",\n        backbone_kwargs: Optional[Dict] = None,\n        head: str = \"fpn\",\n        head_kwargs: Optional[Dict] = None,\n        pretrained: Union[bool, str] = True,\n        loss_fn: Optional[Callable] = None,\n        optimizer: Type[torch.optim.Optimizer] = torch.optim.AdamW,\n        metrics: Union[Metric, Callable, Mapping, Sequence, None] = None,\n        learning_rate: float = 1e-3,\n        multi_label: bool = False,\n        serializer: Optional[Union[Serializer, Mapping[str, Serializer]]] = None,\n        postprocess: Optional[Postprocess] = None,\n    ) -> None:\n        if metrics is None:\n            metrics = IoU(num_classes=num_classes)\n\n        if loss_fn is None:\n            loss_fn = F.cross_entropy\n\n        # TODO: need to check for multi_label\n        if multi_label:\n            raise NotImplementedError(\"Multi-label not supported yet.\")\n\n        super().__init__(\n            model=None,\n            loss_fn=loss_fn,\n            optimizer=optimizer,\n            metrics=metrics,\n            learning_rate=learning_rate,\n            serializer=serializer or SegmentationLabels(),\n            postprocess=postprocess or self.postprocess_cls()\n        )\n\n        self.save_hyperparameters()\n\n        if not backbone_kwargs:\n            backbone_kwargs = {}\n\n        if not head_kwargs:\n            head_kwargs = {}\n\n        if isinstance(backbone, nn.Module):\n            self.backbone = backbone\n        else:\n            self.backbone = self.backbones.get(backbone)(**backbone_kwargs)\n\n        self.head: nn.Module = self.heads.get(head)(\n            backbone=self.backbone, num_classes=num_classes, pretrained=pretrained, **head_kwargs\n        )\n        self.backbone = self.head.encoder\n\n    def training_step(self, batch: Any, batch_idx: int) -> Any:\n        batch = (batch[DefaultDataKeys.INPUT], batch[DefaultDataKeys.TARGET])\n        return super().training_step(batch, batch_idx)\n\n    def validation_step(self, batch: Any, batch_idx: int) -> Any:\n        batch = (batch[DefaultDataKeys.INPUT], batch[DefaultDataKeys.TARGET])\n        return super().validation_step(batch, batch_idx)\n\n    def test_step(self, batch: Any, batch_idx: int) -> Any:\n        batch = (batch[DefaultDataKeys.INPUT], batch[DefaultDataKeys.TARGET])\n        return super().test_step(batch, batch_idx)\n\n    def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> Any:\n        batch_input = (batch[DefaultDataKeys.INPUT])\n        batch[DefaultDataKeys.PREDS] = super().predict_step(batch_input, batch_idx, dataloader_idx=dataloader_idx)\n        return batch\n\n    def forward(self, x) -> torch.Tensor:\n        res = self.head(x)\n\n        # some frameworks like torchvision return a dict.\n        # In particular, torchvision segmentation models return the output logits\n        # in the key `out`.\n        if torch.jit.isinstance(res, Dict[str, torch.Tensor]):\n            out = res['out']\n        elif torch.is_tensor(res):\n            out = res\n        else:\n            raise NotImplementedError(f\"Unsupported output type: {type(res)}\")\n\n        return out\n\n    @classmethod\n    def available_pretrained_weights(cls, backbone: str):\n        result = cls.backbones.get(backbone, with_metadata=True)\n        pretrained_weights = None\n\n        if \"weights_paths\" in result[\"metadata\"]:\n            pretrained_weights = list(result[\"metadata\"][\"weights_paths\"])\n\n        return pretrained_weights\n\n    @staticmethod\n    def _ci_benchmark_fn(history: List[Dict[str, Any]]):\n        \"\"\"\n        This function is used only for debugging usage with CI\n        \"\"\"\n        assert history[-1][\"val_iou\"] > 0.2\n", "path": "flash/image/segmentation/model.py"}, {"content": "# Copyright The PyTorch Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom functools import partial\n\nfrom flash.core.registry import FlashRegistry\nfrom flash.core.utilities.imports import _SEGMENTATION_MODELS_AVAILABLE\n\nif _SEGMENTATION_MODELS_AVAILABLE:\n    import segmentation_models_pytorch as smp\n\nSEMANTIC_SEGMENTATION_BACKBONES = FlashRegistry(\"backbones\")\n\nif _SEGMENTATION_MODELS_AVAILABLE:\n\n    ENCODERS = smp.encoders.get_encoder_names()\n\n    def _load_smp_backbone(backbone: str, **_) -> str:\n        return backbone\n\n    for encoder_name in ENCODERS:\n        short_name = encoder_name\n        if short_name.startswith(\"timm-\"):\n            short_name = encoder_name[5:]\n\n        available_weights = smp.encoders.encoders[encoder_name][\"pretrained_settings\"].keys()\n        SEMANTIC_SEGMENTATION_BACKBONES(\n            partial(_load_smp_backbone, backbone=encoder_name),\n            name=short_name,\n            namespace=\"image/segmentation\",\n            weights_paths=available_weights,\n        )\n", "path": "flash/image/segmentation/backbones.py"}]}