update models

app.py

@@ -1,90 +1,133 @@
 import torch
+import numpy as np
 from PIL import Image
 from torchvision import transforms
 from matplotlib import pyplot as plt
 import gradio as gr
 
-from models import MainModel  # Import class for your main model
-from utils import lab_to_rgb, build_res_unet#, build_mobile_unet  # Utility to convert LAB to RGB
+from models import MainModel, UNetAuto, Autoencoder  
+from utils import lab_to_rgb, build_res_unet, build_mobilenet_unet  # Utility to convert LAB to RGB
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
+# Hàm load models
+def load_autoencoder_model(auto_model_path):
+    unet = UNetAuto(in_channels=1, out_channels=2).to(device)
+    model = Autoencoder(unet).to(device)
+    model.load_state_dict(torch.load(auto_model_path, map_location=device))
+    model.to(device)
+    model.eval()
+    return model
 
-def load_model(generator_model_path, colorization_model_path): #, model_type='resnet')
-
-    #if model_type == 'resnet':
-    net_G = build_res_unet(n_input=1, n_output=2, size=256)
-    # elif model_type == 'mobilenet':
-    #     net_G = build_mobile_unet(n_input=1, n_output=2, size=256)
+def load_model(generator_model_path, colorization_model_path, model_type='resnet'):
+    if model_type == 'resnet':
+        net_G = build_res_unet(n_input=1, n_output=2, size=256)
+    elif model_type == 'mobilenet':
+        net_G = build_mobilenet_unet(n_input=1, n_output=2, size=256)
     
     net_G.load_state_dict(torch.load(generator_model_path, map_location=device))
-    
-    # Create MainModel and load weights
     model = MainModel(net_G=net_G)
     model.load_state_dict(torch.load(colorization_model_path, map_location=device))
-    
-    # Move model to device and set to eval mode
     model.to(device)
     model.eval()
-    
     return model
 
-# Load pretrained models
 resnet_model = load_model(
     "weight/pascal_res18-unet.pt",
-    "weight/pascal_final_model_weights.pt"
-    # model_type='resnet'
+    "weight/pascal_final_model_weights.pt",
+    model_type='resnet'
 )
 
-# mobilenet_model = load_model(
-#     "weight/mobile-unet.pt",
-#     "weight/mobile_pascal_final_model_weights.pt",
-#     model_type='mobilenet'
-# )
+mobilenet_model = load_model(
+    "weight/mobile-unet.pt",
+    "weight/mobile_pascal_final_model_weights.pt",
+    model_type='mobilenet'
+)
+
+autoencoder_model = load_autoencoder_model("weight/autoencoder.pt")
 
 # Transformations
 def preprocess_image(image):
     image = image.resize((256, 256))
-    image = transforms.ToTensor()(image)[:1] * 2. - 1.  # Normalize to [-1, 1]
+    image = transforms.ToTensor()(image)[:1] * 2. - 1.
     return image
 
 def postprocess_image(grayscale, prediction):
     return lab_to_rgb(grayscale.unsqueeze(0), prediction.cpu())[0]
 
-# Prediction function
-def colorize_image(input_image):
-    # Convert input to grayscale
-    input_image = Image.fromarray(input_image).convert('L')
-    grayscale = preprocess_image(input_image).to(device)
+# Prediction function with output control
+def colorize_image(input_image, mode):
+    grayscale_image = Image.fromarray(input_image).convert('L')
+    grayscale = preprocess_image(grayscale_image).to(device)
     
-    # Generate predictions
     with torch.no_grad():
         resnet_output = resnet_model.net_G(grayscale.unsqueeze(0))
-        # mobilenet_output = mobilenet_model.net_G(grayscale.unsqueeze(0))
+        mobilenet_output = mobilenet_model.net_G(grayscale.unsqueeze(0))
+        autoencoder_output = autoencoder_model(grayscale.unsqueeze(0))
     
-    # Post-process results
     resnet_colorized = postprocess_image(grayscale, resnet_output)
-    # mobilenet_colorized = postprocess_image(grayscale, mobilenet_output)
+    mobilenet_colorized = postprocess_image(grayscale, mobilenet_output)
+    autoencoder_colorized = postprocess_image(grayscale, autoencoder_output)
     
-    return (
-        input_image,  # Grayscale image
-        resnet_colorized  # ResNet18 colorized image
-        # mobilenet_colorized  # MobileNet colorized image
-    )
+    if mode == "ResNet":
+        return resnet_colorized, None, None
+    elif mode == "MobileNet":
+        return None, mobilenet_colorized, None
+    elif mode == "Autoencoder":
+        return None, None, autoencoder_colorized
+    elif mode == "Comparison":
+        return resnet_colorized, mobilenet_colorized, autoencoder_colorized
+
 
 # Gradio Interface
