Update README.md

README.md

@@ -1,4 +1,5 @@
 ---
+
 language:
   - "en"
 tags:
@@ -14,22 +15,23 @@ metrics:
   - KL
   - CER
 library_name: diffusers
+
 ---
 
 ## Emuru Convolutional VAE
 
-This repository hosts the **Emuru Convolutional VAE**, described in our paper. The model features a convolutional encoder and decoder, each with four layers. The output channels for these layers are 32, 64, 128, and 256, respectively. The encoder downsamples an input RGB image \( I \in \mathbb{R}^{3 \times W \times H} \) to a latent representation with a single channel and spatial dimensions \( h \times w \) (where \( h = H/8 \) and \( w = W/8 \)). This design compresses the style information in the image, allowing a lightweight Transformer Decoder to efficiently process the latent features.
+This repository hosts the **Emuru Convolutional VAE**, described in our paper. The model features a convolutional encoder and decoder, each with four layers. The output channels for these layers are 32, 64, 128, and 256, respectively. The encoder downsamples an input RGB image (with three channels and dimensions width and height) to a latent representation with a single channel and spatial dimensions that are one-eighth of the original height and width. This design compresses the style information in the image, allowing a lightweight Transformer Decoder to efficiently process the latent features.
 
 ### Training Details
 
 - **Hardware:** NVIDIA RTX 4090  
-- **Iterations:** 60k  
-- **Optimizer:** AdamW with a learning rate of 1e-4  
+- **Iterations:** 60,000  
+- **Optimizer:** AdamW with a learning rate of 0.0001  
 - **Loss Components:**  
-  - **MAE Loss (\(\mathcal{L}_{MAE}\))** with weight 1  
-  - **WID Loss (\(\mathcal{L}_{WID}\))** with weight 0.005  
-  - **HTR Loss (\(\mathcal{L}_{HTR}\))** with weight 0.3 (using noisy teacher-forcing with probability 0.3)  
-  - **KL Loss (\(\mathcal{L}_{KL}\))** with weight \(\beta = 1\text{e-6}\)
+  - **MAE Loss:** weight of 1  
+  - **WID Loss:** weight of 0.005  
+  - **HTR Loss:** weight of 0.3 (using noisy teacher-forcing with a probability of 0.3)  
+  - **KL Loss:** with a beta parameter set to 1e-6
 
 ### Auxiliary Networks
 
@@ -54,29 +56,27 @@ from diffusers import AutoencoderKL
 import torch
 from torchvision.transforms.functional import to_tensor, to_pil_image
 from PIL import Image
-import requests
-from io import BytesIO
 
 # Load the pre-trained Emuru VAE from Hugging Face Hub.
 model = AutoencoderKL.from_pretrained("vpippi/emuru_vae")
 
-# Function to load and preprocess an RGB image from a URL:
-# Fetches the image via requests, converts it to RGB, and transforms it to a tensor normalized to [0, 1].
-def preprocess_image_from_url(url):
-    response = requests.get(url)
-    image = Image.open(BytesIO(response.content)).convert("RGB")
-    image_tensor = to_tensor(image).unsqueeze(0)  # Add batch dimension.
+# Function to preprocess an RGB image:
+# Loads the image, converts it to RGB, and transforms it to a tensor normalized to [0, 1].
+def preprocess_image(image_path):
+    image = Image.open(image_path).convert("RGB")
+    image_tensor = to_tensor(image).unsqueeze(0)  # Add batch dimension
     return image_tensor
 
 # Function to postprocess a tensor back to a PIL image for visualization:
 # Clamps the tensor to [0, 1] and converts it to a PIL image.
 def postprocess_tensor(tensor):
-    tensor = torch.clamp(tensor, 0, 1).squeeze(0)  # Remove batch dimension.
+    tensor = torch.clamp(tensor, 0, 1).squeeze(0)  # Remove batch dimension
     return to_pil_image(tensor)
 
-# Example URL of the image.
-image_url = "https://aimagelab.ing.unimore.it/imagelab/uploadedImages/000883.jpg"
-input_image = preprocess_image_from_url(image_url)
+# Example: Encode and decode an image.
+# Replace with your image path.
+image_path = "/path/to/image"
+input_image = preprocess_image(image_path)
 
 # Encode the image to the latent space.
 # The encode() method returns an object with a 'latent_dist' attribute.
@@ -89,6 +89,8 @@ with torch.no_grad():
 with torch.no_grad():
     reconstructed = model.decode(latents).sample
 
+# Load the original image for comparison.
+original_image = Image.open(image_path).convert("RGB")
 # Convert the reconstructed tensor back to a PIL image.
 reconstructed_image = postprocess_tensor(reconstructed)
 
@@ -106,8 +108,8 @@ If you'd like to test with images hosted directly on the Hugging Face Hub, consi
   from huggingface_hub import hf_hub_download
   from PIL import Image
 
-  # Replace 'vpippi/emuru_vae' and 'samples/example_image.jpg' with your details.
-  image_path = hf_hub_download(repo_id="vpippi/emuru_vae", filename="samples/example_image.jpg")
+  # Replace 'vpippi/emuru_vae' and 'samples/lam_sample.jpg' with your details.
+  image_path = hf_hub_download(repo_id="vpippi/emuru_vae", filename="samples/lam_sample.jpg")
   sample_image = Image.open(image_path).convert("RGB")
   sample_image.show()
   ```