Foreign-Mirrors/Fooocus
1
0
Fork 0
mirror of https://github.com/lllyasviel/Fooocus.git synced 2026-08-16 13:13:16 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix: correctly count sam masks

Browse Source
This commit is contained in:
Manuel Schmid
2024-06-16 15:57:53 +02:00
parent ef9fd293ff
commit b585d9dfa7
1 changed files with 2 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extras/inpaint_mask.py
+2 -4
View File
@@ -85,10 +85,7 @@ def generate_mask_from_image(image: np.ndarray, mask_model: str = 'sam', extras=
sam_predictor = SamPredictor(sam) sam_predictor = SamPredictor(sam)
final_mask_tensor = torch.zeros((image.shape[0], image.shape[1])) final_mask_tensor = torch.zeros((image.shape[0], image.shape[1]))
dino_detection_count = boxes.size(0) dino_detection_count = boxes.size(0)
sam_detection_count = 0
sam_detection_on_mask_count = 0
if dino_detection_count > 0: if dino_detection_count > 0:
sam_predictor.set_image(image) 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) masks = optimize_masks(masks)
sam_detection_count = len(masks)
sam_objects = min(len(logits), sam_options.max_num_boxes) sam_objects = min(len(logits), sam_options.max_num_boxes)
for obj_ind in range(sam_objects): for obj_ind in range(sam_objects):
mask_tensor = masks[obj_ind][0] mask_tensor = masks[obj_ind][0]
final_mask_tensor += mask_tensor final_mask_tensor += mask_tensor
sam_detection_on_mask_count += 1
final_mask_tensor = (final_mask_tensor > 0).to('cpu').numpy() final_mask_tensor = (final_mask_tensor > 0).to('cpu').numpy()
mask_image = np.dstack((final_mask_tensor, final_mask_tensor, final_mask_tensor)) * 255 mask_image = np.dstack((final_mask_tensor, final_mask_tensor, final_mask_tensor)) * 255