File size: 1,750 Bytes
9e14838 | 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 | # -*- coding: utf-8 -*-
import os
from glob import glob
import numpy as np
from .builder import DATASETS
from .common import CommonDataset
@DATASETS.register_module()
class DF40(CommonDataset):
def __init__(self, cfg, **kwargs):
super().__init__(cfg, **kwargs)
def _load_from_path(self, split):
# Check if the root directory exists.
root_dir = self._cfg.DATA[self.split.upper()].ROOT
assert os.path.exists(root_dir), "Root path to dataset cannot be None!"
data = self._cfg["DATA"]
data_type = data.TYPE
fake_types = self._cfg.DATA[split.upper()]["FAKETYPE"]
img_paths, labels, mask_paths, ot_props = [], [], [], []
# Load image data for each type of fake technique.
for ft in fake_types:
# Construct path: ROOT/split/fake_type (skipping data_type)
data_dir = os.path.join(root_dir, self.split, data_type, ft)
if not os.path.exists(data_dir):
raise ValueError("Data Directory is invalid!")
# Define common image extensions.
extensions = ["jpg", "jpeg", "png", "tif", "webp"]
img_paths_ = []
# Recursively search for images in the fake type directory.
for ext in extensions:
pattern = os.path.join(data_dir, "**", f"*.{ext}")
img_paths_.extend(glob(pattern, recursive=True))
# Extend the main lists with the images and corresponding labels.
img_paths.extend(img_paths_)
labels.extend(np.full(len(img_paths_), int(ft != "real_videos")))
print("{} image paths have been loaded from DF40!".format(len(img_paths)))
return img_paths, labels, mask_paths, ot_props
|