sebastiansarasti commited on
Commit
eb29cc1
·
verified ·
1 Parent(s): 883346a

Create utils.py

Browse files
Files changed (1) hide show
  1. utils.py +39 -0
utils.py ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import numpy as np
3
+ from PIL import Image
4
+ from torchvision.transforms import Compose, Resize, Grayscale, ToTensor, ToPILImage
5
+
6
+
7
+ # global variable for the grayscale transform
8
+ transform_gs = Compose(
9
+ [Resize((128, 128)), Grayscale(num_output_channels=1), ToTensor()]
10
+ )
11
+
12
+ def process_gs_image(image):
13
+ """
14
+ Function to process the grayscale image.
15
+ """
16
+ # Save original size for later use
17
+ original_size = image.size # (width, height)
18
+
19
+ # Convert the image to grayscale and resize
20
+ image = transform_gs(image)
21
+
22
+ # Add the batch dimension
23
+ image = image.unsqueeze(0)
24
+
25
+ # Return both the processed image and original size
26
+ return image, original_size
27
+
28
+ def inverse_transform_cs(tensor, original_size):
29
+ """
30
+ Function to convert the tensor back to the color image and resize it to its original size.
31
+ """
32
+ # Convert the tensor back to a PIL image
33
+ to_pil = ToPILImage()
34
+ pil_image = to_pil(tensor.squeeze(0)) # Remove the batch dimension
35
+
36
+ # Resize the image back to the original size
37
+ pil_image = pil_image.resize(original_size)
38
+
39
+ return pil_image