happyz
/
Fooocus
mirror of https://github.com/lllyasviel/Fooocus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix: correctly count sam masks

This commit is contained in:
Manuel Schmid 2024-06-16 15:57:53 +02:00
parent ef9fd293ff
commit b585d9dfa7
No known key found for this signature in database
GPG Key ID: 32C4F7569B40B84B
1 changed files with 2 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
extras/inpaint_mask.py
View File

@ -85,10 +85,7 @@ def generate_mask_from_image(image: np.ndarray, mask_model: str = 'sam', extras=
sam_predictor = SamPredictor(sam)
final_mask_tensor = torch.zeros((image.shape[0], image.shape[1]))
dino_detection_count = boxes.size(0)
sam_detection_count = 0
sam_detection_on_mask_count = 0
if dino_detection_count > 0:
sam_predictor.set_image(image)
@ -118,11 +115,12 @@ def generate_mask_from_image(image: np.ndarray, mask_model: str = 'sam', extras=
)
masks = optimize_masks(masks)
sam_detection_count = len(masks)
sam_objects = min(len(logits), sam_options.max_num_boxes)
for obj_ind in range(sam_objects):
mask_tensor = masks[obj_ind][0]
final_mask_tensor += mask_tensor
sam_detection_on_mask_count += 1
final_mask_tensor = (final_mask_tensor > 0).to('cpu').numpy()
mask_image = np.dstack((final_mask_tensor, final_mask_tensor, final_mask_tensor)) * 255