From e60eb82f8d21baa4dc079c7ff41910e24b5a502f Mon Sep 17 00:00:00 2001
From: pythongosssss <125205205+pythongosssss@users.noreply.github.com>
Date: Sat, 25 Mar 2023 13:03:39 +0000
Subject: [PATCH] Fixed output image

---
 anime-segmentation/anime_segmentation.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/anime-segmentation/anime_segmentation.py b/anime-segmentation/anime_segmentation.py
index 2ba8997..217d9f5 100644
--- a/anime-segmentation/anime_segmentation.py
+++ b/anime-segmentation/anime_segmentation.py
@@ -1,3 +1,4 @@
+import copy
 import sys
 import os
 import torch
@@ -102,11 +103,11 @@ class Segment:
                 img = img[:, :, :3]
 
             img = img.astype(np.float32) / 255.0
-            img = torch.from_numpy(img)
+            img = torch.from_numpy(img)[None,]
             outputs.append(img)
 
-        # if len(outputs) == 1:
-        #     outputs = outputs[0]
+        if len(outputs) == 1:
+            outputs = outputs[0]
 
         return (outputs, first_mask,)