Update README.md

README.md

@@ -22,19 +22,18 @@ cd autoencoder-grayscale2color-landscape
 git lfs pull
 ```
 ### Import Libraries
-```bash
+```python
 from PIL import Image
 import os
 import numpy as np
 import tensorflow as tf
 import requests
 from skimage.color import lab2rgb
-
+import matplotlib.pyplot as plt
 from models.auto_encoder_gray2color import SpatialAttention
-
 ```
 ### Load Model file
-```bash
+```python
 # Load the saved model once at startup
 load_model_path = "./ckpts/best_model.h5"
 
@@ -45,24 +44,8 @@ loaded_autoencoder = tf.keras.models.load_model(
 )
 ```
 
-### Inferenc Step
+### Define Functions
 ```python
-# Assuming process_image function is defined as provided
-# and required imports (tensorflow, skimage.color for lab2rgb) are available
-
-def display_images_pil(input_path):
-    # Read input image
-    input_img = Image.open(input_path)
-    
-    # Process the image using the provided function
-    output_img = process_image(input_img)
-    
-    # Display input image
-    input_img.show(title="Input Image")
-    
-    # Display output image
-    output_img.show(title="Output Image")
-
 def process_image(input_img):
     # Store original input dimensions
     original_width, original_height = input_img.size
@@ -94,10 +77,50 @@ def process_image(input_img):
 
     return rgb_image
 
+def process_and_plot_images(input_path):
+    # Read input image
+    input_img = Image.open(input_path)
+    
+    # Process the image (placeholder for your process_image function)
+    output_img = process_image(input_img)
+    
+    # Save output image to output.jpg
+    output_img.save("output.jpg")
+    
+    return input_img, output_img
 
+def plot_in_out_images(input_img, output_img):
+    # Create a figure with two subplots for input and output images
+    plt.figure(figsize=(17, 8), dpi=300)  # Set DPI to 300
+    
+    # Plot input image
+    plt.subplot(1, 2, 1)
+    plt.imshow(input_img, cmap='gray')
+    plt.title("Input Image")
+    plt.axis('off')  # Hide axes for cleaner display
+    
+    # Plot output image
+    plt.subplot(1, 2, 2)
+    plt.imshow(output_img, cmap='gray')
+    plt.title("Output Image")
+    plt.axis('off')  # Hide axes for cleaner display
+    
+    # Save the figure as output.jpg with 300 DPI
+    plt.savefig("output.jpg", dpi=300, bbox_inches='tight')
+    
+    # Show the plot
+    plt.show()
+```
+### Inference
+```python
 # Example usage
-if __name__ == "__main__":
-    # Replace 'input_image.jpg' with the path to your image
-    image_path = "<input_image.jpg>"
-    display_images_pil(image_path)
-```
+WIDTH, HEIGHT = 512, 512
+# Replace 'input_image.jpg' with the path to your image
+image_path = "<input_image.jpg>"
+input_img, output_img = process_and_plot_images(image_path)
+
+plot_in_out_images(input_img, output_img)
+```
+
+### Example Output
+![Plot Image](./examples/model_output.jpg)