import os import random import os.path as osp import tarfile import zipfile from collections import defaultdict import gdown import errno import warnings import json import numpy as np import h5py from torch.utils.data import Dataset as TorchDataset from PIL import Image class Datum: def __init__(self, impath="", label=0, domain=0, classname=""): self._impath = impath self._label = label self._domain = domain self._classname = classname @property def impath(self): return self._impath @property def label(self): return self._label @property def domain(self): return self._domain @property def classname(self): return self._classname class DatasetBase: dataset_dir = "" domains = [] def __init__(self, train_x=None, train_u=None, val=None, test=None): self._train_x = train_x self._train_u = train_u self._val = val self._test = test self._num_classes = self.get_num_classes(train_x) self._lab2cname, self._classnames = self.get_lab2cname(train_x) @property def train_x(self): return self._train_x @property def train_u(self): return self._train_u @property def val(self): return self._val @property def test(self): return self._test @property def lab2cname(self): return self._lab2cname @property def classnames(self): return self._classnames @property def num_classes(self): return self._num_classes @staticmethod def get_num_classes(data_source): label_set = set() for item in data_source: label_set.add(item.label) return max(label_set) + 1 @staticmethod def get_lab2cname(data_source): container = set() for item in data_source: container.add((item.label, item.classname)) mapping = {label: classname for label, classname in container} labels = list(mapping.keys()) labels.sort() classnames = [mapping[label] for label in labels] return mapping, classnames @property def template(self): return self.templates[0] def check_input_domains(self, source_domains, target_domains): assert len(source_domains) > 0, "source_domains (list) is empty" assert len(target_domains) > 0, "target_domains (list) is empty" self.is_input_domain_valid(source_domains) self.is_input_domain_valid(target_domains) def is_input_domain_valid(self, input_domains): for domain in input_domains: if domain not in self.domains: raise ValueError( "Input domain must belong to {}, " "but got [{}]".format(self.domains, domain) ) def download_data(self, url, dst, from_gdrive=True): if not osp.exists(osp.dirname(dst)): os.makedirs(osp.dirname(dst)) if from_gdrive: gdown.download(url, dst, quiet=False) else: raise NotImplementedError print("Extracting file ...") if dst.endswith(".zip"): zip_ref = zipfile.ZipFile(dst, "r") zip_ref.extractall(osp.dirname(dst)) zip_ref.close() elif dst.endswith(".tar"): tar = tarfile.open(dst, "r:") tar.extractall(osp.dirname(dst)) tar.close() elif dst.endswith(".tar.gz"): tar = tarfile.open(dst, "r:gz") tar.extractall(osp.dirname(dst)) tar.close() else: raise NotImplementedError print("File extracted to {}".format(osp.dirname(dst))) def generate_fewshot_dataset( self, *data_sources, num_shots=-1, repeat=False ): if num_shots < 1: if len(data_sources) == 1: return data_sources[0] return data_sources print(f"Creating a {num_shots}-shot dataset") output = [] for data_source in data_sources: tracker = self.split_dataset_by_label(data_source) dataset = [] for label, items in tracker.items(): if len(items) >= num_shots: sampled_items = random.sample(items, num_shots) else: if repeat: sampled_items = random.choices(items, k=num_shots) else: sampled_items = items dataset.extend(sampled_items) output.append(dataset) if len(output) == 1: return output[0] return output def split_dataset_by_label(self, data_source): output = defaultdict(list) for item in data_source: output[item.label].append(item) return output def split_dataset_by_domain(self, data_source): output = defaultdict(list) for item in data_source: output[item.domain].append(item) return output class DatasetWrapper(TorchDataset): def __init__(self, data_source, transform=None, is_train=False): self.data_source = data_source self.transform = transform self.is_train = is_train self._h5_cache = {} self._h5_info_printed = set() def __len__(self): return len(self.data_source) def __getitem__(self, idx): item = self.data_source[idx] impath = item.impath if isinstance(impath, str): img0 = Image.open(impath).convert("RGB") elif isinstance(impath, tuple) and len(impath) == 4 and impath[0] == 'h5': _, fpath, key, index = impath f = self._h5_cache.get(fpath) if f is None: f = h5py.File(fpath, 'r') self._h5_cache[fpath] = f # Print file info once try: keys = list(f.keys()) print(f"H5 open: {os.path.basename(fpath)} keys={keys[:5]}{'...' if len(keys) > 5 else ''}") except Exception as e: print(f"H5 open (keys) failed for {fpath}: {e}") # Print dataset info per file the first time we see this path if fpath not in self._h5_info_printed: try: ds = f[key] print(f"H5 dataset: {os.path.basename(fpath)}[{key}] shape={getattr(ds, 'shape', '?')} dtype={getattr(ds, 'dtype', '?')}") except Exception as e: print(f"H5 dataset info failed for {fpath}[{key}]: {e}") self._h5_info_printed.add(fpath) arr = f[key][int(index)] arr = np.asarray(arr) # Convert CHW -> HWC if needed if arr.ndim == 3 and arr.shape[0] in (1, 3) and arr.shape[-1] not in (1, 3): arr = np.transpose(arr, (1, 2, 0)) # Ensure HWC and uint8 if arr.ndim == 3 and arr.shape[-1] in (1, 3): pass else: raise ValueError(f"Unexpected H5 image shape: {arr.shape}") if arr.dtype != np.uint8: arr = arr.astype(np.uint8) if arr.shape[-1] == 1: img0 = Image.fromarray(arr.squeeze(-1), mode='L').convert('RGB') else: img0 = Image.fromarray(arr, mode='RGB') else: # Fallback: if already PIL Image if isinstance(impath, Image.Image): img0 = impath else: raise ValueError("Unsupported impath type in DatasetWrapper") if self.transform: img = self.transform(img0) else: img = img0 return img, item.label def check_isfile(fpath): isfile = osp.isfile(fpath) if not isfile: warnings.warn('No file found at "{}"'.format(fpath)) return isfile def mkdir_if_missing(dirname): if not osp.exists(dirname): try: os.makedirs(dirname) except OSError as e: if e.errno != errno.EEXIST: raise def listdir_nohidden(path, sort=False): items = [f for f in os.listdir(path) if not f.startswith(".")] if sort: items.sort() return items def read_json(fpath): with open(fpath, "r") as f: obj = json.load(f) return obj def write_json(obj, fpath): mkdir_if_missing(osp.dirname(fpath)) with open(fpath, "w") as f: json.dump(obj, f, indent=4, separators=(",", ": "))