mirror of
https://github.com/Haoming02/sd-webui-old-photo-restoration.git
synced 2026-04-29 02:31:37 +00:00
89a8626838
commit cd7a9c103d1ea981ecd236d4e9111fd3c1cd6c2b Author: Haoming <hmstudy02@gmail.com> Date: Tue Dec 19 11:33:44 2023 +0800 add README commit 30127cbb2a8e5f461c540729dc7ad457f66eb94c Author: Haoming <hmstudy02@gmail.com> Date: Tue Dec 19 11:12:16 2023 +0800 fix Face Enhancement distortion commit 6d52de5368c6cfbd9342465b5238725c186e00b9 Author: Haoming <hmstudy02@gmail.com> Date: Mon Dec 18 18:27:25 2023 +0800 better? args handling commit 0d1938b59eb77a038ee0a91a66b07fb9d7b3d6d4 Author: Haoming <hmstudy02@gmail.com> Date: Mon Dec 18 17:40:19 2023 +0800 bug fix related to Scratch commit 8315cd05ffeb2d651b4c57d70bf04b413ca8901d Author: Haoming <hmstudy02@gmail.com> Date: Mon Dec 18 17:24:52 2023 +0800 implement step 2 ~ 4 commit a5feb04b3980bdd80c6b012a94c743ba48cdfe39 Author: Haoming <hmstudy02@gmail.com> Date: Mon Dec 18 11:55:20 2023 +0800 process scratch commit3b18f7b042Author: Haoming <hmstudy02@gmail.com> Date: Wed Dec 13 11:57:20 2023 +0800 "init" commitd0148e0e82Author: Haoming <hmstudy02@gmail.com> Date: Wed Dec 13 10:34:39 2023 +0800 clone repo
63 lines
1.6 KiB
Python
63 lines
1.6 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# Licensed under the MIT License.
|
|
|
|
import torch.utils.data as data
|
|
from PIL import Image
|
|
import os
|
|
|
|
IMG_EXTENSIONS = [
|
|
'.jpg', '.JPG', '.jpeg', '.JPEG',
|
|
'.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP', '.tiff'
|
|
]
|
|
|
|
|
|
def is_image_file(filename):
|
|
return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
|
|
|
|
|
|
def make_dataset(dir):
|
|
images = []
|
|
assert os.path.isdir(dir), '%s is not a valid directory' % dir
|
|
|
|
for root, _, fnames in sorted(os.walk(dir)):
|
|
for fname in fnames:
|
|
if is_image_file(fname):
|
|
path = os.path.join(root, fname)
|
|
images.append(path)
|
|
|
|
return images
|
|
|
|
|
|
def default_loader(path):
|
|
return Image.open(path).convert('RGB')
|
|
|
|
|
|
class ImageFolder(data.Dataset):
|
|
|
|
def __init__(self, root, transform=None, return_paths=False,
|
|
loader=default_loader):
|
|
imgs = make_dataset(root)
|
|
if len(imgs) == 0:
|
|
raise(RuntimeError("Found 0 images in: " + root + "\n"
|
|
"Supported image extensions are: " +
|
|
",".join(IMG_EXTENSIONS)))
|
|
|
|
self.root = root
|
|
self.imgs = imgs
|
|
self.transform = transform
|
|
self.return_paths = return_paths
|
|
self.loader = loader
|
|
|
|
def __getitem__(self, index):
|
|
path = self.imgs[index]
|
|
img = self.loader(path)
|
|
if self.transform is not None:
|
|
img = self.transform(img)
|
|
if self.return_paths:
|
|
return img, path
|
|
else:
|
|
return img
|
|
|
|
def __len__(self):
|
|
return len(self.imgs)
|