Upload 4 files

app.py

@@ -0,0 +1,69 @@
+import torch
+
+import numpy as np
+import gradio as gr
+from PIL import Image
+from unet import UNet
+from torchvision import transforms
+from torchvision.transforms.functional import to_tensor, to_pil_image
+import matplotlib.pyplot as plt
+
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
+device = torch.device(device)
+# Load the trained model
+model_path = 'cityscapes_dataUNet.pth'
+num_classes = 10
+model = UNet(num_classes=num_classes)
+model.load_state_dict(torch.load(model_path))
+model.to(device)
+model.eval()
+
+
+
+
+# Define the prediction function that takes an input image and returns the segmented image
+def predict_segmentation(image):
+    print(device)
+    # Convert the input image to a PyTorch tensor and normalize it
+    image = Image.fromarray(image, 'RGB')
+    # image = transforms.functional.resize(image, (256, 256))
+    image = to_tensor(image).unsqueeze(0)
+    image = transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))(image)
+    image=image.to(device)
+    
+    print("input shape",image.shape) # input shape torch.Size([1, 3, 256, 256])
+    print("input dtype",image.dtype) # input dtype torch.float32
+
+    # Make a prediction using the model
+    with torch.no_grad():
+        
+        print(image.shape, image.dtype) # torch.Size([1, 3, 256, 256]) torch.float32
+
+        output= model(image)
+        # print(output.shape,output.dtype) # torch.Size([1, 10, 256, 256]) torch.float32
+
+        predicted_class = torch.argmax(output, dim=1).squeeze(0)                                                                                                                                                                                                                                                                                                                                                                                                                        
+        predicted_class = predicted_class.cpu().detach().numpy().astype(np.uint8)
+        print(predicted_class.dtype , predicted_class.shape) #  int64 (256, 256)
+
+
+# Visualize the predicted segmentation mask
+    plt.imshow(predicted_class)
+    plt.show()                                                                                                                                                                                                                                                                                                                                                                   
+    # Apply the inverse transform to convert the normalized image back to RGB
+    # predicted_class = inverse_transform(torch.from_numpy(predicted_class))
+
+    print("predicted class ",predicted_class)
+    
+    predicted_class = to_pil_image(predicted_class)
+
+    # Return the predicted segmentation                                                                                                                                                                                                         
+    return predicted_class
+
+# Define the Gradio interface
+input_image = gr.inputs.Image()
+output_image = gr.outputs.Image(type='numpy')
+
+gr.Interface(fn=predict_segmentation, inputs=input_image, outputs=output_image, 
+title='UNet Image Segmentation', 
+description='Segment an image using a UNet model').launch()

cityscapes_dataUNet.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9f41259bc794efd66defb8c0029ea054dcf4bb0b98dcc3229e36267782132315
+size 138216745

cityscapes_dataUnet.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2a9d4dec513aa5ef4a2873cfc1eba34fb5e7a15b0dad50041043ff80d23b2b30
+size 138216745

unet.py

@@ -0,0 +1,57 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import transforms
+class UNet(nn.Module):
+    
+    def __init__(self, num_classes):
+        super(UNet, self).__init__()
+        self.num_classes = num_classes
+        self.contracting_11 = self.conv_block(in_channels=3, out_channels=64)
+        self.contracting_12 = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.contracting_21 = self.conv_block(in_channels=64, out_channels=128)
+        self.contracting_22 = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.contracting_31 = self.conv_block(in_channels=128, out_channels=256)
+        self.contracting_32 = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.contracting_41 = self.conv_block(in_channels=256, out_channels=512)
+        self.contracting_42 = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.middle = self.conv_block(in_channels=512, out_channels=1024)
+        self.expansive_11 = nn.ConvTranspose2d(in_channels=1024, out_channels=512, kernel_size=3, stride=2, padding=1, output_padding=1)
+        self.expansive_12 = self.conv_block(in_channels=1024, out_channels=512)
+        self.expansive_21 = nn.ConvTranspose2d(in_channels=512, out_channels=256, kernel_size=3, stride=2, padding=1, output_padding=1)
+        self.expansive_22 = self.conv_block(in_channels=512, out_channels=256)
+        self.expansive_31 = nn.ConvTranspose2d(in_channels=256, out_channels=128, kernel_size=3, stride=2, padding=1, output_padding=1)
+        self.expansive_32 = self.conv_block(in_channels=256, out_channels=128)
+        self.expansive_41 = nn.ConvTranspose2d(in_channels=128, out_channels=64, kernel_size=3, stride=2, padding=1, output_padding=1)
+        self.expansive_42 = self.conv_block(in_channels=128, out_channels=64)
+        self.output = nn.Conv2d(in_channels=64, out_channels=num_classes, kernel_size=3, stride=1, padding=1)
+        
+    def conv_block(self, in_channels, out_channels):
+        block = nn.Sequential(nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1),
+                                    nn.ReLU(),
+                                    nn.BatchNorm2d(num_features=out_channels),
+                                    nn.Conv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1),
+                                    nn.ReLU(),
+                                    nn.BatchNorm2d(num_features=out_channels))
+        return block
+    
+    def forward(self, X):
+        contracting_11_out = self.contracting_11(X) # [-1, 64, 256, 256]
+        contracting_12_out = self.contracting_12(contracting_11_out) # [-1, 64, 128, 128]
+        contracting_21_out = self.contracting_21(contracting_12_out) # [-1, 128, 128, 128]
+        contracting_22_out = self.contracting_22(contracting_21_out) # [-1, 128, 64, 64]
+        contracting_31_out = self.contracting_31(contracting_22_out) # [-1, 256, 64, 64]
+        contracting_32_out = self.contracting_32(contracting_31_out) # [-1, 256, 32, 32]
+        contracting_41_out = self.contracting_41(contracting_32_out) # [-1, 512, 32, 32]
+        contracting_42_out = self.contracting_42(contracting_41_out) # [-1, 512, 16, 16]
+        middle_out = self.middle(contracting_42_out) # [-1, 1024, 16, 16]
+        expansive_11_out = self.expansive_11(middle_out) # [-1, 512, 32, 32]
+        expansive_12_out = self.expansive_12(torch.cat((expansive_11_out, contracting_41_out), dim=1)) # [-1, 1024, 32, 32] -> [-1, 512, 32, 32]
+        expansive_21_out = self.expansive_21(expansive_12_out) # [-1, 256, 64, 64]
+        expansive_22_out = self.expansive_22(torch.cat((expansive_21_out, contracting_31_out), dim=1)) # [-1, 512, 64, 64] -> [-1, 256, 64, 64]
+        expansive_31_out = self.expansive_31(expansive_22_out) # [-1, 128, 128, 128]
+        expansive_32_out = self.expansive_32(torch.cat((expansive_31_out, contracting_21_out), dim=1)) # [-1, 256, 128, 128] -> [-1, 128, 128, 128]
+        expansive_41_out = self.expansive_41(expansive_32_out) # [-1, 64, 256, 256]
+        expansive_42_out = self.expansive_42(torch.cat((expansive_41_out, contracting_11_out), dim=1)) # [-1, 128, 256, 256] -> [-1, 64, 256, 256]
+        output_out = self.output(expansive_42_out) # [-1, num_classes, 256, 256]
+        return output_out