diff --git a/py/extand_canvas.py b/py/extand_canvas.py index a792d48..d044fa2 100644 --- a/py/extand_canvas.py +++ b/py/extand_canvas.py @@ -56,7 +56,8 @@ class ExtendCanvas: m = 1 - m l_masks.append(tensor2pil(torch.unsqueeze(m, 0)).convert('L')) else: - l_masks.append(Image.new('L', size=tensor2pil(l_images[0]).size, color='white')) + if len(l_masks) == 0: + l_masks.append(Image.new('L', size=tensor2pil(l_images[0]).size, color='white')) max_batch = max(len(l_images), len(l_masks)) for i in range(max_batch): diff --git a/py/mask_edge_ultrl_detail.py b/py/mask_edge_ultrl_detail.py index 5d7ec63..9da0452 100644 --- a/py/mask_edge_ultrl_detail.py +++ b/py/mask_edge_ultrl_detail.py @@ -49,6 +49,7 @@ class MaskEdgeUltraDetail: for i in range(len(l_images)): _image = l_images[i] orig_image = tensor2pil(_image).convert('RGB') + _image = pil2tensor(orig_image) _mask = l_masks[i] if mask_grow != 0: _mask = expand_mask(_mask, mask_grow, mask_grow//2) diff --git a/py/mask_edge_ultrl_detail_v2.py b/py/mask_edge_ultrl_detail_v2.py index 6d118ce..65e9c65 100644 --- a/py/mask_edge_ultrl_detail_v2.py +++ b/py/mask_edge_ultrl_detail_v2.py @@ -50,6 +50,7 @@ class MaskEdgeUltraDetailV2: for i in range(len(l_images)): _image = l_images[i] orig_image = tensor2pil(_image).convert('RGB') + _image = pil2tensor(orig_image) _mask = l_masks[i] if mask_grow != 0: _mask = expand_mask(_mask, mask_grow, mask_grow//2) diff --git a/py/person_mask_Ultra.py b/py/person_mask_Ultra.py index dbc71d2..3420952 100644 --- a/py/person_mask_Ultra.py +++ b/py/person_mask_Ultra.py @@ -76,8 +76,9 @@ class PersonMaskUltra: # image = torch.unsqueeze(image, 0) orig_image = tensor2pil(image.unsqueeze(0)).convert('RGB') # Convert the Tensor to a PIL image - i = 255. * image.cpu().numpy() - image_pil = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + # i = 255. * image.cpu().numpy() + # image_pil = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + image_pil = tensor2pil(image.unsqueeze(0)).convert('RGB') # create our foreground and background arrays for storing the mask results mask_background_array = np.zeros((image_pil.size[0], image_pil.size[1], 4), dtype=np.uint8) mask_background_array[:] = (0, 0, 0, 255) @@ -127,7 +128,7 @@ class PersonMaskUltra: tensor_mask = tensor_mask.squeeze(3)[..., 0] _mask = tensor2pil(tensor_mask).convert('L') if process_detail: - _mask = tensor2pil(mask_edge_detail(image.unsqueeze(0), pil2tensor(_mask), detail_range, black_point, white_point)) + _mask = tensor2pil(mask_edge_detail(pil2tensor(orig_image), pil2tensor(_mask), detail_range, black_point, white_point)) ret_image = RGB2RGBA(orig_image, _mask) ret_images.append(pil2tensor(ret_image)) ret_masks.append(image2mask(_mask)) diff --git a/py/person_mask_ultra_v2.py b/py/person_mask_ultra_v2.py index af97852..b3341f8 100644 --- a/py/person_mask_ultra_v2.py +++ b/py/person_mask_ultra_v2.py @@ -134,11 +134,11 @@ class PersonMaskUltraV2: detail_range = detail_erode + detail_dilate if process_detail: if detail_method == 'GuidedFilter': - _mask = guided_filter_alpha(_image, _mask, detail_range // 6 + 1) + _mask = guided_filter_alpha(pil2tensor(orig_image), _mask, detail_range // 6 + 1) _mask = tensor2pil(histogram_remap(_mask, black_point, white_point)) elif detail_method == 'PyMatting': _mask = tensor2pil( - mask_edge_detail(_image, _mask, + mask_edge_detail(pil2tensor(orig_image), _mask, detail_range // 8 + 1, black_point, white_point)) else: _trimap = generate_VITMatte_trimap(_mask, detail_erode, detail_dilate) diff --git a/py/rembg_ultra.py b/py/rembg_ultra.py index 3e22798..c50e865 100644 --- a/py/rembg_ultra.py +++ b/py/rembg_ultra.py @@ -32,6 +32,7 @@ class RemBgUltra: for i in image: i = torch.unsqueeze(i, 0) + i = pil2tensor(tensor2pil(i).convert('RGB')) orig_image = tensor2pil(i).convert('RGB') _mask = RMBG(orig_image) if process_detail: diff --git a/py/rembg_ultra_v2.py b/py/rembg_ultra_v2.py index 1d0eeba..abf7680 100644 --- a/py/rembg_ultra_v2.py +++ b/py/rembg_ultra_v2.py @@ -37,6 +37,7 @@ class RmBgUltraV2: for i in image: i = torch.unsqueeze(i, 0) + i = pil2tensor(tensor2pil(i).convert('RGB')) orig_image = tensor2pil(i).convert('RGB') _mask = RMBG(orig_image) _mask = pil2tensor(_mask) diff --git a/py/segment_anything_ultra.py b/py/segment_anything_ultra.py index 3d7a51f..46d987f 100644 --- a/py/segment_anything_ultra.py +++ b/py/segment_anything_ultra.py @@ -47,6 +47,7 @@ class SegmentAnythingUltra: for i in image: i = torch.unsqueeze(i, 0) + i = pil2tensor(tensor2pil(i).convert('RGB')) item = tensor2pil(i).convert('RGBA') boxes = groundingdino_predict(DINO_MODEL, item, prompt, threshold) if boxes.shape[0] == 0: diff --git a/py/segment_anything_ultra_v2.py b/py/segment_anything_ultra_v2.py index ce1a95a..300de6b 100644 --- a/py/segment_anything_ultra_v2.py +++ b/py/segment_anything_ultra_v2.py @@ -61,6 +61,7 @@ class SegmentAnythingUltraV2: for i in image: i = torch.unsqueeze(i, 0) + i = pil2tensor(tensor2pil(i).convert('RGB')) _image = tensor2pil(i).convert('RGBA') boxes = groundingdino_predict(DINO_MODEL, _image, prompt, threshold) if boxes.shape[0] == 0: