new image loading code

app.py

@@ -94,24 +94,22 @@ def show_images(batch: th.Tensor):
     reshaped = scaled.permute(2, 0, 3, 1).reshape([batch.shape[2], -1, 3])
     return Image.fromarray(reshaped.numpy())
 
-def read_image(path: str, size: int = 256) -> Tuple[th.Tensor, th.Tensor]:
-    pil_img = Image.open(path).convert('RGB')
+def read_image(pil_img, size: int = 256) -> Tuple[th.Tensor, th.Tensor]:
+    pil_img = pil_img.convert('RGB')
     pil_img = pil_img.resize((size, size), resample=Image.BICUBIC)
+
     img = np.array(pil_img)
     return th.from_numpy(img)[None].permute(0, 3, 1, 2).float() / 127.5 - 1
 
-def read_mask(path: str, size: int = 256) -> Tuple[th.Tensor, th.Tensor]:
-    #pil_img = PIL.Image.open(path).convert('L')
-    pil_img_full = PIL.Image.open(path).convert('RGBA')
-    #image = Image.open( inputImagePath ).convert( 'RGBA' )
-    pil_img = pil_img_full.getchannel( 'A' ) # Mode 'L'
+def read_mask(pil_img_full, size: int = 256) -> Tuple[th.Tensor, th.Tensor]:
+    pil_img_full = pil_img_full.convert('RGBA')
 
-    # pil_img = PIL.ImageOps.invert(pil_img)
+    pil_img = pil_img_full.getchannel( 'A' ) # Mode 'L'
     pil_img = pil_img.resize((size, size), resample=PIL.Image.BICUBIC)
+
     img = np.array(pil_img)[..., np.newaxis]
     return th.from_numpy(img)[None].permute(0, 3, 1, 2).float() / 255.0
 
-
 def pil_to_numpy(pil_img: Image) -> Tuple[th.Tensor, th.Tensor]:
     img = np.array(pil_img)
     return th.from_numpy(img)[None].permute(0, 3, 1, 2).float() / 127.5 - 1
@@ -122,34 +120,44 @@ def inpaint(input_img, prompt):
     print(prompt)
 
     # Save as png for later mask detection :)
-    input_img_256 = input_img.convert('RGB').resize((256, 256), resample=Image.BICUBIC)
-    input_img_64 = input_img.convert('RGB').resize((64, 64), resample=Image.BICUBIC)
+    # input_img_256 = input_img.convert('RGB').resize((256, 256), resample=Image.BICUBIC)
+    # input_img_64 = input_img.convert('RGB').resize((64, 64), resample=Image.BICUBIC)
 
-    input_img_with_mask_converted = input_img.convert('RGBA').getchannel( 'A' )  # Mode 'L'
+    # input_img_with_mask_converted = input_img.convert('RGBA').getchannel( 'A' )  # Mode 'L'
 
-    input_img_with_mask_64 = input_img_with_mask_converted.resize((64, 64), resample=Image.BICUBIC)
+    # input_img_with_mask_64 = input_img_with_mask_converted.resize((64, 64), resample=Image.BICUBIC)
     # TODO: make 256x256 mask more accurate when upscaling?
-    input_img_with_mask_256 = input_img_with_mask_converted.resize((256, 256), resample=Image.BICUBIC)
+    # input_img_with_mask_256 = input_img_with_mask_converted.resize((256, 256), resample=Image.BICUBIC)
     # return input_img, input_img_with_mask_64
     # Source image we are inpainting
-    source_image_256 = pil_to_numpy(input_img_256)
-    source_image_64 = pil_to_numpy(input_img_64)
+    source_image_256 = read_image(input_img, size=256)
+    source_image_64 = read_image(input_img, size=64)
+
+    source_mask_64 = read_mask(input_img, size=64)
+    source_mask_64 = (source_mask_64>0.5).float()
+
+    # these are better but will leave a "mark"
+    source_mask_256 = read_mask(input_img, size=256)
+    source_mask_256 = (source_mask_256>0.5).float()
+    # source_mask_256 = F.interpolate(source_mask_64, (256, 256), mode='nearest')
+    # source_image_256 = pil_to_numpy(input_img_256)
+    # source_image_64 = pil_to_numpy(input_img_64)
 
     # Since gradio doesn't supply which pixels were drawn, we need to find it ourselves!
     # Assuming that all black pixels are meant for inpainting.
     # input_img_with_mask_64 = input_img_with_mask.convert('L').resize((64, 64), resample=Image.BICUBIC)
-    gray_scale_source_image_64 = image_to_tensor(input_img_with_mask_64)
-    gray_scale_source_image_256 = image_to_tensor(input_img_with_mask_256)
+    # gray_scale_source_image_64 = image_to_tensor(input_img_with_mask_64)
+    # gray_scale_source_image_256 = image_to_tensor(input_img_with_mask_256)
 
-    source_mask_64 = (gray_scale_source_image_64!=0).float()
-    source_mask_256 = (gray_scale_source_image_256!=0).float()
+    # source_mask_64 = (gray_scale_source_image_64!=0).float()
+    # source_mask_256 = (gray_scale_source_image_256!=0).float()
 
-    source_mask_64_img = tensor_to_image(source_mask_64)
+    # source_mask_64_img = tensor_to_image(source_mask_64)
 
     # The mask should always be a boolean 64x64 mask, and then we
     # can upsample it for the second stage.
-    source_mask_64 = source_mask_64.unsqueeze(0)
-    source_mask_256 = source_mask_256.unsqueeze(0)
+    # source_mask_64 = source_mask_64.unsqueeze(0)
+    # source_mask_256 = source_mask_256.unsqueeze(0)
     # source_mask_256 = F.interpolate(source_mask_64, (256, 256), mode='nearest')