Update texture_transfer.py

texture_transfer.py

@@ -36,15 +36,15 @@ def create_image_tile(input_patch, x_dim, y_dim):
 
 
 def create_image_cutout(texture_transfer_image, actual_mask):
-    image = Image.open(texture_transfer_image)
+    image = Image.open(texture_transfer_image).convert('RGB')
     mask = Image.open(actual_mask).convert('L')
     if mask.size != image.size:
         image = image.resize(mask.size, Image.Resampling.NEAREST)
     image_np = np.array(image)
     mask_np = np.array(mask)
-    binary_mask = (mask_np > 127).astype(np.uint8) * 255
-    masked_image_np = image_np * np.expand_dims(binary_mask, axis=-1) // 255
-    masked_image = Image.fromarray(masked_image_np)
+    binary_mask = (mask_np > 127).astype(np.uint8)
+    masked_image_np = image_np * np.expand_dims(binary_mask, axis=-1)
+    masked_image = Image.fromarray(masked_image_np.astype(np.uint8))
     masked_image.save('cut_out_image.png')
 
 
@@ -56,16 +56,16 @@ def paste_image(base_image_path, cutout_image_path, mask_path):
         cutout_rgb = cutout.convert("RGB")
         cutout_alpha = cutout.split()[-1]
     else:
-        cutout_rgb = cutout
+        cutout_rgb = cutout.convert("RGB")
         cutout_alpha = mask
     cutout_rgb = cutout_rgb.resize(background.size, Image.Resampling.NEAREST)
     cutout_alpha = cutout_alpha.resize(background.size, Image.Resampling.NEAREST)
-    cutout_alpha_np = np.array(cutout_alpha)
-    binary_mask = (cutout_alpha_np > 128).astype(np.uint8) * 255
     cutout_rgb_np = np.array(cutout_rgb)
     background_np = np.array(background)
-    cutout_masked = cutout_rgb_np * np.expand_dims(binary_mask, axis=-1) // 255
-    background_masked = background_np * np.expand_dims(255 - binary_mask, axis=-1) // 255
+    cutout_alpha_np = np.array(cutout_alpha)
+    cutout_alpha_np = cutout_alpha_np / 255.0
+    cutout_masked = (cutout_rgb_np * cutout_alpha_np[..., np.newaxis]).astype(np.uint8)
+    background_masked = (background_np * (1 - cutout_alpha_np[..., np.newaxis])).astype(np.uint8)
     result_np = cutout_masked + background_masked
-    result = Image.fromarray(result_np.astype(np.uint8))
-    result.save('result.png')
+    result = Image.fromarray(result_np)
+    result.save('result.png')