Added the notebook and weights for the PgIC (Progressive Image Compression) conv autoencoder

.gitignore

@@ -0,0 +1 @@
+venv

Image_Transfer_Scripts/With_AE/client_decoder.py

@@ -0,0 +1,96 @@
+import io
+import socket
+import time
+from io import BytesIO
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image
+
+
+class Decoder(nn.Module):
+    def __init__(self):
+        super(Decoder, self).__init__()
+        self.deconv1 = nn.ConvTranspose2d(64, 256, kernel_size=3, stride=2, padding=1, output_padding=1)
+        self.deconv2 = nn.ConvTranspose2d(256, 128, kernel_size=3, stride=2, padding=1, output_padding=1)
+        self.deconv3 = nn.ConvTranspose2d(128, 64, kernel_size=3, stride=2, padding=1, output_padding=1)
+        self.deconv4 = nn.ConvTranspose2d(64, 32, kernel_size=3, stride=2, padding=1, output_padding=1)
+        self.deconv5 = nn.ConvTranspose2d(32, 3, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x):
+        x = F.relu(self.deconv1(x))
+        x = F.relu(self.deconv2(x))
+        x = F.relu(self.deconv3(x))
+        x = F.relu(self.deconv4(x))
+        x = torch.sigmoid(self.deconv5(x))
+        return x
+
+decoder = Decoder()
+# here if Tests/Scripts/decoder_model.pth is not found then try using Tests\Scripts\decoder_model.pth
+decoder.load_state_dict(
+    torch.load(r"PgIC_decoder_b8.pth", map_location=torch.device("cpu"))
+    #torch.load(r"AutoEncoded_Image_Transfer\AutoEncoder_Weights\PgIC_decoder_9M.pth", map_location=torch.device("cpu"))
+)   
+decoder.eval()
+print(decoder)
+
+
+def receive_image_client(server_ip, server_port):
+    client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    client_socket.connect((server_ip, server_port))
+
+    start_time = time.time()
+
+    image_data = b""
+    while True:
+        chunk = client_socket.recv(1024)
+        if not chunk:
+            break
+        image_data += chunk
+    end_time = time.time()
+    # The start time is sent from the server as a string which is decoded here
+    start_time = (image_data.split(b"   ")[1]).decode()
+    image_data = image_data.split(b"   ")[0]
+
+    # Transfer time is calculated here
+    transfer_time = end_time - float(start_time)
+    received_encoded_output = torch.load(io.BytesIO(image_data))
+
+    print(f"Image received successfully in {transfer_time} seconds")
+    client_socket.close()
+
+    display_received_image(received_encoded_output)
+
+
+def display_received_image(encoded_output):
+    # result=torch.Tensor(numpy.frombuffer(encoded_output, dtype=numpy.int32))
+    decoded_output = decoder(encoded_output)
+    # image = Image.open(BytesIO(decoded_output))
+    decoded_output_np = decoded_output.detach().numpy()
+    img = np.transpose(decoded_output_np[0], (1, 2, 0))
+    plt.imsave('received_image.png', img)
+    fig, ax = plt.subplots(figsize=(5, 5))
+    ax.imshow(img)
+    ax.get_xaxis().set_visible(False)
+    ax.get_yaxis().set_visible(False)
+
+    plt.title("recieved image")
+    plt.show()
+
+def get_host_ip():
+    try:
+        host_name = socket.gethostname()
+        host_ip = socket.gethostbyname(host_name)
+        return host_ip
+    except:
+        print("Unable to get Hostname and IP")
+        return None
+
+if __name__ == "__main__":
+    server_ip = get_host_ip()  # Replace with server IP when not working locally
+    server_port = 55555  # Random port which is free at any time
+
+    receive_image_client(server_ip, server_port)

server_encoder.py → Image_Transfer_Scripts/With_AE/server_encoder.py

@@ -10,35 +10,30 @@
 import io
 import torchvision.transforms as transforms
 
-
-class Encoder(nn.Module):
+class Encoder(nn.Module): 
     def __init__(self):
         super(Encoder, self).__init__()
-
-        # Encoder layers
-         # conv layer (depth from 3 --> 64), 3x3 kernels
-        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
-         # conv layer (depth from 3 --> 16), 3x3 kernels
-        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
-         # conv layer (depth from 3 --> 16), 3x3 kernels
-        self.conv3 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
-        self.pool = nn.MaxPool2d(2, 2)
+        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=2, padding=1)
+        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1)
+        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)
+        self.conv4 = nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1)
+        self.bottleneck = nn.Conv2d(256, 64, kernel_size=3, stride=1, padding=1)
 
     def forward(self, x):
