CalvinBackup/facefusion-labs
1
0
Fork 0
mirror of https://github.com/facefusion/facefusion-labs.git synced 2026-06-25 07:59:55 +02:00
Code Issues Packages Projects Releases Wiki Activity

This is not a TensorDataset

Browse Source
This commit is contained in:
henryruhs
2025-02-21 11:03:31 +01:00
parent 3d9ff4add0
commit b33281425a
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
face_swapper/src/data_loader.py
+3 -2
View File
@@ -5,13 +5,14 @@ from typing import Tuple
import cv2
import torch
from torch.utils.data import TensorDataset
from torch import Tensor
from torch.utils.data import Dataset
from torchvision import transforms
from .types import Batch, ImagePathList, ImagePathSet
class DataLoader(TensorDataset):
class DataLoader(Dataset[Tensor]):
def __init__(self, dataset_path : str, dataset_image_pattern : str, dataset_directory_pattern : str, same_person_probability : float) -> None:
self.same_person_probability = same_person_probability
self.directory_paths = glob.glob(dataset_directory_pattern.format(dataset_path))