(close #56) Merge pull request #59 from sushmanthreddy/opencv

use imageio instead of opencv

requirements.txt

@@ -1,6 +1,6 @@
 torch>=1.8.0
 torchvision>=0.8.1
-opencv-python
+imageio
 numpy
 tqdm
 mkl

torch_dreams/custom_image_param.py

@@ -1,7 +1,8 @@
 from .auto_image_param import BaseImageParam
 
-import cv2 
+
 import torch
+import imageio
 
 from .utils import (
     lucid_colorspace_to_rgb, 
@@ -42,9 +43,10 @@ def __init__(self, image, device):
 
         super().__init__()
         self.device = device
+        #use imageio to read the image
         if isinstance(image, str):
-            image = cv2.cvtColor(cv2.imread(image), cv2.COLOR_BGR2RGB)/255.
-            image = torch.tensor(image).permute(-1,0,1).unsqueeze(0)
+            image = imageio.imread(image) / 255.0
+            image = torch.tensor(image).permute(2, 0, 1).unsqueeze(0)
         self.set_param(image)
 
     def normalize(self,x, device):

torch_dreams/masked_image_param.py

@@ -1,4 +1,4 @@
-import cv2
+import imageio
 import torch
 from .custom_image_param import CustomImageParam
 from .transforms import imagenet_transform
@@ -48,8 +48,8 @@ def __init__(self, mask_tensor, image=None, device="cuda"):
         self.mask = mask_tensor.to(self.device)
 
         if isinstance(image, str):
-            image = cv2.cvtColor(cv2.imread(image), cv2.COLOR_BGR2RGB) / 255.0
-            image = torch.tensor(image).permute(-1, 0, 1).unsqueeze(0)
+            image = imageio.imread(image) / 255.0
+            image = torch.tensor(image).permute(2, 0, 1).unsqueeze(0)
 
         self.original_nchw_image_tensor = image.to(device)