Upload test_dataset.py
Browse files- deepfillv2/test_dataset.py +10 -4
deepfillv2/test_dataset.py
CHANGED
@@ -7,10 +7,16 @@ from config import *
|
|
7 |
|
8 |
|
9 |
class InpaintDataset(Dataset):
|
10 |
-
def __init__(self):
|
11 |
-
|
12 |
-
|
13 |
-
self.
|
|
|
|
|
|
|
|
|
|
|
|
|
14 |
|
15 |
def __len__(self):
|
16 |
return len(self.imglist)
|
|
|
7 |
|
8 |
|
9 |
class InpaintDataset(Dataset):
|
10 |
+
def __init__(self, in_image, mask_image, resize_to):
|
11 |
+
if resize_to is None:
|
12 |
+
resize_to = RESIZE_TO
|
13 |
+
self.imglist = [in_image]
|
14 |
+
self.masklist = [mask_image]
|
15 |
+
self.setsize = resize_to
|
16 |
+
# def __init__(self):
|
17 |
+
# self.imglist = [INIMAGE]
|
18 |
+
# self.masklist = [MASKIMAGE]
|
19 |
+
# self.setsize = RESIZE_TO
|
20 |
|
21 |
def __len__(self):
|
22 |
return len(self.imglist)
|