Upload folder using huggingface_hub

app.py

@@ -11,7 +11,7 @@ import gradio as gr
 from gradio_imageslider import ImageSlider
 
 from utils import preprocess_img, preprocess_img_from_path, postprocess_img
-from vgg19 import VGG_19
+from vgg.vgg19 import VGG_19
 from inference import inference
 
 if torch.cuda.is_available(): device = 'cuda'

u2net/model.py

@@ -0,0 +1,151 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+def init_weight(layer):
+    nn.init.xavier_uniform_(layer.weight)
+    if layer.bias is not None:
+        nn.init.constant_(layer.bias, 0)
+
+    
+class ConvBlock(nn.Module):
+    def __init__(self, in_channel, out_channel, dilation=1):
+        super(ConvBlock, self).__init__()
+        self.conv = nn.Conv2d(in_channel, out_channel, kernel_size=3, stride=1, padding=dilation, dilation=dilation)
+        self.bn = nn.BatchNorm2d(out_channel)
+        self.relu = nn.ReLU(inplace=True)
+        init_weight(self.conv)
+        
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.relu(x)
+        return x
+
+
+class RSU(nn.Module):
+    def __init__(self, L, C_in, C_out, M):
+        super(RSU, self).__init__()
+        self.conv = ConvBlock(C_in, C_out)
+        
+        self.enc = nn.ModuleList([ConvBlock(C_out, M)])
+        for i in range(L-2):
+            self.enc.append(ConvBlock(M, M))
+            
+        self.mid = ConvBlock(M, M, dilation=2)
+        
+        self.dec = nn.ModuleList([ConvBlock(2*M, M) for _ in range(L-2)])
+        self.dec.append(ConvBlock(2*M, C_out))
+    
+        self.downsample = nn.MaxPool2d(2, stride=2)
+        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear')
+        
+    def forward(self, x):
+        x = self.conv(x)
+        
+        out = []
+        for i, enc in enumerate(self.enc):
+            if i == 0: out.append(enc(x))
+            else: out.append(enc(self.downsample(out[i-1])))
+            
+        y = self.mid(out[-1])
+        
+        for i, dec in enumerate(self.dec):
+            if i > 0: y = self.upsample(y)
+            y = dec(torch.cat((out[len(self.dec)-i-1], y), dim=1))
+            
+        return x + y
+    
+
+class RSU4F(nn.Module):
+    def __init__(self, C_in, C_out, M):
+        super(RSU4F, self).__init__()
+        self.conv = ConvBlock(C_in, C_out)
+        
+        self.enc = nn.ModuleList([
+            ConvBlock(C_out, M),
+            ConvBlock(M, M, dilation=2),
+            ConvBlock(M, M, dilation=4)
+        ])
+        
+        self.mid = ConvBlock(M, M, dilation=8)
+        
+        self.dec = nn.ModuleList([
+            ConvBlock(2*M, M, dilation=4),
+            ConvBlock(2*M, M, dilation=2),
+            ConvBlock(2*M, C_out)
+        ])
+        
+    def forward(self, x):
+        x = self.conv(x)
+        
+        out = []
+        for i, enc in enumerate(self.enc):
+            if i == 0: out.append(enc(x))
+            else: out.append(enc(out[i-1]))
+            
+        y = self.mid(out[-1])
+        
+        for i, dec in enumerate(self.dec):
+            y = dec(torch.cat((out[len(self.dec)-i-1], y), dim=1))
+            
+        return x + y
+    
+    
+class U2Net(nn.Module):
+  def __init__(self):
+    super(U2Net, self).__init__()
+    self.enc = nn.ModuleList([
+      RSU(L=7, C_in=3, C_out=64, M=32),
+      RSU(L=6, C_in=64, C_out=128, M=32),
+      RSU(L=5, C_in=128, C_out=256, M=64),
+      RSU(L=4, C_in=256, C_out=512, M=128),
+      RSU4F(C_in=512, C_out=512, M=256),
+      RSU4F(C_in=512, C_out=512, M=256)
+    ])
+
+    self.dec = nn.ModuleList([
+      RSU4F(C_in=1024, C_out=512, M=256),
+      RSU(L=4, C_in=1024, C_out=256, M=128),
+      RSU(L=5, C_in=512, C_out=128, M=64),
+      RSU(L=6, C_in=256, C_out=64, M=32),
+      RSU(L=7, C_in=128, C_out=64, M=16)
+    ])
+
+    self.convs = nn.ModuleList([
+      nn.Conv2d(64, 1, 3, padding=1),
+      nn.Conv2d(64, 1, 3, padding=1),
+      nn.Conv2d(128, 1, 3, padding=1),
+      nn.Conv2d(256, 1, 3, padding=1),
+      nn.Conv2d(512, 1, 3, padding=1),
+      nn.Conv2d(512, 1, 3, padding=1)
+    ])
+
+    self.lastconv = nn.Conv2d(6, 1, 1)
+    self.downsample = nn.MaxPool2d(2, stride=2)
+
+    init_weight(self.lastconv)
+    for conv in self.convs:
+      init_weight(conv)
+
+  def upsample(self, x, target):
+    return F.interpolate(x, size=target.shape[2:], mode='bilinear')
+
+  def forward(self, x):
+    enc_out = []
+    for i, enc in enumerate(self.enc):
+      if i == 0: enc_out.append(enc(x))
+      else: enc_out.append(enc(self.downsample(enc_out[i-1])))
+
+    dec_out = [enc_out[-1]]
+    for i, dec in enumerate(self.dec):
+      dec_out.append(dec(torch.cat((self.upsample(dec_out[i], enc_out[4-i]), enc_out[4-i]), dim=1)))
+    
+    side_out = []
+    for i, conv in enumerate(self.convs):
+      if i == 0: side_out.append(conv(dec_out[5]))
+      else: side_out.append(self.upsample(conv(dec_out[5-i]), side_out[0]))
+    
+    side_out.append(self.lastconv(torch.cat(side_out, dim=1)))       
+
+    return [torch.sigmoid(s.squeeze(1)) for s in side_out]