-
-        # Encoder
         x = F.relu(self.conv1(x))
-        x = self.pool(x)
         x = F.relu(self.conv2(x))
-        x = self.pool(x)
         x = F.relu(self.conv3(x))
-        x_encoded = self.pool(x)
-        return x_encoded
+        x = F.relu(self.conv4(x))
+        x = self.bottleneck(x)
+        return x
 
 
 encoder = Encoder()
 # here if Tests/Scripts/encoder_model.pth is not found then try using Tests\Scripts\encoder_model.pth
-encoder.load_state_dict(torch.load('Tests/Scripts/encoder_model.pth', map_location=torch.device('cpu')))
+encoder.load_state_dict(
+    torch.load(r'PgIC_encoder_b8.pth', map_location=torch.device('cpu'))
+    #torch.load(r'AutoEncoded_Image_Transfer\AutoEncoder_Weights\PgIC_encoder_9M.pth', map_location=torch.device('cpu'))
+    )
 encoder.eval()
 
 def send_image_server(ip, port, image_path):
@@ -66,6 +61,8 @@ def send_image_server(ip, port, image_path):
         pil_image = pil_image.convert('RGB')
 
         image_tensor = transform(pil_image).unsqueeze(0)
+        pil_resized = pil_image.resize((256, 256))
+        pil_resized.save('resized_image.png')
         # Pass the resized image tensor to the encoder
         encoded_output = encoder(image_tensor)
 
@@ -82,8 +79,18 @@ def send_image_server(ip, port, image_path):
         data_connection.sendall(encoded_output_bytes)
         data_connection.close()
 
+def get_host_ip():
+    try:
+        host_name = socket.gethostname()
+        host_ip = socket.gethostbyname(host_name)
+        return host_ip
+    except:
+        print("Unable to get Hostname and IP")
+        return None
+
 if __name__ == "__main__":
-    server_ip = '127.0.0.1' # Keep the server ip
+
+    server_ip = get_host_ip() # Keep the server ip
     server_port = 55555 # any random always free port
     root = tk.Tk()
     root.withdraw()

client.py → Image_Transfer_Scripts/Without_AE/client.py

@@ -1,13 +1,16 @@
 import socket
+import time
+from io import BytesIO
+
 import matplotlib.pyplot as plt
 from PIL import Image
-from io import BytesIO
-import time
 
 def receive_image_client(server_ip, server_port):
     client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
     client_socket.connect((server_ip, server_port))
 
+    start_time = time.time()
+
     image_data = b""
     while True:
         chunk = client_socket.recv(1024)
@@ -16,25 +19,32 @@ def receive_image_client(server_ip, server_port):
         image_data += chunk
 
     end_time = time.time()
+    # 5 spaces(tab characters to separate image data from the start_time engraved in image)
+
+    start_time = (image_data.split(b"     ")[1]).decode()
+    image_data = image_data.split(b"     ")[0]
 
-    # take the start time from the image_data
-    start_time=(image_data.split(b"   ")[1]).decode()
-    image_data=image_data.split(b"   ")[0]
+    # Transfer time is calculated here
     transfer_time = end_time - float(start_time)
+    received_encoded_output = BytesIO(image_data)
 
     print(f"Image received successfully in {transfer_time} seconds")
-    client_socket.close()
 
+
+    client_socket.close()
     display_received_image(image_data)
 
+
+
 def display_received_image(image_data):
     image = Image.open(BytesIO(image_data))
     plt.imshow(image)
     plt.title("Received Image")
     plt.show()
 
+
 if __name__ == "__main__":
-    server_ip = '127.0.0.1' # Keep the server IP here
-    server_port = 55555 # Random port which is free at any time
+    server_ip = "10.20.61.160"  # Keep the server IP here
+    server_port = 55555  # Random port which is free at any time
 
-    receive_image_client(server_ip, server_port)
+    receive_image_client(server_ip, server_port)

Image_Transfer_Scripts/Without_AE/server.py

