loaded new colorizer weights

app.py

@@ -4,27 +4,26 @@ import torch
 import torch.nn as nn
 
 from PIL import Image
-from model import model, image_transforms
+from model import model, image_transforms_gs, image_transforms_rgb
 
 
-
-def col_select(value):
-  print(value)
-
-
-st.title("Banan Image Colorizer")
+st.title("UNET Image Colorizer")
 
 upload_file = st.file_uploader("Upload Image")
 
 if upload_file:
+  
   image = upload_file
+  image_gs = 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()
-
+  if len(np.array(image).shape) < 3: 
+    image = image_transforms_gs(image)
+  else:
+    image = image_transforms_rgb(image)
   
+  image_color = model(image.unsqueeze(0)).squeeze().permute(1, 2, 0).detach().cpu().numpy()
+
   col1, col2 = st.columns(2)
-  col1.image(image_gs_prev)
+  col1.image(image_gs)
   col2.image(image_color, clamp=True, channels='RGB')

model.py

@@ -8,14 +8,21 @@ import torch.nn.functional as F
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+image_transforms_rgb = torchvision.transforms.Compose([
+    torchvision.transforms.Resize((256, 256)),
+    torchvision.transforms.ToTensor(),
+    torchvision.transforms.Normalize(mean=[0.0,0.0,0.0], std=[1.0,1.0,1.0]),
+    torchvision.transforms.Grayscale()
+])
 
-image_transforms = torchvision.transforms.Compose([
+image_transforms_gs = torchvision.transforms.Compose([
     torchvision.transforms.Resize((256, 256)),
-    torchvision.transforms.Grayscale(),
     torchvision.transforms.ToTensor(),
-    torchvision.transforms.Normalize(mean=[0.0], std=[1.0])  
+    torchvision.transforms.Normalize(mean=[0.0], std=[1.0]),
 ])
 
+
+
 class ConvBlock(nn.Module):
     def __init__(self, in_channel, out_channel):
         super(ConvBlock, self).__init__()
@@ -99,5 +106,6 @@ class UNETFruitColor(nn.Module):
       
       
 model = UNETFruitColor()
-model.load_state_dict(torch.load("banana_colorizer_unet.pth", map_location=device),strict=True)
+model = nn.DataParallel(model).to(device)
+model.load_state_dict(torch.load("unet_colorizer_flickr_5_93_Ploss_10_14K.pth", map_location=device),strict=True)
 model.eval()

unet_colorizer_flickr_5_93_Ploss_10_14K.pth

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