update files

.scivision/model.yml

@@ -1 +1,8 @@
-
+name: catdog_segmentation_model
+url: https://github.com/pwochner/catdog_segmentation_model
+import: catdog_segmentation_model
+model: CatDogUNet
+prediction_fn:
+    call: predict
+    args:
+        X: image

LICENSE

@@ -1 +1,21 @@
-TODO: Please add the appropriate license for this dataset.
+MIT License
+
+Copyright (c) 2022 pwochner
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

catdog_segmentation_model/__init__.py

@@ -1 +1 @@
-
+from .model import CatDogUNet

catdog_segmentation_model/model.py

@@ -1 +1,39 @@
+import os
+import requests
+import torch
+from torchvision import transforms
+import numpy as np
+import unet
 
+
+class CatDogUNet:
+    def __init__(self):
+        filename = "unet_model.ckpt"
+        if not os.path.exists(filename):
+            model_path = os.path.join(
+                "https://connectionsworkshop.blob.core.windows.net/pets", filename
+            )
+            r = requests.get(model_path)
+            with open(filename, "wb") as outfile:
+                outfile.write(r.content)
+        self.model = unet.CatDogUNet(num_classes=1)
+        self.model.load_state_dict(torch.load(filename))
+        self.model.eval()
+
+    def predict(self, image):
+        # transform input image (as required by model)
+        transform_input = transforms.Compose([transforms.Resize((192, 192)),])
+        image = image.values
+        image = image[:, :, 0:3]  # make sure we have only 3 channels
+        image = np.transpose(image, (2, 0, 1))
+        image = image / 255
+        image = torch.from_numpy(image).type(torch.float32)
+        image = transform_input(image)
+
+        # make prediction
+        prediction = self.model(image)
+        return prediction
+
+
+#  -------------
+# - We need to return an image with class labels

catdog_segmentation_model/unet.py

@@ -1 +1,139 @@
+import torch
+import torch.nn as nn
+
+
+class Block(nn.Module):
+    """
+    Class for the basic convolutional building block of the unet
+    """
+
+    def __init__(self, in_ch, out_ch):
+        """
+        Constructor. 
+        :param in_ch: number of input channels to the block
+        :param out_ch: number of output channels of the block 
+        """
+        super().__init__()
+        self.conv1 = nn.Conv2d(
+            in_ch, out_ch, 3, padding=1
+        )  # (in_channels, out_channels, kernel_size, padding)
+        self.relu = nn.ReLU()
+        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)
+
+    def forward(self, x):
+        """
+        Returns the output of a forward pass of the block
+        :param x: the input tensor
+        :return: the output tensor of the block
+        """
+        # a block consists of two convolutional layers
+        # with ReLU activations
+        x = self.conv1(x)
+        x = self.relu(x)
+        x = self.conv2(x)
+        x = self.relu(x)
+        return x
+
+
+class Encoder(nn.Module):
+    """
+    Class for the encoder part of the unet.
+    """
+
+    def __init__(self, chs=(3, 64, 128, 256, 512, 1024)):
+        """
+        Constructor.
+        :param chs: tuple giving the number of input channels of each block in the encoder
+        """
+        super().__init__()
+        self.enc_blocks = nn.ModuleList(
+            [Block(chs[i], chs[i + 1]) for i in range(len(chs) - 1)]
+        )
+        self.pool = nn.MaxPool2d(2)
+
+    def forward(self, x):
+        """
+        Returns the list of the outputs of all the blocks in the encoder
+        :param x: input image tensor
+        """
+        ftrs = []  # a list to store features
+        cnt = 0
+        for block in self.enc_blocks:
+            x = block(x)
+            # save features to concatenate to decoder blocks
+            ftrs.append(x)
+            x = self.pool(x)
+            cnt += 1
+        return ftrs
+
+
+class Decoder(nn.Module):
+    """
+    Class for the decoder part of the unet.
+    """
+
+    def __init__(self, chs=(1024, 512, 256, 128, 64)):
+        """
+        Constructor.
+        :param chs: tuple giving the number of input channels of each block in the decoder
+        """
+        super().__init__()
+        self.chs = chs
+        self.upconvs = nn.ModuleList(
+            [nn.ConvTranspose2d(chs[i], chs[i + 1], 2, 2) for i in range(len(chs) - 1)]
+        )
+        self.dec_blocks = nn.ModuleList(
+            [Block(chs[i], chs[i + 1]) for i in range(len(chs) - 1)]
+        )
+
+    def forward(self, x, encoder_features):
+        """
+        Returns the output of the decoder part of the unet
+        :param x: input tensor to the decoder
+        :param encoder_features: list of the encoder features to be concatenated to the corresponding level of the decoder
+        """
+        for i in range(len(self.chs) - 1):
+            # logging.warning("Decoder level: %d", i)
+            x = self.upconvs[i](x)
+            # logging.warning(f"Dimensions x:{x.shape}")
+            # get the features from the corresponding level of the encoder
+            enc_ftrs = encoder_features[i]
+            # logging.warning(f"Dimensions encoder features:{enc_ftrs.shape}")
+            # concatenate these features to x
+            x = torch.cat([x, enc_ftrs], dim=1)
+            x = self.dec_blocks[i](x)
+        return x
+
+
+class UNet(nn.Module):
+    """
+    Class for the unet
+    """
+
+    def __init__(
+        self,
+        enc_chs=(3, 64, 128, 256, 512, 1024),
+        dec_chs=(1024, 512, 256, 128, 64),
+        num_classes=1,
+    ):
+        """
+        Constructor.
+        :param enc_chs: tuple giving the number of input channels of each block in the encoder
+        :param dec_chs: tuple giving the number of input channels of each block in the encoder
+        :param num_classes: number of output classes of the segmentation
+        """
+        super().__init__()
+        self.encoder = Encoder(enc_chs)
+        self.decoder = Decoder(dec_chs)
+        self.head = nn.Conv2d(dec_chs[-1], num_classes, 1)  # output layer
+
+    def forward(self, x):
+        """
+        Returns the output of a forward pass of the unet
+        :param x: the input tensor to the unet
+        """
+        enc_ftrs = self.encoder(x)
+        out = self.decoder(enc_ftrs[::-1][0], enc_ftrs[::-1][1:])
+        out = self.head(out)
+        return out
 

setup.py

@@ -9,10 +9,10 @@
             requirements.append(stripped)
 
 setup(
-    name="catdog_classification_model",
+    name="catdog_segmentation_model",
     version="0.0.1",
-    description="scivision plugin, using a cat/dog classifier CNN model",
-    url="TODO: INSERT URL OF THE MODEL GITHUB REPO HERE",
+    description="scivision plugin, using a UNet to segment cat/dog images",
+    url="https://github.com/pwochner/catdog_segmentation_model",
     packages=find_packages(),
     install_requires=requirements,
     python_requires=">=3.7",