added model and code

Unet_acc_94.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4ddd029a071966365f75ea31ac19ce929dc888c9df801dcd6831888fd15444ec
+size 124247189

app.py

@@ -0,0 +1,52 @@
+import os
+os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
+
+import torch
+import torch.nn as nn
+import cv2
+import numpy as np
+import torch
+import torch.nn as nn
+import torchvision.transforms.functional as TF
+from torchvision import transforms
+
+from model import DoubleConv,UNET
+    
+convert_tensor = transforms.ToTensor()
+device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = UNET(in_channels=3, out_channels=1).to(device)
+# model=torch.load("Unet_acc_94.pth",map_location=torch.device('cpu'))
+
+model=torch.load("Unet_acc_94.pth",map_location=device)
+
+def predict(img):
+    img=cv2.resize(img,(240,160))
+    test_img=convert_tensor(img).unsqueeze(0)
+    print(test_img.shape)
+    preds=model(test_img.float())
+    preds=torch.sigmoid(preds)
+    preds=(preds > 0.5).float()
+    print(preds.shape)
+    im=preds.squeeze(0).permute(1,2,0).detach()
+    print(im.shape)
+    im=im.numpy()
+    return im
+
+import streamlit as st
+st.title("Image Colorizer")
+        
+file=st.file_uploader("Please upload the B/W image",type=["jpg","jpeg","png"])
+print(file)
+if file is None:
+  st.text("Please Upload an image")
+else:
+    file_bytes = np.asarray(bytearray(file.read()), dtype=np.uint8)
+    opencv_image = cv2.imdecode(file_bytes, 1)
+    im=predict(opencv_image)
+    st.text("Original")
+    st.image(file)
+    st.text("Colorized!!")
+    st.image(im)
+
+
+

dataset.py

@@ -0,0 +1,36 @@
+import torch
+from torch.utils.data.dataloader import DataLoader,Dataset
+import torch.optim as optim
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+
+import numpy as np
+import matplotlib.pyplot as plt
+import os
+from PIL import Image
+
+class Segmentation_Dataset(Dataset):
+    def __init__(self,img_dir,mask_dir,transform=None):
+        self.img_dir=img_dir
+        self.mask_dir=mask_dir
+        self.transform=transform
+        self.images=os.listdir(img_dir)
+        self.images=[im for im in self.images if ".jpg" in im]
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self,idx):
+        img_path=os.path.join(self.img_dir,self.images[idx])
+        mask_path=os.path.join(self.mask_dir,self.images[idx].replace(".jpg",".png"))
+
+        image=np.array(Image.open(img_path).convert("RGB"))
+        mask=np.array(Image.open(mask_path).convert("L"),dtype=np.float32)
+        mask[mask==255]=1.0
+
+        if self.transform is not None:
+            augmentations=self.transform(image=image,mask=mask)
+            image=augmentations["image"]
+            mask=augmentations["mask"]
+
+        return image, mask
+        

model.py

@@ -0,0 +1,68 @@
+import torch
+import torch.nn as nn
+import torchvision.transforms.functional as TF
+
+class DoubleConv(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(DoubleConv, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, 3, 1, 1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, 3, 1, 1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+        )
+
+    def forward(self, x):
+        return self.conv(x)
+
+class UNET(nn.Module):
+    def __init__(
+            self, in_channels=3, out_channels=1, features=[64, 128, 256, 512],
+    ):
+        super(UNET, self).__init__()
+        self.ups = nn.ModuleList()
+        self.downs = nn.ModuleList()
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+
+        # Down part of UNET
+        for feature in features:
+            self.downs.append(DoubleConv(in_channels, feature))
+            in_channels = feature
+
+        # Up part of UNET
+        for feature in reversed(features):
+            self.ups.append(
+                nn.ConvTranspose2d(
+                    feature*2, feature, kernel_size=2, stride=2,
+                )
+            )
+            self.ups.append(DoubleConv(feature*2, feature))
+
+        self.bottleneck = DoubleConv(features[-1], features[-1]*2)
+        self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1)
+
+    def forward(self, x):
+        skip_connections = []
+
+        for down in self.downs:
+            x = down(x)
+            skip_connections.append(x)
+            x = self.pool(x)
+
+        x = self.bottleneck(x)
+        skip_connections = skip_connections[::-1]
+
+        for idx in range(0, len(self.ups), 2):
+            x = self.ups[idx](x)
+            skip_connection = skip_connections[idx//2]
+
+            if x.shape != skip_connection.shape:
+                x = TF.resize(x, size=skip_connection.shape[2:])
+
+            concat_skip = torch.cat((skip_connection, x), dim=1)
+            x = self.ups[idx+1](concat_skip)
+
+        return self.final_conv(x)
+

predict.py

@@ -0,0 +1,38 @@
+import torch
+import torchvision
+from torchvision import transforms
+
+import numpy as np
+import matplotlib.pyplot as plt
+from PIL import Image
+from model import DoubleConv,UNET
+
+import os
+os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
+
+
+convert_tensor = transforms.ToTensor()
+device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# print(device)
+
+model = UNET(in_channels=3, out_channels=1).to(device)
+model=torch.load("Unet_acc_94.pth",map_location=torch.device('cpu'))
+
+# test_img=np.array(Image.open("profilepic - Copy.jpeg").resize((160,240)))
+test_img=Image.open("104.jpg").resize((240,160))
+
+# test_img=torch.tensor(test_img).permute(2,1,0)
+# test_img=test_img.unsqueeze(0)
+test_img=convert_tensor(test_img).unsqueeze(0)
+print(test_img.shape)
+preds=model(test_img.float())
+preds=torch.sigmoid(preds)
+preds=(preds > 0.5).float()
+print(preds.shape)
+im=preds.squeeze(0).permute(1,2,0).detach()
+print(im.shape)
+fig,axs=plt.subplots(1,2)
+
+axs[0].imshow(im)
+axs[1].imshow(test_img.squeeze(0).permute(1,2,0).detach())
+plt.show()

tempCodeRunnerFile.py

@@ -0,0 +1,38 @@
+import torch
+import torchvision
+from torchvision import transforms
+
+import numpy as np
+import matplotlib.pyplot as plt
+from PIL import Image
+from model import DoubleConv,UNET
+
+import os
+os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
+
+
+convert_tensor = transforms.ToTensor()
+device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# print(device)
+
+model = UNET(in_channels=3, out_channels=1).to(device)
+model=torch.load("Unet_acc_94.pth",map_location=torch.device('cpu'))
+
+# test_img=np.array(Image.open("profilepic - Copy.jpeg").resize((160,240)))
+test_img=Image.open("104.jpg").resize((240,160))
+
+# test_img=torch.tensor(test_img).permute(2,1,0)
+# test_img=test_img.unsqueeze(0)
+test_img=convert_tensor(test_img).unsqueeze(0)
+print(test_img.shape)
+preds=model(test_img.float())
+preds=torch.sigmoid(preds)
+preds=(preds > 0.5).float()
+print(preds.shape)
+im=preds.squeeze(0).permute(1,2,0).detach()
+print(im.shape)
+fig,axs=plt.subplots(1,2)
+
+axs[0].imshow(im)
+axs[1].imshow(test_img.squeeze(0).permute(1,2,0).detach())
+plt.show()