File size: 2,377 Bytes
06017be | 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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | import os
import sys
# Ensure CodeFormer directory is on sys.path
tools_dir = os.path.dirname(os.path.abspath(__file__))
project_dir = os.path.dirname(tools_dir)
codeformer_dir = os.path.join(project_dir, "models", "CodeFormer")
for p in (codeformer_dir, tools_dir, project_dir):
if p not in sys.path:
sys.path.insert(0, p)
import yaml
import torch
from basicsr.data import build_dataset
def main():
print("=== Phase 4: Dataset Loader & Game Character Mixing Verification ===")
config_path = os.path.join(codeformer_dir, "options", "CodeFormer_stage3_custom.yml")
print(f"Reading configuration from: {config_path}")
with open(config_path, "r", encoding="utf-8") as f:
config = yaml.safe_load(f)
dataset_opt = config["datasets"]["train"]
# Adjust relative dataroot_gt to absolute path
dataset_opt["dataroot_gt"] = os.path.join(codeformer_dir, dataset_opt["dataroot_gt"])
dataset_opt["prefetch_mode"] = None
dataset_opt["num_worker_per_gpu"] = 0
print(f"Dataset root: {dataset_opt['dataroot_gt']}")
# Count images in dataset folder
gt_dir = dataset_opt["dataroot_gt"]
all_files = []
for root, _, files in os.walk(gt_dir):
for file in files:
if file.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
all_files.append(os.path.join(root, file))
print(f"Total image files found on disk: {len(all_files)}")
# Build BasicSR Dataset
dataset = build_dataset(dataset_opt)
print(f"Dataset object instantiated. Total items reported by dataset loader: {len(dataset)}")
# Fetch first 5 items to verify shapes and pipeline execution
print("\nFetching sample items from dataset loader:")
for i in range(min(5, len(dataset))):
item = dataset[i]
gt_tensor = item['gt']
in_tensor = item['in']
print(f" Item #{i}: GT shape={gt_tensor.shape}, Input shape={in_tensor.shape}, GT min/max=[{gt_tensor.min():.2f}, {gt_tensor.max():.2f}]")
assert gt_tensor.shape == (3, 512, 512), f"Invalid GT shape: {gt_tensor.shape}"
assert in_tensor.shape == (3, 512, 512), f"Invalid Input shape: {in_tensor.shape}"
print("\nSUCCESS: Phase 4 Dataset Expansion & Game Character Mixing verified with exit code 0!")
if __name__ == "__main__":
main()
|