| import torch |
| import torch.nn as nn |
| from torch.utils import data |
| import os |
| from PIL import Image |
| from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize |
| try: |
| from torchvision.transforms import InterpolationMode |
| BICUBIC = InterpolationMode.BICUBIC |
| except ImportError: |
| BICUBIC = Image.BICUBIC |
| import glob |
|
|
| def image_transform(n_px): |
| return Compose([ |
| Resize(n_px, interpolation=BICUBIC), |
| CenterCrop(n_px), |
| ToTensor(), |
| Normalize((0.48145466, 0.4578275, 0.40821073), |
| (0.26862954, 0.26130258, 0.27577711)), |
| ]) |
|
|
| class Image_dataset(data.Dataset): |
| def __init__(self,dataset_folder="/data1/haolin/datasets",categories=['03001627'],n_px=224): |
| self.dataset_folder=dataset_folder |
| self.image_folder=os.path.join(self.dataset_folder,'other_data') |
| self.preprocess=image_transform(n_px) |
| self.image_path=[] |
| for cat in categories: |
| subpath=os.path.join(self.image_folder,cat,"6_images") |
| model_list=os.listdir(subpath) |
| for folder in model_list: |
| model_folder=os.path.join(subpath,folder) |
| image_list=os.listdir(model_folder) |
| for image_filename in image_list: |
| image_filepath=os.path.join(model_folder,image_filename) |
| self.image_path.append(image_filepath) |
| def __len__(self): |
| return len(self.image_path) |
|
|
| def __getitem__(self,index): |
| path=self.image_path[index] |
| basename=os.path.basename(path)[:-4] |
| model_id=path.split(os.sep)[-2] |
| category=path.split(os.sep)[-4] |
| image=Image.open(path) |
| image_tensor=self.preprocess(image) |
|
|
| return {"images":image_tensor,"image_name":basename,"model_id":model_id,"category":category} |
|
|
| class Image_InTheWild_dataset(data.Dataset): |
| def __init__(self,dataset_dir="/data1/haolin/data/real_scene_process_data",scene_id="letian-310",n_px=224): |
| self.dataset_dir=dataset_dir |
| self.preprocess = image_transform(n_px) |
| self.image_path = [] |
| if scene_id=="all": |
| scene_list=os.listdir(self.dataset_dir) |
| for id in scene_list: |
| image_folder=os.path.join(self.dataset_dir,id,"6_images") |
| self.image_path+=glob.glob(image_folder+"/*/*jpg") |
| else: |
| image_folder = os.path.join(self.dataset_dir, scene_id, "6_images") |
| self.image_path += glob.glob(image_folder + "/*/*jpg") |
| def __len__(self): |
| return len(self.image_path) |
|
|
| def __getitem__(self,index): |
| path=self.image_path[index] |
| basename=os.path.basename(path)[:-4] |
| model_id=path.split(os.sep)[-2] |
| scene_id=path.split(os.sep)[-4] |
| image=Image.open(path) |
| image_tensor=self.preprocess(image) |
|
|
| return {"images":image_tensor,"image_name":basename,"model_id":model_id,"scene_id":scene_id} |
|
|
|
|
