update

app.py

@@ -11,7 +11,7 @@ from utils import lab_to_rgb, build_res_unet, build_mobilenet_unet  # Utility to
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 # Hàm load models
-def load_autoencoder_model(auto_model_path):
+def load_unet_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))
@@ -44,7 +44,7 @@ mobilenet_model = load_model(
     model_type='mobilenet'
 )
 
-autoencoder_model = load_autoencoder_model("weight/autoencoder.pt")
+unet_model = load_unet_model("weight/autoencoder.pt")
 
 # Transformations
 def preprocess_image(image):
@@ -67,21 +67,21 @@ def colorize_image(input_image, mode):
     with torch.no_grad():
         resnet_output = resnet_model.net_G(grayscale.unsqueeze(0))
         mobilenet_output = mobilenet_model.net_G(grayscale.unsqueeze(0))
-        autoencoder_output = autoencoder_model(grayscale.unsqueeze(0))
+        unet_output = unet_model(grayscale.unsqueeze(0))
     
     # Resize outputs to match the original size
     resnet_colorized = postprocess_image(grayscale, resnet_output, original_size)
     mobilenet_colorized = postprocess_image(grayscale, mobilenet_output, original_size)
-    autoencoder_colorized = postprocess_image(grayscale, autoencoder_output, original_size)
+    unet_colorized = postprocess_image(grayscale, unet_output, original_size)
     
     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 == "Unet":
+        return None, None, unet_colorized
     elif mode == "Comparison":
-        return resnet_colorized, mobilenet_colorized, autoencoder_colorized
+        return resnet_colorized, mobilenet_colorized, unet_colorized
 
 
 # Gradio Interface
@@ -90,7 +90,7 @@ def gradio_interface():
         # Input components
         input_image = gr.Image(type="numpy", label="Upload an Image")
         output_modes = gr.Radio(
-            choices=["ResNet", "MobileNet", "Autoencoder", "Comparison"],
+            choices=["ResNet", "MobileNet", "Unet", "Comparison"],
             value="ResNet",
             label="Output Mode"
         )
@@ -101,7 +101,7 @@ def gradio_interface():
         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)
+            unet_output = gr.Image(label="Colorized Image (Unet)", visible=False)
 
         # Output mode logic
         def update_visibility(mode):
@@ -109,7 +109,7 @@ def gradio_interface():
                 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":
+            elif mode == "Unet":
                 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)
@@ -118,14 +118,14 @@ def gradio_interface():
         output_modes.change(
             fn=update_visibility,
             inputs=[output_modes],
-            outputs=[resnet_output, mobilenet_output, autoencoder_output]
+            outputs=[resnet_output, mobilenet_output, unet_output]
         )
 
         # Submit logic
         submit_button.click(
             fn=colorize_image,
             inputs=[input_image, output_modes],
-            outputs=[resnet_output, mobilenet_output, autoencoder_output]
+            outputs=[resnet_output, mobilenet_output, unet_output]
         )
 
     return demo