Update app.py

app.py

@@ -157,7 +157,12 @@ def predict(images, resolution, weights_file):
             save_file_path = os.path.join(save_dir, "{}.png".format(os.path.splitext(os.path.basename(image_src))[0]))
             image_masked.save(save_file_path)
             save_paths.append(save_file_path)
-        image_masked = np.array(image_masked.convert('RGB'))
+
+        # Apply the prediction mask to the original image
+        pred = torch.nn.functional.interpolate(pred, size=image_shape, mode='bilinear', align_corners=True).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))