| 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) |
| |
| |
|
|
| 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 |
|
|
|
|
|
|
|
|
|
|
|
|