-interface = gr.Interface(
-    fn=colorize_image,
-    inputs=gr.Image(type="numpy", label="Upload a Color Image"),
-    outputs=[
-        gr.Image(label="Grayscale Image"),
-        gr.Image(label="Colorized Image (ResNet18)")
-        # gr.Image(label="Colorized Image (MobileNet)")
-    ],
-    title="Image Colorization",
-    description="Upload a color image"
-)
+def gradio_interface():
+    with gr.Blocks() as demo:
+        # Input components
+        input_image = gr.Image(type="numpy", label="Upload an Image")
+        output_modes = gr.Radio(
+            choices=["ResNet", "MobileNet", "Autoencoder", "Comparison"],
+            value="ResNet",
+            label="Output Mode"
+        )
+        
+        submit_button = gr.Button("Submit")
+
+        # Output components
+        with gr.Row():  # Place output images in a single row
+            resnet_output = gr.Image(label="Colorized Image (ResNet18)", visible=False)
+            mobilenet_output = gr.Image(label="Colorized Image (MobileNet)", visible=False)
+            autoencoder_output = gr.Image(label="Colorized Image (Autoencoder)", visible=False)
+
+        # Output mode logic
+        def update_visibility(mode):
+            if mode == "ResNet":
+                return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)
+            elif mode == "MobileNet":
+                return gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)
+            elif mode == "Autoencoder":
+                return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
+            elif mode == "Comparison":
+                return gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
+
+        # Dynamic event listener for output mode changes
+        output_modes.change(
+            fn=update_visibility,
+            inputs=[output_modes],
+            outputs=[resnet_output, mobilenet_output, autoencoder_output]
+        )
+
+        # Submit logic
+
+        submit_button.click(
+            fn=colorize_image,
+            inputs=[input_image, output_modes],
+            outputs=[resnet_output, mobilenet_output, autoencoder_output]
+        )
+
+    return demo
+
+
 
-# Launch Gradio app
-if __name__ == '__main__':
-    interface.launch()
+# Launch
+if __name__ == "__main__":
+    gradio_interface().launch()

models.py

@@ -171,4 +171,124 @@ class MainModel(nn.Module):
         self.set_requires_grad(self.net_D, False)
         self.opt_G.zero_grad()
         self.backward_G()
-        self.opt_G.step()
+        self.opt_G.step()
+
+
+class UNetAuto(nn.Module):
+
+    def __init__(self, in_channels=1, out_channels=2, features=[64, 128, 256, 512]):
+
+        super(UNetAuto, self).__init__()
+
+        self.encoder = nn.ModuleList()
+
+        self.decoder = nn.ModuleList()
+
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+
+
+
+        # Encoder part
+
+        for feature in features:
+
+            self.encoder.append(self._block(in_channels, feature))
+
+            in_channels = feature
+
+
+
+        # Decoder part (Upsampling)
+
+        for feature in reversed(features):
+
+            self.decoder.append(nn.ConvTranspose2d(feature * 2, feature, kernel_size=2, stride=2))
+
+            self.decoder.append(self._block(feature * 2, feature))
+
+
+
+        # Final Convolution
+
+        self.bottleneck = self._block(features[-1], features[-1] * 2)
+
+        self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1)
+
+
+
+    def forward(self, x): #, t):
+
+        skip_connections = []
+
+
+
+        # Encode
+
+        for layer in self.encoder:
+
+            x = layer(x)
+
+            skip_connections.append(x)
+
+            x = self.pool(x)
+
+
+
+        # Bottleneck
+
+        x = self.bottleneck(x)
+
+
+
+        # Decode
+
+        skip_connections = skip_connections[::-1]
+
+        for idx in range(0, len(self.decoder), 2):
+
+            x = self.decoder[idx](x)
+
+            skip_connection = skip_connections[idx // 2]
+
+            x = torch.cat((x, skip_connection), dim=1)  # Skip connection
+
+            x = self.decoder[idx + 1](x)
+
+
+
+        return self.final_conv(x)
+
+
+
+    def _block(self, in_channels, out_channels):
+
+        return nn.Sequential(
+
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False),
+
+            nn.BatchNorm2d(out_channels),
+
+            nn.ReLU(inplace=True),
+
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False),
+
+            nn.BatchNorm2d(out_channels),
+
+            nn.ReLU(inplace=True),
+
+        )
+    
+
+class Autoencoder(nn.Module):
+
+    def __init__(self, model):
+
+        super(Autoencoder, self).__init__()
+
+        self.model = model
+
+
+
+    def forward(self, x): #, t):
+
+        return self.model(x)#, t)

utils.py

@@ -28,12 +28,12 @@ def build_res_unet(n_input=1, n_output=2, size=256):
     net_G = DynamicUnet(body, n_output, (size, size)).to(device)
     return net_G
 
-# def build_mobile_unet(n_input=1, n_output=2, size=256):
-#     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-#     mobilenet_model = mobilenet_v2(pretrained=True)  
-#     body = create_body(mobilenet_model, n_in=n_input, cut=-2)
-#     net_G = DynamicUnet(body, n_output, (size, size)).to(device)
-#     return net_G
+def build_mobilenet_unet(n_input=1, n_output=2, size=256):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    mobilenet = mobilenet_v2(pretrained=True)
+    body = create_body(mobilenet.features, pretrained=True, n_in=n_input, cut=-2)
+    net_G = DynamicUnet(body, n_output, (size, size)).to(device)
+    return net_G
 
 def create_loss_meters():
     loss_D_fake = AverageMeter()

weight/autoencoder.pt

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