u2net/train.py

@@ -0,0 +1 @@
+# for training u2net

vgg/vgg16.py

@@ -0,0 +1,72 @@
+import torch.nn as nn
+import torchvision.models as models
+
+""" VGG_16 Architecture
+VGG(
+  (features): Sequential(
+    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (1): ReLU(inplace=True)
+    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (3): ReLU(inplace=True)
+    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
+    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (6): ReLU(inplace=True)
+    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (8): ReLU(inplace=True)
+    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
+    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (11): ReLU(inplace=True)
+    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (13): ReLU(inplace=True)
+    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (15): ReLU(inplace=True)
+    (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
+    (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (18): ReLU(inplace=True)
+    (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (20): ReLU(inplace=True)
+    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (22): ReLU(inplace=True)
+    (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
+    (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (25): ReLU(inplace=True)
+    (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (27): ReLU(inplace=True)
+    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (29): ReLU(inplace=True)
+    (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
+  )
+  (avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
+  (classifier): Sequential(
+    (0): Linear(in_features=25088, out_features=4096, bias=True)
+    (1): ReLU(inplace=True)
+    (2): Dropout(p=0.5, inplace=False)
+    (3): Linear(in_features=4096, out_features=4096, bias=True)
+    (4): ReLU(inplace=True)
+    (5): Dropout(p=0.5, inplace=False)
+    (6): Linear(in_features=4096, out_features=1000, bias=True)
+  )
+)
+"""
+
+class VGG_16(nn.Module):
+    def __init__(self):
+        super(VGG_16, self).__init__()
+        self.model = models.vgg16(weights=models.VGG16_Weights.IMAGENET1K_V1).features[:30]
+        
+        for i, _ in enumerate(self.model):
+            if i in [4, 9, 16, 23]:
+                self.model[i] = nn.AvgPool2d(kernel_size=2, stride=2, padding=0)
+                
+    def forward(self, x):
+        features = []
+        for i, layer in enumerate(self.model):
+            x = layer(x)
+            if i in [0, 5, 10, 17, 24]:
+                features.append(x)
+        return features
+    
+
+if __name__ == '__main__':
+    model = VGG_16()
+    print(model)

vgg/vgg19.py

@@ -0,0 +1,78 @@
+import torch.nn as nn
+import torchvision.models as models
+
+""" VGG_19 Architecture
+VGG(
+  (features): Sequential(
+    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (1): ReLU(inplace=True)
+    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (3): ReLU(inplace=True)
+    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
+    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (6): ReLU(inplace=True)
+    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (8): ReLU(inplace=True)
+    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
+    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (11): ReLU(inplace=True)
+    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (13): ReLU(inplace=True)
+    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (15): ReLU(inplace=True)
+    (16): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (17): ReLU(inplace=True)
+    (18): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
+    (19): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (20): ReLU(inplace=True)
+    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (22): ReLU(inplace=True)
+    (23): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (24): ReLU(inplace=True)
+    (25): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (26): ReLU(inplace=True)
+    (27): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
+    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (29): ReLU(inplace=True)
+    (30): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (31): ReLU(inplace=True)
+    (32): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (33): ReLU(inplace=True)
+    (34): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
+    (35): ReLU(inplace=True)
+    (36): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
+  )
+  (avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
+  (classifier): Sequential(
+    (0): Linear(in_features=25088, out_features=4096, bias=True)
+    (1): ReLU(inplace=True)
+    (2): Dropout(p=0.5, inplace=False)
+    (3): Linear(in_features=4096, out_features=4096, bias=True)
+    (4): ReLU(inplace=True)
+    (5): Dropout(p=0.5, inplace=False)
+    (6): Linear(in_features=4096, out_features=1000, bias=True)
+  )
+)
+"""
+
+class VGG_19(nn.Module):
+    def __init__(self):
+        super(VGG_19, self).__init__()
+        self.model = models.vgg19(weights=models.VGG19_Weights.IMAGENET1K_V1).features[:30]
+        
+        for i, _ in enumerate(self.model):
+            if i in [4, 9, 18, 27]:
+                self.model[i] = nn.AvgPool2d(kernel_size=2, stride=2, padding=0)
+                
+    def forward(self, x):
+        features = []
+        for i, layer in enumerate(self.model):
+            x = layer(x)
+            if i in [0, 5, 10, 19, 28]:
+                features.append(x)
+        return features
+    
+
+if __name__ == '__main__':
+    model = VGG_19()
+    print(model)