231 lines
7.4 KiB
Python
231 lines
7.4 KiB
Python
import os
|
||
import sys
|
||
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
|
||
|
||
import folder_paths
|
||
import json
|
||
import cv2
|
||
|
||
import logging
|
||
import numpy as np
|
||
import torch
|
||
from PIL import Image, ImageDraw, ImageFont
|
||
from urllib.parse import urlparse
|
||
from torch.hub import download_url_to_file
|
||
|
||
from omdet.infernece.det_engine import DetEngine
|
||
from omdet.utils.plots import Annotator
|
||
|
||
logger = logging.getLogger("ComfyUI-OmDet")
|
||
|
||
OmDet_model_list = {
|
||
"OmDet-Turbo_tiny_SWIN_T.pth": {
|
||
"model_url": "https://huggingface.co/omlab/OmDet-Turbo_tiny_SWIN_T/resolve/main/OmDet-Turbo_tiny_SWIN_T.pth?download=true",
|
||
},
|
||
"ViT-B-16.pt": {
|
||
"model_url": "https://huggingface.co/omlab/OmDet-Turbo_tiny_SWIN_T/resolve/main/ViT-B-16.pt?download=true",
|
||
},
|
||
}
|
||
|
||
def download_model(model_url, model_file_path):
|
||
try:
|
||
torch.hub.download_url_to_file(model_url, model_file_path)
|
||
except Exception as e:
|
||
print(f"Failed to download model from {model_url}. Error: {e}")
|
||
|
||
def check_and_download_models():
|
||
model_path = os.path.join(folder_paths.models_dir, "OmDet")
|
||
if not os.path.exists(model_path):
|
||
os.makedirs(model_path)
|
||
|
||
for model_name, model_info in OmDet_model_list.items():
|
||
model_filename = model_name
|
||
model_file_path = os.path.join(model_path, model_filename)
|
||
|
||
if not os.path.exists(model_file_path):
|
||
print(f"{model_filename} not found. Downloading from {model_info['model_url']}...")
|
||
download_model(model_info["model_url"], model_file_path)
|
||
|
||
# modified from https://stackoverflow.com/questions/22058048/hashing-a-file-in-python
|
||
def calculate_file_hash(filename: str, hash_every_n: int = 1):
|
||
# Larger video files were taking >.5 seconds to hash even when cached,
|
||
# so instead the modified time from the filesystem is used as a hash
|
||
h = hashlib.sha256()
|
||
h.update(filename.encode())
|
||
h.update(str(os.path.getmtime(filename)).encode())
|
||
return h.hexdigest()
|
||
|
||
def plot_boxes_to_image(image_pil, tgt):
|
||
H, W = tgt["size"]
|
||
result = tgt["result"]
|
||
|
||
res_mask = []
|
||
res_image = []
|
||
|
||
box_color = (255, 0, 0) # Red color for the box
|
||
text_color = (255, 255, 255) # White color for the text
|
||
|
||
draw = ImageDraw.Draw(image_pil)
|
||
|
||
# Get the current file path and use it to create a relative path to the font file
|
||
current_file_path = os.path.dirname(os.path.abspath(__file__))
|
||
font_path = os.path.join(current_file_path, "docs", "PingFangRegular.ttf")
|
||
font_size = 20
|
||
font = ImageFont.truetype(font_path, font_size)
|
||
|
||
labelme_data = {
|
||
"version": "4.5.6",
|
||
"flags": {},
|
||
"shapes": [],
|
||
"imagePath": None,
|
||
"imageData": None,
|
||
"imageHeight": H,
|
||
"imageWidth": W,
|
||
}
|
||
for item in result:
|
||
x1,y1,x2,y2 = item['xmin'],item['ymin'],item['xmax'],item['ymax']
|
||
label = item["label"]
|
||
threshold = round(item["conf"], 2)
|
||
|
||
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
|
||
points = [[x1, y1], [x2, y2]]
|
||
|
||
# Save labelme json
|
||
shape = {
|
||
"label": label,
|
||
"points": points,
|
||
"group_id": None,
|
||
"shape_type": "rectangle",
|
||
"flags": {},
|
||
}
|
||
labelme_data["shapes"].append(shape)
|
||
|
||
# Change label
|
||
label = label + ":" + str(threshold)
|
||
shape["threshold"] = str(threshold)
|
||
|
||
# Draw rectangle on the image using PIL
|
||
draw.rectangle([(x1, y1), (x2, y2)], outline=box_color, width=3)
|
||
|
||
# Draw label on the image using PIL
|
||
text_size = draw.textsize(label, font=font)
|
||
label_ymin = max(y1, text_size[1] + 10)
|
||
draw.rectangle([(x1, y1 - text_size[1] - 10), (x1 + text_size[0], y1)], fill=box_color)
|
||
draw.text((x1, y1 - text_size[1] - 10), label, font=font, fill=text_color)
|
||
|
||
# Draw mask
|
||
mask = np.zeros((H, W, 1), dtype=np.uint8)
|
||
cv2.rectangle(mask, (int(x1), int(y1)), (int(x2), int(y2)), (255, 255, 255), -1)
|
||
mask_tensor = torch.from_numpy(mask).permute(2, 0, 1).float() / 255.0
|
||
res_mask.append(mask_tensor)
|
||
|
||
if len(res_mask) == 0:
|
||
mask = np.zeros((H, W, 1), dtype=np.uint8)
|
||
mask_tensor = torch.from_numpy(mask).permute(2, 0, 1).float() / 255.0
|
||
res_mask.append(mask_tensor)
|
||
|
||
# Convert the PIL image back to a numpy array
|
||
image_with_boxes = np.array(image_pil)
|
||
|
||
# Convert the modified image to a torch tensor
|
||
image_with_boxes_tensor = torch.from_numpy(image_with_boxes.astype(np.float32) / 255.0)
|
||
image_with_boxes_tensor = torch.unsqueeze(image_with_boxes_tensor, 0)
|
||
res_image.append(image_with_boxes_tensor)
|
||
|
||
return res_image, res_mask, labelme_data
|
||
|
||
def OmDet_detect(data, labels, conf_threshold=0.30, nms_threshold=0.5,device="cuda"):
|
||
prompt = "Detect {}.".format(",".join(labels))
|
||
engine = DetEngine(batch_size=1, device=device)
|
||
res = engine.inf_predict(
|
||
"OmDet-Turbo_tiny_SWIN_T",
|
||
task=prompt,
|
||
data=data,
|
||
labels=labels,
|
||
src_type="pil", # type of the image_paths, "local/tensor/pil/"url"
|
||
conf_threshold=conf_threshold,
|
||
nms_threshold=nms_threshold,
|
||
)
|
||
return res
|
||
|
||
def get_labels(input_string):
|
||
modified_string = input_string.replace(", ", ",").replace("，", ",")
|
||
return modified_string.split(",")
|
||
|
||
class ApplyOmDet:
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"image": ("IMAGE",),
|
||
"prompt": (
|
||
"STRING",
|
||
{"default": "person,cat,dog"},
|
||
),
|
||
"conf_threshold": (
|
||
"FLOAT",
|
||
{"default": 0.3, "min": 0, "max": 1.0, "step": 0.01},
|
||
),
|
||
"nms_threshold": (
|
||
"FLOAT",
|
||
{"default": 0.5, "min": 0, "max": 1.0, "step": 0.01},
|
||
),
|
||
"device": (["cuda", "cpu"],),
|
||
},
|
||
}
|
||
|
||
CATEGORY = "ComfyUI-OmDet"
|
||
FUNCTION = "main"
|
||
RETURN_TYPES = (
|
||
"IMAGE",
|
||
"MASK",
|
||
"JSON",
|
||
)
|
||
|
||
def main(
|
||
self,
|
||
image,
|
||
prompt,
|
||
conf_threshold,
|
||
nms_threshold,
|
||
device
|
||
):
|
||
check_and_download_models()
|
||
|
||
data = []
|
||
labels = get_labels(prompt)
|
||
for item in image:
|
||
image_pil = Image.fromarray(np.clip(255.0 * item.cpu().numpy(), 0, 255).astype(np.uint8)).convert("RGB")
|
||
data.append(image_pil)
|
||
results = OmDet_detect(data, labels, conf_threshold, nms_threshold, device)
|
||
|
||
res_images = []
|
||
res_masks = []
|
||
res_labels = []
|
||
|
||
for image_pil, result in zip(data, results):
|
||
size = image_pil.size
|
||
pred_dict = {
|
||
"size": [size[1], size[0]],
|
||
"result":result
|
||
}
|
||
|
||
image_tensor, mask_tensor, labelme_data = plot_boxes_to_image(image_pil, pred_dict)
|
||
|
||
res_images.extend(image_tensor)
|
||
res_masks.extend(mask_tensor)
|
||
res_labels.append(labelme_data)
|
||
|
||
if len(res_images) == 0:
|
||
res_images.extend(item)
|
||
if len(res_masks) == 0:
|
||
mask = np.zeros((height, width, 1), dtype=np.uint8)
|
||
empty_mask = torch.from_numpy(mask).permute(2, 0, 1).float() / 255.0
|
||
res_masks.extend(empty_mask)
|
||
|
||
return (
|
||
torch.cat(res_images, dim=0),
|
||
torch.cat(res_masks, dim=0),
|
||
res_labels,
|
||
)
|