mirror of
https://github.com/vladmandic/automatic
synced 2026-09-06 21:10:45 +02:00
Vladimir Mandic
114 lines
4.2 KiB
Python
114 lines
4.2 KiB
Python
import os
|
|
import re
|
|
from PIL import Image, ImageDraw
|
|
from modules.logger import log
|
|
|
|
|
|
class_tag_re = re.compile(r'^\[class\s*=\s*([^\]]+)\]\s*(.*)$', re.IGNORECASE)
|
|
|
|
|
|
def list_models(self):
|
|
from modules.detailer import detailer_models
|
|
from modules import shared
|
|
self.list.clear()
|
|
files = []
|
|
downloaded = 0
|
|
for m in detailer_models:
|
|
name = os.path.splitext(os.path.basename(m))[0]
|
|
self.list[name] = m
|
|
files.append(name)
|
|
if os.path.exists(shared.opts.yolo_dir):
|
|
for f in os.listdir(shared.opts.yolo_dir):
|
|
if f.endswith('.pt'):
|
|
downloaded += 1
|
|
name = os.path.splitext(os.path.basename(f))[0]
|
|
if name not in files:
|
|
self.list[name] = os.path.join(shared.opts.yolo_dir, f)
|
|
log.info(f'Available Detailer: path="{shared.opts.yolo_dir}" items={len(list(self.list))} downloaded={downloaded}')
|
|
return list(self.list)
|
|
|
|
|
|
def detailer_opt(p, attr, opts_attr=None):
|
|
"""Read detailer param from processing object if set, otherwise fall back to shared.opts."""
|
|
from modules import shared
|
|
if p is not None:
|
|
val = getattr(p, attr, None)
|
|
if val is not None:
|
|
return val
|
|
return getattr(shared.opts, opts_attr or attr, None)
|
|
|
|
|
|
def parse_prompt_lines(text: str):
|
|
"""Split a detailer prompt into class-tagged templates and positional fallback lines.
|
|
|
|
A line starting with '[CLASS=name]' or '[CLASS=name1,name2]' assigns its text to every
|
|
detection whose label matches one of the given class names (case-insensitive). All
|
|
other non-empty lines are kept, in order, as the legacy positional fallback used for
|
|
detections that don't match any class tag.
|
|
"""
|
|
class_map: dict[str, str] = {}
|
|
fallback: list[str] = []
|
|
for line in (text or '').split('\n'):
|
|
line = line.strip()
|
|
if len(line) == 0:
|
|
continue # blank spacer lines don't count as a fallback entry
|
|
m = class_tag_re.match(line)
|
|
if m:
|
|
names = [n.strip().lower() for n in m.group(1).split(',') if n.strip()]
|
|
for name in names:
|
|
class_map[name] = m.group(2).strip()
|
|
else:
|
|
fallback.append(line)
|
|
return class_map, fallback
|
|
|
|
|
|
def assign_prompts(text: str, items: list) -> list[str]:
|
|
"""Resolve a detailer prompt/negative-prompt string into one entry per detection.
|
|
|
|
Detections whose YOLO label matches a '[CLASS=name]' tag get that tag's text.
|
|
Remaining detections fall back to the untagged lines, applied positionally in
|
|
detection order and cycling if there are more detections than fallback lines
|
|
(matching prior behavior when no class tags are used).
|
|
"""
|
|
class_map, fallback = parse_prompt_lines(text)
|
|
if len(fallback) == 0:
|
|
fallback = ['']
|
|
resolved = []
|
|
fallback_idx = 0
|
|
for item in items:
|
|
label = (getattr(item, 'label', None) or '').strip().lower()
|
|
if label in class_map:
|
|
resolved.append(class_map[label])
|
|
else:
|
|
resolved.append(fallback[fallback_idx % len(fallback)])
|
|
fallback_idx += 1
|
|
return resolved
|
|
|
|
|
|
def get_mask(box: list[int], image: Image.Image, include_mask: bool = True) -> tuple[Image.Image | None, Image.Image]:
|
|
cropped = image.crop(box)
|
|
if not include_mask:
|
|
return None, cropped
|
|
mask = Image.new('L', image.size, 0)
|
|
draw_mask = ImageDraw.Draw(mask)
|
|
draw_mask.rectangle(box, fill="white", outline=None, width=0)
|
|
return mask, cropped
|
|
|
|
|
|
class DetailerResult:
|
|
def __init__(self, cls: int, label: str, score: float, box: list[int], mask: Image.Image = None, item: Image.Image = None, width = 0, height = 0, args = None):
|
|
if args is None:
|
|
args = {}
|
|
self.cls = cls
|
|
self.label = label
|
|
self.score = score
|
|
self.box = box
|
|
self.mask = mask
|
|
self.item = item
|
|
self.width = width if width > 0 else box[2] - box[0]
|
|
self.height = height if height > 0 else box[3] - box[1]
|
|
self.args = args
|
|
|
|
def __str__(self):
|
|
return f'DetailerResult(cls={self.cls} label="{self.label}" score={self.score:.2f} box={self.box} size={self.width}x{self.height} args={self.args})'
|