File size: 3,634 Bytes
8395863 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 |
import cv2
from PIL import Image
import numpy as np
import os
import torch
from torch.utils.data import Dataset, DataLoader
def preprocess_img(img_dir, channels=3):
if channels == 1:
img = cv2.imread(img_dir, 0)
elif channels == 3:
img = cv2.imread(img_dir)
shape_r = 288
shape_c = 384
img_padded = np.ones((shape_r, shape_c, channels), dtype=np.uint8)
if channels == 1:
img_padded = np.zeros((shape_r, shape_c), dtype=np.uint8)
original_shape = img.shape
rows_rate = original_shape[0] / shape_r
cols_rate = original_shape[1] / shape_c
if rows_rate > cols_rate:
new_cols = (original_shape[1] * shape_r) // original_shape[0]
img = cv2.resize(img, (new_cols, shape_r))
if new_cols > shape_c:
new_cols = shape_c
img_padded[:,
((img_padded.shape[1] - new_cols) // 2):((img_padded.shape[1] - new_cols) // 2 + new_cols)] = img
else:
new_rows = (original_shape[0] * shape_c) // original_shape[1]
img = cv2.resize(img, (shape_c, new_rows))
if new_rows > shape_r:
new_rows = shape_r
img_padded[((img_padded.shape[0] - new_rows) // 2):((img_padded.shape[0] - new_rows) // 2 + new_rows),
:] = img
return img_padded
def postprocess_img(pred, org_dir):
pred = np.array(pred)
org = cv2.imread(org_dir, 0)
shape_r = org.shape[0]
shape_c = org.shape[1]
predictions_shape = pred.shape
rows_rate = shape_r / predictions_shape[0]
cols_rate = shape_c / predictions_shape[1]
if rows_rate > cols_rate:
new_cols = (predictions_shape[1] * shape_r) // predictions_shape[0]
pred = cv2.resize(pred, (new_cols, shape_r))
img = pred[:, ((pred.shape[1] - shape_c) // 2):((pred.shape[1] - shape_c) // 2 + shape_c)]
else:
new_rows = (predictions_shape[0] * shape_c) // predictions_shape[1]
pred = cv2.resize(pred, (shape_c, new_rows))
img = pred[((pred.shape[0] - shape_r) // 2):((pred.shape[0] - shape_r) // 2 + shape_r), :]
return img
class MyDataset(Dataset):
"""Load dataset."""
def __init__(self, ids, stimuli_dir, saliency_dir, fixation_dir, transform=None):
"""
Args:
csv_file (string): Path to the csv file with annotations.
root_dir (string): Directory with all the images.
transform (callable, optional): Optional transform to be applied
on a sample.
"""
self.ids = ids
self.stimuli_dir = stimuli_dir
self.saliency_dir = saliency_dir
self.fixation_dir = fixation_dir
self.transform = transform
def __len__(self):
return len(self.ids)
def __getitem__(self, idx):
if torch.is_tensor(idx):
idx = idx.tolist()
im_path = self.stimuli_dir + self.ids.iloc[idx, 0]
image = Image.open(im_path).convert('RGB')
img = np.array(image) / 255.
img = np.transpose(img, (2, 0, 1))
img = torch.from_numpy(img)
# if self.transform:
# img = self.transform(image)
smap_path = self.saliency_dir + self.ids.iloc[idx, 1]
saliency = Image.open(smap_path)
smap = np.expand_dims(np.array(saliency) / 255., axis=0)
smap = torch.from_numpy(smap)
fmap_path = self.fixation_dir + self.ids.iloc[idx, 2]
fixation = Image.open(fmap_path)
fmap = np.expand_dims(np.array(fixation) / 255., axis=0)
fmap = torch.from_numpy(fmap)
sample = {'image': img, 'saliency': smap, 'fixation': fmap}
return sample
|