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

feat: add debug dino and mask dilate and erode

Browse Source
This commit is contained in:
Manuel Schmid
2024-06-09 22:31:41 +02:00
parent f2e7b65ed3
commit 57c049858c
3 changed files with 38 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extras/inpaint_mask.py
+18 -27
View File
@@ -1,26 +1,9 @@
from PIL import Image
import numpy as np
import torch
from rembg import remove, new_session
from extras.GroundingDINO.util.inference import default_groundingdino
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
def run_grounded_sam(input_image, text_prompt, box_threshold, text_threshold):
# run grounding dino model
detections, _, _, _ = default_groundingdino(
image=np.array(input_image),
caption=text_prompt,
box_threshold=box_threshold,
text_threshold=text_threshold
)
return detections.xyxy
def generate_mask_from_image(image, mask_model, extras, box_erode_or_dilate: int=0):
def generate_mask_from_image(image: np.ndarray, mask_model: str, extras: dict, box_erode_or_dilate: int=0, debug_dino: bool=False) -> np.ndarray | None:
if image is None:
return
@@ -28,29 +11,37 @@ def generate_mask_from_image(image, mask_model, extras, box_erode_or_dilate: int
image = image['image']
if mask_model == 'sam':
img = Image.fromarray(image)
boxes = run_grounded_sam(img, extras['sam_prompt_text'], box_threshold=extras['box_threshold'], text_threshold=extras['text_threshold'])
detections, _, _, _ = default_groundingdino(
image=image,
caption=extras['sam_prompt_text'],
box_threshold=extras['box_threshold'],
text_threshold=extras['text_threshold']
)
detection_boxes = detections.xyxy
# use full image if no box has been found
boxes = np.array([[0, 0, image.shape[1], image.shape[0]]]) if len(boxes) == 0 else boxes
detection_boxes = np.array([[0, 0, image.shape[1], image.shape[0]]]) if len(detection_boxes) == 0 else detection_boxes
extras['sam_prompt'] = []
# from PIL import ImageDraw
# draw = ImageDraw.Draw(img)
for idx, box in enumerate(boxes):
for idx, box in enumerate(detection_boxes):
box_list = box.tolist()
if box_erode_or_dilate != 0:
box_list[0] -= box_erode_or_dilate
box_list[1] -= box_erode_or_dilate
box_list[2] += box_erode_or_dilate
box_list[3] += box_erode_or_dilate
# draw.rectangle(box_list, fill=128, outline ="red")
extras['sam_prompt'] += [{"type": "rectangle", "data": box_list}]
# img.show()
if debug_dino:
from PIL import ImageDraw, Image
image_with_boxes = Image.new("RGB", (image.shape[1], image.shape[0]), color="black")
draw = ImageDraw.Draw(image_with_boxes)
for box in extras['sam_prompt']:
draw.rectangle(box['data'], fill="white")
return np.array(image_with_boxes)
return remove(
image,
session=new_session(mask_model, **extras),
only_mask=True,
# post_process_mask=True,
**extras
)