Fixed output image

2023-03-25 13:03:39 +00:00
parent ee0d773bdf
commit e60eb82f8d
1 changed files with 4 additions and 3 deletions
@@ -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,)