Update app.py

app.py

@@ -3,64 +3,68 @@ import torch
 from torchvision import transforms
 from PIL import Image
 
-
+# Path to your exported TorchScript models (.pt)
 model_paths = {
-    "All colors model": "unet_generator.pt",
-    "20 colors model": "20color_generator.pt"
+    "All colors": "unet_generator.pt",
+    "20 colors only": "20color_generator.pt"
 }
 
-
+# Check if a GPU is available
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+# Image transformations (resize and convert to tensor)
 transform = transforms.Compose([
     transforms.Resize((512, 512)),
     transforms.ToTensor(),
 ])
 
+# Function to load the selected model
 def load_model(path):
     model = torch.jit.load(path, map_location=device)
     model.eval()
     return model
 
-
+# Main colorization function
 def colorize(image, selected_model):
     """
-    Converts the image to greyscale, it shows it to the user and then it generates the colored version depending on the model selected.
+    Converts the input image to grayscale, displays it,
+    and generates the colorized version using the selected model.
     """
-    #in greyscale
+    # Convert to grayscale
     gray = image.convert("L")
 
-    # preprocess
+    # Preprocess for model input
     gray_tensor = transform(gray).unsqueeze(0).to(device)
 
-    # model selected
+    # Load the selected model
     model = load_model(model_paths[selected_model])
 
-    #colors the image
+    # Generate the colorized image
     with torch.no_grad():
         output = model(gray_tensor)
 
+    # Process output and convert to PIL image
     output = output.squeeze(0).permute(1, 2, 0).clamp(0, 1).cpu().numpy()
     output_image = Image.fromarray((output * 255).astype('uint8'))
 
-    return gray, output_image 
+    return gray, output_image  # Return grayscale and colorized images
 
+# Create Gradio interface
 gr.Interface(
     fn=colorize,
     inputs=[
-        gr.Image(type="pil", label="Input Umage"),
-        gr.Radio(choices=["All colors", "20 colors"], label="Model")
+        gr.Image(type="pil", label="Input Image"),
+        gr.Radio(choices=["All colors", "20 colors only"], label="Model")
     ],
     outputs=[
-        gr.Image(type="pil", label="Black & White"),
-        gr.Image(type="pil", label="Colored Image")
+        gr.Image(type="pil", label="Grayscale Image"),
+        gr.Image(type="pil", label="Colorized Image")
     ],
-    title="Image colorization",
+    title="Image Colorization",
     description=(
-        "Upload a color image and choose a model to see it colorized from a greyscale version. "
+        "Upload a color image and choose a model to see it colorized from a grayscale version. "
         "The system first converts the input image to black and white, then uses a trained deep learning model "
-        "to generate a colored version. You can experiment with two models: one trained on a full color palette "
+        "to generate a colorized version. You can experiment with two models: one trained on a full color palette "
         "and another limited to just 20 colors."
     )
-
 ).launch()