feat: add box_erode_or_dilate to generate_mask_from_image, expose more sam return values

This commit is contained in:
Manuel Schmid
2024-06-09 18:45:13 +02:00
parent 9affa32583
commit ff9fa6c837
2 changed files with 21 additions and 7 deletions
+2 -2
View File
@@ -25,7 +25,7 @@ class GroundingDinoModel(Model):
caption: str,
box_threshold: float = 0.35,
text_threshold: float = 0.25
) -> Tuple[sv.Detections, List[str]]:
) -> Tuple[sv.Detections, torch.Tensor, torch.Tensor, List[str]]:
if self.model is None:
filename = load_file_from_url(
url="https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth",
@@ -56,7 +56,7 @@ class GroundingDinoModel(Model):
source_w=source_w,
boxes=boxes,
logits=logits)
return detections, phrases
return detections, boxes, logits, phrases
def predict(