Change all predicted results as refined RGBA images. Fix a typo in device specification.

app.py

@@ -21,7 +21,7 @@ import zipfile
 torch.set_float32_matmul_precision('high')
 torch.jit.script = lambda f: f
 
-device = "cuda" if torch.cuda.is_available() else "CPU"
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 ### image_proc.py
 def refine_foreground(image, mask, r=90):
@@ -125,20 +125,18 @@ def predict(images, resolution, weights_file):
     for idx_image, image_src in enumerate(images):
         if isinstance(image_src, str):
             if os.path.isfile(image_src):
-                image = np.array(Image.open(image_src))
+                image_ori = Image.open(image_src)
             else:
                 response = requests.get(image_src)
                 image_data = BytesIO(response.content)
-                image = np.array(Image.open(image_data))
+                image_ori = Image.open(image_data)
         else:
-            image = image_src
-
-        image_shape = image.shape[:2]
-        image_pil = array_to_pil_image(image, tuple(resolution))
+            image_ori = Image.fromarray(image_src)
 
+        image = image_ori.convert('RGB')
         # Preprocess the image
         image_preprocessor = ImagePreprocessor(resolution=tuple(resolution))
-        image_proc = image_preprocessor.proc(image_pil)
+        image_proc = image_preprocessor.proc(image)
         image_proc = image_proc.unsqueeze(0)
 
         # Prediction
@@ -148,8 +146,8 @@ def predict(images, resolution, weights_file):
 
         # Show Results
         pred_pil = transforms.ToPILImage()(pred)
-        image_masked = refine_foreground(Image.fromarray(image), pred_pil)
-        image_masked.putalpha(pred_pil.resize(Image.fromarray(image).size))
+        image_masked = refine_foreground(image, pred_pil)
+        image_masked.putalpha(pred_pil.resize(image.size))
 
         torch.cuda.empty_cache()
 
@@ -158,12 +156,6 @@ def predict(images, resolution, weights_file):
             image_masked.save(save_file_path)
             save_paths.append(save_file_path)
 
-        # Apply the prediction mask to the original image
-        pred = torch.nn.functional.interpolate(preds, size=image_shape, mode='bilinear', align_corners=True).squeeze().numpy()
-        image_pil = image_pil.resize(pred.shape[::-1])
-        pred = np.repeat(np.expand_dims(pred, axis=-1), 3, axis=-1)
-        image_masked = (pred * np.array(image_pil)).astype(np.uint8)
-
     if tab_is_batch:
         zip_file_path = os.path.join(save_dir, "{}.zip".format(save_dir))
         with zipfile.ZipFile(zip_file_path, 'w') as zipf:
@@ -171,7 +163,7 @@ def predict(images, resolution, weights_file):
                 zipf.write(file, os.path.basename(file))
         return save_paths, zip_file_path
     else:
-        return image, image_masked
+        return (image_ori, image_masked)
 
 
 examples = [[_] for _ in glob('examples/*')][:]