49 lines
1.7 KiB
Python
49 lines
1.7 KiB
Python
from transformers import pipeline as hf_pipeline
|
|
from PIL import Image, ImageDraw
|
|
import numpy as np
|
|
import torch
|
|
import json
|
|
|
|
object_detection_model_name_list = [
|
|
"facebook/detr-resnet-50",
|
|
"hustvl/yolos-tiny",
|
|
]
|
|
|
|
|
|
class ObjectDetectionPipeline:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"model_name": (object_detection_model_name_list, {"default": object_detection_model_name_list[0]}),
|
|
"threshold": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING",)
|
|
RETURN_NAMES = ("annotated_image", "detections_json",)
|
|
FUNCTION = "object_detection_pipeline"
|
|
CATEGORY = "Transformers/ComputerVision/ObjectDetection"
|
|
|
|
def object_detection_pipeline(self, image, model_name, threshold):
|
|
img = 255.0 * image[0].cpu().numpy()
|
|
pil_image = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8))
|
|
|
|
pipe = hf_pipeline("object-detection", model=model_name)
|
|
results = pipe(pil_image, threshold=threshold)
|
|
|
|
draw = ImageDraw.Draw(pil_image)
|
|
for det in results:
|
|
box = det["box"]
|
|
x, y, x2, y2 = box["xmin"], box["ymin"], box["xmax"], box["ymax"]
|
|
draw.rectangle((x, y, x2, y2), outline="red", width=2)
|
|
label_text = f"{det['label']} ({det['score']:.2f})"
|
|
draw.text((x, y), label_text, fill="white")
|
|
|
|
arr = np.array(pil_image).astype(np.float32) / 255.0
|
|
tensor_image = torch.from_numpy(arr)[None,]
|
|
detections_json = json.dumps(results, indent=2, default=str)
|
|
|
|
return (tensor_image, detections_json,)
|