NeuralFalcon commited on
Commit
e214599
·
verified ·
1 Parent(s): f7fccca

Upload test_dataset.py

Browse files
Files changed (1) hide show
  1. 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
- self.imglist = [INIMAGE]
12
- self.masklist = [MASKIMAGE]
13
- self.setsize = RESIZE_TO
 
 
 
 
 
 
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)