-
Notifications
You must be signed in to change notification settings - Fork 0
/
utils.py
20 lines (17 loc) · 623 Bytes
/
utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
import torch
import cv2
import numpy as np
import random
def image_to_tensor(image):
image_tensor = image.transpose(2, 0, 1)
image_tensor = image_tensor.astype(np.float32)
image_tensor = torch.from_numpy(image_tensor)
if torch.cuda.is_available(): # put on GPU if CUDA is available
image_tensor = image_tensor.cuda()
return image_tensor
def resize_and_bgr2gray(image):
image = image[0:288, 0:404]
image_data = cv2.cvtColor(cv2.resize(image, (84, 84)), cv2.COLOR_BGR2GRAY)
image_data[image_data > 0] = 255
image_data = np.reshape(image_data, (84, 84, 1))
return image_data