fix: SegmDetectorCombined, BboxDetectorCombined - return empty mask instead of None

This commit is contained in:
Dr.Lt.Data
2024-04-15 11:06:20 +09:00
parent 8ffc4efcd6
commit 971c4a37aa
2 changed files with 9 additions and 1 deletions
+1 -1
View File
@@ -2,7 +2,7 @@ import configparser
import os
version_code = [4, 87, 5]
version_code = [4, 87, 6]
version = f"V{version_code[0]}.{version_code[1]}" + (f'.{version_code[2]}' if len(version_code) > 2 else '')
dependency_version = 20
+8
View File
@@ -151,6 +151,10 @@ class SegmDetectorCombined:
def doit(self, segm_detector, image, threshold, dilation):
mask = segm_detector.detect_combined(image, threshold, dilation)
if mask is None:
mask = torch.zeros((image.shape[2], image.shape[1]), dtype=torch.float32, device="cpu")
return (mask.unsqueeze(0),)
@@ -167,6 +171,10 @@ class BboxDetectorCombined(SegmDetectorCombined):
def doit(self, bbox_detector, image, threshold, dilation):
mask = bbox_detector.detect_combined(image, threshold, dilation)
if mask is None:
mask = torch.zeros((image.shape[2], image.shape[1]), dtype=torch.float32, device="cpu")
return (mask.unsqueeze(0),)