Upload 2 files

app.py

@@ -0,0 +1,93 @@
+import streamlit as st
+from PIL import Image
+import cv2
+import numpy as np
+import time
+import models
+import torch
+
+from torchvision import transforms
+from torchvision import transforms
+
+def load_model(path, model):
+    model.load_state_dict(torch.load(path, map_location=torch.device('cpu')))
+    return model
+
+def predict(img):
+    model = models.unet(3, 1)
+    model = load_model('model.pth',model)
+
+    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])
+    img = cv2.resize(img, (512, 512))
+    convert_tensor = transforms.ToTensor()
+    img =  convert_tensor(img).float()
+    img = normalize(img)
+    img = torch.unsqueeze(img, dim=0)
+
+    output = model(img)
+    result = torch.sigmoid(output)
+
+    threshold = 0.5
+    result = (result >= threshold).float()
+    prediction = result[0].cpu()  # Move tensor to CPU if it's on GPU
+    # Convert tensor to a numpy array
+    prediction_array = prediction.numpy()
+    # Rescale values to the range [0, 255]
+    prediction_array = (prediction_array * 255).astype('uint8').transpose(1, 2, 0)
+    cv2.imwrite("test.png",prediction_array)
+    return prediction_array
+
+def predicjt(img):
+    model1 = models.SAunet(3, 1)
+    model1 = load_model('saunet.pth',model1)
+
+    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])
+    img = cv2.resize(img, (512, 512))
+    convert_tensor = transforms.ToTensor()
+    img =  convert_tensor(img).float()
+    img = normalize(img)
+    img = torch.unsqueeze(img, dim=0)
+
+    output = model1(img)
+    result = torch.sigmoid(output)
+
+    threshold = 0.5
+    result = (result >= threshold).float()
+    prediction = result[0].cpu()  # Move tensor to CPU if it's on GPU
+    # Convert tensor to a numpy array
+    prediction_array = prediction.numpy()
+    # Rescale values to the range [0, 255]
+    prediction_array = (prediction_array * 255).astype('uint8').transpose(1, 2, 0)
+    cv2.imwrite("test1.png",prediction_array)
+    return prediction_array
+def main():
+    st.title("Image Segmentation Demo")
+
+    # Predefined list of image names
+    image_names = ["01_test.tif", "02_test.tif", "03_test.tif"]
+
+    # Create a selection box for the images
+    selected_image_name = st.selectbox("Select an Image", image_names)
+
+    # Load the selected image
+    selected_image = cv2.imread(selected_image_name)
+
+    # Display the selected image
+    st.image(selected_image, channels="RGB")
+
+    # Create a button for segmentation
+    if st.button("Segment"):
+        # Perform segmentation on the selected image
+        segmented_image = predict(selected_image)
+        segmented_image1 = predicjt(selected_image)
+
+
+        # Display the segmented image
+        st.image(segmented_image, channels="RGB",caption='U-Net segmentation')
+        st.image(segmented_image1, channels="RGB",caption='Spatial Attention U-Net segmentation ')
+
+# Function to perform segmentation on the selected image
+
+
+if __name__ == "__main__":
+    main()

models.py

@@ -0,0 +1,210 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import  Tensor
+
+
+class DropBlock(nn.Module):
+    def __init__(self, block_size: int = 5, p: float = 0.1):
+        super().__init__()
+        self.block_size = block_size
+        self.p = p
+
+    def calculate_gamma(self, x: Tensor) -> float:
+        
+
+        invalid = (1 - self.p) / (self.block_size ** 2)
+        valid = (x.shape[-1] ** 2) / ((x.shape[-1] - self.block_size + 1) ** 2)
+        return invalid * valid
+
+    def forward(self, x: Tensor) -> Tensor:
+        N, C, H, W = x.size()
+        if self.training:
+            gamma = self.calculate_gamma(x)
+            mask_shape = (N, C, H - self.block_size + 1, W - self.block_size + 1)
+            mask = torch.bernoulli(torch.full(mask_shape, gamma, device=x.device))
+            mask = F.pad(mask, [self.block_size // 2] * 4, value=0)
+            mask_block = 1 - F.max_pool2d(
+                mask,
+                kernel_size=(self.block_size, self.block_size),
+                stride=(1, 1),
+                padding=(self.block_size // 2, self.block_size // 2),
+            )
+            x = mask_block * x * (mask_block.numel() / mask_block.sum())
+        return x
+
+
+class double_conv(nn.Module):
+  def __init__(self,intc,outc):
+    super().__init__()
+    self.conv1=nn.Conv2d(intc,outc,kernel_size=3,padding=1,stride=1)
+    self.drop1= DropBlock(7, 0.9)
+    self.bn1=nn.BatchNorm2d(outc)
+    self.relu1=nn.ReLU()
+    self.conv2=nn.Conv2d(outc,outc,kernel_size=3,padding=1,stride=1)
+    self.drop2=DropBlock(7, 0.9)
+    self.bn2=nn.BatchNorm2d(outc)
+    self.relu2=nn.ReLU()
+
+  def forward(self,input):
+    x=self.relu1(self.bn1(self.drop1(self.conv1(input))))
+    x=self.relu2(self.bn2(self.drop2(self.conv2(x))))
+
+    return x
+class upconv(nn.Module):
+  def __init__(self,intc,outc) -> None:
+    super().__init__()
+    self.up=nn.ConvTranspose2d(intc, outc, kernel_size=2, stride=2, padding=0)
+   # self.relu=nn.ReLU()
+
+  def forward(self,input):
+    x=self.up(input)
+    #x=self.relu(self.up(input))
+    return x
+class unet(nn.Module):
+  def __init__(self,int,out) -> None:
+    'int: represent the number of image channels'
+    'out: number of the desired final channels'
+
+    super().__init__()
+    'encoder'
+    self.convlayer1=double_conv(int,64)
+    self.down1=nn.MaxPool2d((2, 2))
+    self.convlayer2=double_conv(64,128)
+    self.down2=nn.MaxPool2d((2, 2))
+    self.convlayer3=double_conv(128,256)
+    self.down3=nn.MaxPool2d((2, 2))
+    self.convlayer4=double_conv(256,512)
+    self.down4=nn.MaxPool2d((2, 2))
+
+    'bridge'
+    self.bridge=double_conv(512,1024)
+    'decoder'
+    self.up1=upconv(1024,512)
+    self.convlayer5=double_conv(1024,512)
+    self.up2=upconv(512,256)
+    self.convlayer6=double_conv(512,256)
+    self.up3=upconv(256,128)
+    self.convlayer7=double_conv(256,128)
+    self.up4=upconv(128,64)
+    self.convlayer8=double_conv(128,64)
+    'output'
+    self.outputs = nn.Conv2d(64, out, kernel_size=1, padding=0)
+    self.sig=nn.Sigmoid()
+  def forward(self,input):
+    'encoder'
+    l1=self.convlayer1(input)
+    d1=self.down1(l1)
+    l2=self.convlayer2(d1)
+    d2=self.down2(l2)
+    l3=self.convlayer3(d2)
+    d3=self.down3(l3)
+    l4=self.convlayer4(d3)
+    d4=self.down4(l4)
+    'bridge'
+    bridge=self.bridge(d4)
+    'decoder'
+    up1=self.up1(bridge)
+    up1 = torch.cat([up1, l4], axis=1)
+    l5=self.convlayer5(up1)
+
+    up2=self.up2(l5)
+    up2 = torch.cat([up2, l3], axis=1)
+    l6=self.convlayer6(up2)
+
+    up3=self.up3(l6)
+    up3= torch.cat([up3, l2], axis=1)
+    l7=self.convlayer7(up3)
+
+    up4=self.up4(l7)
+    up4 = torch.cat([up4, l1], axis=1)
+    l8=self.convlayer8(up4)
+    out=self.outputs(l8)
+
+    #out=self.sig(self.outputs(l8))
+    return out
+class spatialAttention(nn.Module):
+  def __init__(self) -> None:
+    super().__init__()
+
+    self.conv77=nn.Conv2d(2,1,kernel_size=7,padding=3)
+    self.sig=nn.Sigmoid()
+  def forward(self,input):
+    x=torch.cat( (torch.max(input,1)[0].unsqueeze(1), torch.mean(input,1).unsqueeze(1)), dim=1 )
+    x=self.sig(self.conv77(x))
+    x=input*x
+    return x
+class SAunet(nn.Module):
+  def __init__(self,int,out) -> None:
+    'int: represent the number of image channels'
+    'out: number of the desired final channels'
+
+    super().__init__()
+    'encoder'
+    self.convlayer1=double_conv(int,64)
+    self.down1=nn.MaxPool2d((2, 2))
+    self.convlayer2=double_conv(64,128)
+    self.down2=nn.MaxPool2d((2, 2))
+    self.convlayer3=double_conv(128,256)
+    self.down3=nn.MaxPool2d((2, 2))
+    self.convlayer4=double_conv(256,512)
+    self.down4=nn.MaxPool2d((2, 2))
+
+    'bridge'
+    self.attmodule=spatialAttention()
+    self.bridge1=nn.Conv2d(512,1024,kernel_size=3,stride=1,padding=1)
+    self.bn1=nn.BatchNorm2d(1024)
+    self.act1=nn.ReLU()
+    self.bridge2=nn.Conv2d(1024,1024,kernel_size=3,stride=1,padding=1)
+    self.bn2=nn.BatchNorm2d(1024)
+    self.act2=nn.ReLU()
+    'decoder'
+    self.up1=upconv(1024,512)
+    self.convlayer5=double_conv(1024,512)
+    self.up2=upconv(512,256)
+    self.convlayer6=double_conv(512,256)
+    self.up3=upconv(256,128)
+    self.convlayer7=double_conv(256,128)
+    self.up4=upconv(128,64)
+    self.convlayer8=double_conv(128,64)
+    'output'
+    self.outputs = nn.Conv2d(64, out, kernel_size=1, padding=0)
+    self.sig=nn.Sigmoid()
+  def forward(self,input):
+    'encoder'
+    l1=self.convlayer1(input)
+    d1=self.down1(l1)
+    l2=self.convlayer2(d1)
+    d2=self.down2(l2)
+    l3=self.convlayer3(d2)
+    d3=self.down3(l3)
+    l4=self.convlayer4(d3)
+    d4=self.down4(l4)
+    'bridge'
+    b1=self.act1(self.bn1(self.bridge1(d4)))
+    att=self.attmodule(b1)
+    b2=self.act2(self.bn2(self.bridge2(att)))
+    'decoder'
+    up1=self.up1(b2)
+    up1 = torch.cat([up1, l4], axis=1)
+    l5=self.convlayer5(up1)
+
+    up2=self.up2(l5)
+    up2 = torch.cat([up2, l3], axis=1)
+    l6=self.convlayer6(up2)
+
+    up3=self.up3(l6)
+    up3= torch.cat([up3, l2], axis=1)
+    l7=self.convlayer7(up3)
+
+    up4=self.up4(l7)
+    up4 = torch.cat([up4, l1], axis=1)
+    l8=self.convlayer8(up4)
+    out=self.outputs(l8)
+
+    #out=self.sig(self.outputs(l8))
+    return out
+
+
+
+