@@ -0,0 +1,55 @@
+import socket
+import time
+import cv2
+
+def send_image_server(ip, port, image_path):
+    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    server_socket.bind((ip, port))
+    server_socket.listen(1)
+
+    print(f"Server listening on {ip}:{port}")
+
+    while True:
+        data_connection, address = server_socket.accept()
+        print(f"Connection from {address}")
+
+        # Open the image using cv2 and resize to 256x256
+        try:
+            image = cv2.imread(image_path)
+            resized_image = cv2.resize(image, (256, 256))  # Use cv2 for resizing
+
+            # Convert to RGB format for byte conversion (if necessary)
+            if resized_image.shape[2] == 3:  # Already RGB
+                resized_image_data = resized_image.tobytes()
+            else:  # Convert to RGB
+                resized_image_data = cv2.cvtColor(resized_image, cv2.COLOR_BGR2RGB).tobytes()
+
+        except FileNotFoundError:
+            print(f"Error: File '{image_path}' not found.")
+            data_connection.close()
+            continue
+
+        # Combine image data with start time
+        start_time = time.time()
+        image_data_with_time = resized_image_data + b"     " + (str(start_time)).encode()
+
+
+        # Send the combined data directly
+        data_connection.sendall(image_data_with_time)
+        data_connection.close()
+
+def get_host_ip():
+    try:
+        host_name = socket.gethostname()
+        host_ip = socket.gethostbyname(host_name)
+        return host_ip
+    except:
+        print("Unable to get Hostname and IP")
+        return None
+
+if __name__ == "__main__":
+    server_ip = get_host_ip()
+    server_port = 55555
+    image_path = "./Lake.jpg"
+
+    send_image_server(server_ip, server_port, image_path)

Image_Transfer_Scripts/ssim_checker.py

@@ -0,0 +1,41 @@
+from skimage.metrics import structural_similarity as ssim_sk
+import matplotlib.pyplot as plt
+import numpy as np
+from sewar.full_ref import mse, rmse, psnr, uqi, ssim, ergas, scc, rase, sam, msssim, vifp
+
+def compute_ssim(image1,image2):
+    if image1.shape[2] == 4:
+        # Convert RGBA image to RGB
+        image1 = image1[:, :, :3]
+
+    if image2.shape[2] == 4:
+        # Convert RGBA image to RGB
+        image2 = image2[:, :, :3]
+    ssim_index = ssim_sk(image1, image2, data_range=1.0,multichannel=True, channel_axis=2)
+    image1 = (image1 * 255).astype(np.uint8)
+    image2 = (image2 * 255).astype(np.uint8)
+    print("MSE: ", mse(image1,image2))
+    print("RMSE: ", rmse(image1, image2))
+    #print("PSNR: ", psnr(image1, image2))
+    #print("SSIM: ", ssim(image1, image2))
+    print("UQI: ", uqi(image1, image2))
+    print("MSSSIM: ", msssim(image1, image2))
+    #print("ERGAS: ", ergas(image1, image2))
+    #print("SCC: ", scc(image1, image2))
+    #print("RASE: ", rase(image1, image2))
+    #print("SAM: ", sam(image1, image2))
+    print("VIF: ", vifp(image1, image2))
+    return ssim_index
+def main():
+
+    resized = plt.imread('resized_image.png')
+    recieved = plt.imread('received_image.png')
+    # print(resized.shape)
+    # print(recieved.shape)
+    return compute_ssim(resized, recieved)
+
+if __name__ == '__main__':
+    try:
+        print(f"SSIM: {main()}")
+    except KeyboardInterrupt:
+        print('Interrupted')

README.md

@@ -1,2 +1,26 @@
 # AutoEncoded_Image_Transfer
-Compressing images using Autoencoders and transferring them over the network
+Compressing images using Autoencoders and transferring them over the network. Currently supports colour image transfer at 256x256 image resolution. The encoder and decoder have been separated after training to deploy on different systems.
+
+### Results: 
+- Compression ratio of 12:1
+- Image transfer times 10 times faster
+
+![image](https://github.com/05kashyap/AutoEncoded_Image_Transfer/assets/120780494/20b596dd-2682-44e3-b09a-0ba1b7eb16d3)
+#### IMG Source : https://medium.com/@birla.deepak26/autoencoders-76bb49ae6a8f
+
+### Training Dataset: ```Stanford Dogs Dataset, Animals-10```
+
+## Directories
+
+### 1. Image_Transfer_Scripts
+> Contains 2 sub directories ```With_AE``` which houses the python scripts containing the model along image transfer across network. Also contains directory ```Without_AE``` housing python scripts without containing the model but only image transfer across network.
+
+### 2. AutoEncoder_Weights
+#### - There are 3 subdirectories
+> - Working : Contains the most stable and usable autoencoder weights (Currently for the 1M architecture).
+> - Experimental : Contains unstable autoencoder weights (Currently for the 9M architecture). At the moment has highly variable results.
+> - Legacy : Contains older versions of the autoencoder weights (Unsupported versions, 1M with less training cycles).
+
+### 3. Train-Test_Notebooks
+> Results notebook for different batch sizes : ```Batch_Tests.ipynb```
+> The rest are training notebooks