mrq/DL-Art-School
1
3
Fork 5
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Return zeros rather than None when image cant be read

Browse Source
This commit is contained in:
James Betker 2020-10-11 08:33:18 -06:00
parent e785029936
commit 92cb83958a
1 changed files with 8 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
codes/data/chunk_with_reference.py
View File

@ -9,10 +9,16 @@ class ChunkWithReference:
self.path = path.path
self.tiles, _ = util.get_image_paths('img', self.path)
# Odd failures occur at times. Rather than crashing, report the error and just return zeros.
def read_image_or_get_zero(self, img_path):
img = util.read_img(None, img_path, rgb=True)
if img is None:
return np.zeros(128, 128, 3)
def __getitem__(self, item):
centers = torch.load(osp.join(self.path, "centers.pt"))
ref = util.read_img(None, osp.join(self.path, "ref.jpg"), rgb=True)
tile = util.read_img(None, self.tiles[item], rgb=True)
ref = self.read_image_or_get_zero(osp.join(self.path, "ref.jpg"))
tile = self.read_image_or_get_zero(self.tiles[item])
tile_id = int(osp.splitext(osp.basename(self.tiles[item]))[0])
center, tile_width = centers[tile_id]
mask = np.full(tile.shape[:2] + (1,), fill_value=.1, dtype=tile.dtype)