|
- import os
- import numpy as np
- import cv2
- import csv
- from mindspore import Tensor
- from mindspore import ops
- from PIL import Image
- import mindspore.dataset.vision as CV
- import mindspore as ms
-
- class HDR_Folder:
- def __init__(self, root, index, patch_num):
-
- imgname = []
- mos_all = []
-
- # csv_file = os.path.join('J:/', 'train_mos.csv')
- csv_file = '/cache/data/TrainInfoPhaseI/train_mos_20.csv' # for server
- with open(csv_file, encoding='utf-8-sig') as f:
- reader = csv.DictReader(f)
- for row in reader:
- imgname.append(row['img_name'])
- mos = np.array(float(row['mos'])).astype(np.float32)
- mos_all.append(mos)
-
- sample = []
- for i, item in enumerate(index):
- for aug in range(patch_num):
- sample.append((os.path.join(root, imgname[item]),
- mos_all[item],
- round(mos_all[item] / 1)
- ))
- self.samples = sample
-
- def __getitem__(self, index):
- """
- Args:
- index (int): Index
-
- Returns:
- tuple: (sample, target) where target is class_index of the target class.
- """
- hdr_path, mos, cls = self.samples[index]
- hdr_img = cv2.imread(hdr_path, flags=-1)[:, :, ::-1] / 65535
- # hdr_img = cv2.resize(hdr_img, dsize=(512, 512), interpolation=cv2.INTER_AREA)
- # hdr_img = np.array(hdr_img, np.float32)
- # # hdr_img = Tensor(hdr_img)
- # input_perm = (2, 0, 1)
- # hdr_img = np.transpose(hdr_img, input_perm)
-
-
- return hdr_img, mos, cls
-
- def __len__(self):
- length = len(self.samples)
- return length
-
- class test_Folder:
-
- def __init__(self, root, index, transform, patch_num):
-
- imgname = []
-
-
- # csv_file = os.path.join('J:/', 'test_mos.csv')
- csv_file = '/cache/data/TestInfoPhaseI/test_frame_10_patch.csv' # for server
-
- with open(csv_file, encoding='utf-8-sig') as f:
- reader = csv.DictReader(f)
- for row in reader:
- imgname.append(row['img_name'])
-
- sample = []
- for i, item in enumerate(index):
- for aug in range(patch_num):
- sample.append((os.path.join(root, imgname[item])))
-
- self.samples = sample
- self.transform = transform
-
- def __getitem__(self, index):
- """
- Args:
- index (int): Index
-
- Returns:
- tuple: (sample, target) where target is class_index of the target class.
- """
- hdr_path = self.samples[index]
- hdr_img = cv2.imread(hdr_path, flags=-1)[:, :, ::-1] / 65535
- mos = Tensor(mos)
- cls = Tensor(cls)
- # hdr_img = cv2.resize(hdr_img, dsize=(1080, 1080), interpolation=cv2.INTER_AREA)
- hdr_img = CV.ToTensor()(hdr_img)
- hdr_img = Tensor(hdr_img, ms.float32)
- input_perm = (2, 0, 1)
- hdr_img = ops.transpose(hdr_img, input_perm)
-
- return hdr_img
-
- def __len__(self):
- length = len(self.samples)
- return length
-
- def pil_loader(path):
- with open(path, 'rb') as f:
- img = Image.open(f)
- return img.convert('RGB')
|