implement unet and deplot

app.py

@@ -0,0 +1,30 @@
+import streamlit as st
+import numpy as np
+import torch
+import torch.nn as nn
+
+from PIL import Image
+from model import model, image_transforms
+
+
+
+def col_select(value):
+  print(value)
+
+
+st.title("Banan Image Colorizer")
+
+upload_file = st.file_uploader("Upload Image")
+
+if upload_file:
+  image = upload_file
+  image = Image.open(image)
+  image_gs = image_transforms(image)
+  image_gs_prev = image_gs.permute(1, 2, 0).detach().cpu().numpy()
+
+  image_color = model(image_gs.unsqueeze(0)).squeeze().permute(1, 2, 0).detach().cpu().numpy()
+
+  
+  col1, col2 = st.columns(2)
+  col1.image(image_gs_prev)
+  col2.image(image_color, clamp=True, channels='RGB')

banana_colorizer_unet.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:97d8406c9bae3c4ccda6962d483d23e1e649e46a8bbe25e1e2bef95e5abc13b3
+size 124265610

model.py

@@ -0,0 +1,103 @@
+import torch 
+import torch.nn as nn
+import torchvision
+
+import torch.nn.functional as F
+
+
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+image_transforms = torchvision.transforms.Compose([
+    torchvision.transforms.Resize((256, 256)),
+    torchvision.transforms.Grayscale(),
+    torchvision.transforms.ToTensor(),
+    torchvision.transforms.Normalize(mean=[0.0], std=[1.0])  
+])
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_channel, out_channel):
+        super(ConvBlock, self).__init__()
+        self.main = nn.Sequential(
+            nn.Conv2d(in_channel, out_channel, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(out_channel),
+            nn.ReLU(True),
+            nn.Conv2d(out_channel, out_channel, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(out_channel),
+            nn.ReLU(True)
+        )
+        
+    def forward(self, x):
+        return self.main(x)
+
+class UNETFruitColor(nn.Module):
+    def __init__(self):
+        super(UNETFruitColor, self).__init__()
+        
+        self.convs = [64, 128, 256, 512]
+        self.convEncoder = nn.ModuleList()
+        
+        in_feature = 1
+        for conv in self.convs:            
+            self.convEncoder.append(ConvBlock(in_feature, conv))
+            in_feature = conv
+        
+        self.bottleNeck = ConvBlock(self.convs[-1], self.convs[-1]*2)
+        
+        in_feature = self.convs[-1]*2
+        
+        self.convDecoder = nn.ModuleList()
+        self.decoderUpConvs = nn.ModuleList()
+        
+        for conv in self.convs[::-1]:
+            self.convDecoder.append(ConvBlock(in_feature, conv))
+            self.decoderUpConvs.append(nn.ConvTranspose2d(in_feature, conv, kernel_size=2, stride=2, padding=0))
+            in_feature = conv
+        
+        
+        # final conv and deconv
+        self.finalUpConv = nn.Conv2d(in_feature, 3, (1, 1))
+        self.sigmoid = nn.Sigmoid()
+        
+    def forward(self,x):
+        skip_conns = []
+        for conv in self.convEncoder:
+            # conv
+            x = conv(x)
+            # append for skip conns
+            skip_conns.append(x)
+            # max pool
+            x = F.max_pool2d(x, (2,2), stride=2)
+            
+        x = self.bottleNeck(x)
+
+        skip_conns = skip_conns[::-1]
+                
+        for idx in range(len(self.convDecoder)):
+            # do upsample here
+            upconv = self.decoderUpConvs[idx]
+            deconv = self.convDecoder[idx]
+            skp = skip_conns[idx]
+            
+            # do up conv
+            x = upconv(x)
+            
+            # crop and cat
+            x_cat = torchvision.transforms.Resize((x.shape[2], x.shape[3]))(skp)
+            x = torch.cat([x_cat, x], dim=1)
+
+            # do deconv
+            x = deconv(x)
+        
+        # final 
+
+        x = self.finalUpConv(x)
+#         x = self.sigmoid(x)
+        return x
+      
+      
+      
+model = UNETFruitColor()
+model.load_state_dict(torch.load("banana_colorizer_unet.pth", map_location=device),strict=True)
+model.eval()

requirements.txt

@@ -0,0 +1 @@
+streamlit