adt2 / adetailer /common.py
ddoc's picture
Upload 36 files
187d856
from __future__ import annotations
from collections import OrderedDict
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional, Union
from huggingface_hub import hf_hub_download
from PIL import Image, ImageDraw
from rich import print
repo_id = "Bingsu/adetailer"
@dataclass
class PredictOutput:
bboxes: list[list[int | float]] = field(default_factory=list)
masks: list[Image.Image] = field(default_factory=list)
preview: Optional[Image.Image] = None
def hf_download(file: str):
try:
path = hf_hub_download(repo_id, file)
except Exception:
msg = f"[-] ADetailer: Failed to load model {file!r} from huggingface"
print(msg)
path = "INVALID"
return path
def get_models(
model_dir: Union[str, Path], huggingface: bool = True
) -> OrderedDict[str, Optional[str]]:
model_dir = Path(model_dir)
if model_dir.is_dir():
model_paths = [
p
for p in model_dir.rglob("*")
if p.is_file() and p.suffix in (".pt", ".pth")
]
else:
model_paths = []
models = OrderedDict()
if huggingface:
models.update(
{
"face_yolov8n.pt": hf_download("face_yolov8n.pt"),
"face_yolov8s.pt": hf_download("face_yolov8s.pt"),
"hand_yolov8n.pt": hf_download("hand_yolov8n.pt"),
"person_yolov8n-seg.pt": hf_download("person_yolov8n-seg.pt"),
"person_yolov8s-seg.pt": hf_download("person_yolov8s-seg.pt"),
}
)
models.update(
{
"mediapipe_face_full": None,
"mediapipe_face_short": None,
"mediapipe_face_mesh": None,
"mediapipe_face_mesh_eyes_only": None,
}
)
invalid_keys = [k for k, v in models.items() if v == "INVALID"]
for key in invalid_keys:
models.pop(key)
for path in model_paths:
if path.name in models:
continue
models[path.name] = str(path)
return models
def create_mask_from_bbox(
bboxes: list[list[float]], shape: tuple[int, int]
) -> list[Image.Image]:
"""
Parameters
----------
bboxes: list[list[float]]
list of [x1, y1, x2, y2]
bounding boxes
shape: tuple[int, int]
shape of the image (width, height)
Returns
-------
masks: list[Image.Image]
A list of masks
"""
masks = []
for bbox in bboxes:
mask = Image.new("L", shape, 0)
mask_draw = ImageDraw.Draw(mask)
mask_draw.rectangle(bbox, fill=255)
masks.append(mask)
return masks
def create_bbox_from_mask(
masks: list[Image.Image], shape: tuple[int, int]
) -> list[list[int]]:
"""
Parameters
----------
masks: list[Image.Image]
A list of masks
shape: tuple[int, int]
shape of the image (width, height)
Returns
-------
bboxes: list[list[float]]
A list of bounding boxes
"""
bboxes = []
for mask in masks:
mask = mask.resize(shape)
bbox = mask.getbbox()
if bbox is not None:
bboxes.append(list(bbox))
return bboxes