260 lines
9.5 KiB
Python
260 lines
9.5 KiB
Python
"""
|
|
Auto-Labeling Assistant for DirtyLeague Dataset.
|
|
Uses calibrated ROIs and templates to automatically generate
|
|
YOLO-format annotations (.txt) for the collected raw screenshots in seconds.
|
|
|
|
Outputs:
|
|
- dataset/labels/<image_name>.txt (normalized YOLO bbox: class_id cx cy w h)
|
|
- dataset/annotated/<image_name>.jpg (visual verification with labeled boxes)
|
|
- dataset/data.yaml (dataset configuration for YOLO training)
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import cv2
|
|
import numpy as np
|
|
|
|
# Class definitions for DirtyLeague UI detection
|
|
CLASSES = [
|
|
"btn_fight", # 0: Yellow/Gold fight button in lobby
|
|
"btn_exit", # 1: White battle exit button (bottom bar)
|
|
"btn_leave", # 2: Red confirm surrender button
|
|
"btn_ok", # 3: Defeat screen OK button
|
|
"btn_collect", # 4: Victory / Chest collect button
|
|
"btn_turn_all", # 5: Chest turn cards over button
|
|
"btn_close", # 6: Offer / modal close 'X' button
|
|
"btn_remove_chest", # 7: Discard extra chest button
|
|
"btn_open_chest", # 8: Open chest button on no-slots screen
|
|
"crown_chest", # 9: 5/5 Crown chest banner in lobby
|
|
"card", # 10: Generic creature card / hero slot
|
|
]
|
|
|
|
CLASS_MAP = {name: i for i, name in enumerate(CLASSES)}
|
|
|
|
# Template mappings to class names
|
|
TEMPLATE_TO_CLASS = {
|
|
"btn_fight": "btn_fight",
|
|
"btn_exit": "btn_exit",
|
|
"btn_leave": "btn_leave",
|
|
"btn_ok": "btn_ok",
|
|
"btn_collect": "btn_collect",
|
|
"btn_collect_victory": "btn_collect",
|
|
"btn_turn_all_cards": "btn_turn_all",
|
|
"btn_close": "btn_close",
|
|
"btn_close_offer": "btn_close",
|
|
"btn_remove_chest": "btn_remove_chest",
|
|
"btn_open_chest": "btn_open_chest",
|
|
}
|
|
|
|
# Calibrated Relative ROIs (rx1, ry1, rx2, ry2) for lightning-fast localized search (<1ms)
|
|
TEMPLATE_ROIS = {
|
|
"btn_fight": [(0.75, 0.35, 0.95, 0.50)],
|
|
"btn_exit": [(0.40, 0.90, 0.52, 1.00)],
|
|
"btn_leave": [(0.48, 0.54, 0.64, 0.70)],
|
|
"btn_ok": [(0.20, 0.70, 0.40, 0.85)],
|
|
"btn_collect": [(0.38, 0.63, 0.60, 0.78)],
|
|
"btn_collect_victory": [(0.18, 0.73, 0.33, 0.87)],
|
|
"btn_turn_all_cards": [(0.38, 0.63, 0.60, 0.78)],
|
|
"btn_close": [(0.85, 0.00, 1.00, 0.16)],
|
|
"btn_close_offer": [(0.85, 0.00, 1.00, 0.16)],
|
|
"btn_remove_chest": [(0.33, 0.80, 0.52, 0.92)],
|
|
"btn_open_chest": [(0.46, 0.80, 0.63, 0.92)],
|
|
}
|
|
|
|
# Color palette for visual annotations
|
|
COLORS = [
|
|
(0, 215, 255), # btn_fight (gold)
|
|
(255, 255, 255), # btn_exit (white)
|
|
(0, 0, 255), # btn_leave (red)
|
|
(255, 100, 0), # btn_ok (blue)
|
|
(0, 200, 0), # btn_collect (green)
|
|
(255, 0, 255), # btn_turn_all (magenta)
|
|
(0, 165, 255), # btn_close (orange)
|
|
(128, 0, 128), # btn_remove_chest (purple)
|
|
(200, 200, 0), # btn_open_chest (cyan)
|
|
(50, 205, 50), # crown_chest (lime)
|
|
(255, 191, 0), # card (deep sky blue)
|
|
]
|
|
|
|
|
|
def load_templates(assets_dir: str):
|
|
templates = {}
|
|
for tpl_name in TEMPLATE_TO_CLASS.keys():
|
|
path = os.path.join(assets_dir, f"{tpl_name}.png")
|
|
if os.path.exists(path):
|
|
img = cv2.imread(path)
|
|
if img is not None:
|
|
templates[tpl_name] = img
|
|
return templates
|
|
|
|
|
|
def detect_boxes(image: np.ndarray, templates: dict, base_height: int = 2050):
|
|
fh, fw = image.shape[:2]
|
|
scale = fh / float(base_height)
|
|
boxes = [] # list of (class_id, x1, y1, x2, y2, conf)
|
|
|
|
# 1. Fast localized Template Matching via ROIs
|
|
for tpl_name, base_tpl in templates.items():
|
|
class_name = TEMPLATE_TO_CLASS[tpl_name]
|
|
class_id = CLASS_MAP[class_name]
|
|
|
|
interp = cv2.INTER_AREA if scale < 1.0 else cv2.INTER_CUBIC
|
|
tpl = cv2.resize(base_tpl, (0, 0), fx=scale, fy=scale, interpolation=interp)
|
|
th, tw = tpl.shape[:2]
|
|
|
|
rois = TEMPLATE_ROIS.get(tpl_name, [(0.0, 0.0, 1.0, 1.0)])
|
|
threshold = 0.55 if "close" in tpl_name else 0.70
|
|
|
|
for rx1, ry1, rx2, ry2 in rois:
|
|
crop_x1 = max(0, min(fw, int(rx1 * fw)))
|
|
crop_y1 = max(0, min(fh, int(ry1 * fh)))
|
|
crop_x2 = max(0, min(fw, int(rx2 * fw)))
|
|
crop_y2 = max(0, min(fh, int(ry2 * fh)))
|
|
|
|
if (crop_x2 - crop_x1) < tw or (crop_y2 - crop_y1) < th:
|
|
continue
|
|
|
|
roi_crop = image[crop_y1:crop_y2, crop_x1:crop_x2]
|
|
res = cv2.matchTemplate(roi_crop, tpl, cv2.TM_CCOEFF_NORMED)
|
|
min_v, max_v, min_l, max_l = cv2.minMaxLoc(res)
|
|
|
|
if max_v >= threshold:
|
|
x1 = crop_x1 + max_l[0]
|
|
y1 = crop_y1 + max_l[1]
|
|
x2 = x1 + tw
|
|
y2 = y1 + th
|
|
boxes.append((class_id, x1, y1, x2, y2, float(max_v)))
|
|
|
|
# 2. Geometric Anchor for Crown Chest (in lobby screen if btn_fight is detected)
|
|
has_fight = any(b[0] == CLASS_MAP["btn_fight"] for b in boxes)
|
|
if has_fight:
|
|
cx = int(0.2279 * fw)
|
|
cy = int(0.6146 * fh)
|
|
cw = int(0.12 * fw)
|
|
ch = int(0.14 * fh)
|
|
x1, y1 = cx - cw // 2, cy - ch // 2
|
|
x2, y2 = cx + cw // 2, cy + ch // 2
|
|
boxes.append((CLASS_MAP["crown_chest"], x1, y1, x2, y2, 0.95))
|
|
|
|
# 3. Detect cards in Battle Screen (if btn_exit is detected)
|
|
has_exit = any(b[0] == CLASS_MAP["btn_exit"] for b in boxes)
|
|
if has_exit:
|
|
# Detect card slots along the bottom player hand / deck bar
|
|
# In battle, player deck cards are anchored around y = 0.78..0.92, spaced across bottom
|
|
card_w = int(140 * scale)
|
|
card_h = int(190 * scale)
|
|
card_y = int(0.85 * fh)
|
|
card_xs = [int(x_ratio * fw) for x_ratio in [0.22, 0.32, 0.68, 0.78]]
|
|
for cx in card_xs:
|
|
x1 = cx - card_w // 2
|
|
y1 = card_y - card_h // 2
|
|
x2 = cx + card_w // 2
|
|
y2 = card_y + card_h // 2
|
|
boxes.append((CLASS_MAP["card"], x1, y1, x2, y2, 0.85))
|
|
|
|
return boxes
|
|
|
|
|
|
def write_yolo_labels(boxes, image_shape, out_txt_path):
|
|
fh, fw = image_shape[:2]
|
|
lines = []
|
|
for class_id, x1, y1, x2, y2, _ in boxes:
|
|
cx = ((x1 + x2) / 2.0) / fw
|
|
cy = ((y1 + y2) / 2.0) / fh
|
|
w = (x2 - x1) / float(fw)
|
|
h = (y2 - y1) / float(fh)
|
|
lines.append(f"{class_id} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}\n")
|
|
|
|
with open(out_txt_path, "w", encoding="utf-8") as f:
|
|
f.writelines(lines)
|
|
|
|
|
|
def draw_annotations(image, boxes):
|
|
annotated = image.copy()
|
|
for class_id, x1, y1, x2, y2, conf in boxes:
|
|
color = COLORS[class_id % len(COLORS)]
|
|
cv2.rectangle(annotated, (x1, y1), (x2, y2), color, 3)
|
|
|
|
label = f"{CLASSES[class_id]} {conf:.2f}"
|
|
(tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)
|
|
cv2.rectangle(annotated, (x1, y1 - th - 10), (x1 + tw + 6, y1), color, -1)
|
|
cv2.putText(annotated, label, (x1 + 3, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 2)
|
|
return annotated
|
|
|
|
|
|
def main():
|
|
root_dir = os.path.dirname(__file__)
|
|
raw_dir = os.path.join(root_dir, "dataset", "raw")
|
|
labels_dir = os.path.join(root_dir, "dataset", "labels")
|
|
annotated_dir = os.path.join(root_dir, "dataset", "annotated")
|
|
assets_dir = os.path.join(root_dir, "assets")
|
|
|
|
os.makedirs(labels_dir, exist_ok=True)
|
|
os.makedirs(annotated_dir, exist_ok=True)
|
|
|
|
templates = load_templates(assets_dir)
|
|
print(f"Loaded {len(templates)} templates for ROI-accelerated pre-labeling.")
|
|
|
|
image_files = sorted([f for f in os.listdir(raw_dir) if f.endswith(".png")])
|
|
print(f"Found {len(image_files)} raw images in '{raw_dir}' to process.")
|
|
|
|
total_boxes = 0
|
|
annotated_images_count = 0
|
|
|
|
for i, fname in enumerate(image_files, 1):
|
|
fpath = os.path.join(raw_dir, fname)
|
|
img = cv2.imread(fpath)
|
|
if img is None:
|
|
continue
|
|
|
|
boxes = detect_boxes(img, templates)
|
|
base_name = os.path.splitext(fname)[0]
|
|
|
|
# 1. Save YOLO label .txt
|
|
txt_path = os.path.join(labels_dir, f"{base_name}.txt")
|
|
write_yolo_labels(boxes, img.shape, txt_path)
|
|
|
|
# 2. Save annotated verification image
|
|
annotated_img = draw_annotations(img, boxes)
|
|
preview_h = 1080
|
|
preview_w = int(img.shape[1] * (preview_h / img.shape[0]))
|
|
preview_img = cv2.resize(annotated_img, (preview_w, preview_h), interpolation=cv2.INTER_AREA)
|
|
jpg_path = os.path.join(annotated_dir, f"{base_name}.jpg")
|
|
cv2.imwrite(jpg_path, preview_img, [int(cv2.IMWRITE_JPEG_QUALITY), 90])
|
|
|
|
total_boxes += len(boxes)
|
|
if len(boxes) > 0:
|
|
annotated_images_count += 1
|
|
|
|
found_classes = ", ".join(sorted(set(CLASSES[b[0]] for b in boxes))) if boxes else "None"
|
|
print(f"[{i:02d}/{len(image_files)}] {fname} -> {len(boxes)} boxes ({found_classes})")
|
|
|
|
# 3. Generate data.yaml
|
|
yaml_path = os.path.join(root_dir, "dataset", "data.yaml")
|
|
yaml_content = f"""# DirtyLeague YOLOv8 Dataset Configuration
|
|
path: {os.path.abspath(os.path.join(root_dir, 'dataset'))}
|
|
train: images/train
|
|
val: images/val
|
|
|
|
names:
|
|
"""
|
|
for idx, cname in enumerate(CLASSES):
|
|
yaml_content += f" {idx}: {cname}\n"
|
|
|
|
with open(yaml_path, "w", encoding="utf-8") as f:
|
|
f.write(yaml_content)
|
|
|
|
print("\n" + "=" * 65)
|
|
print(f" AUTO-LABELING COMPLETE!")
|
|
print(f" Images processed : {len(image_files)}")
|
|
print(f" Images with boxes : {annotated_images_count} / {len(image_files)}")
|
|
print(f" Total boxes placed : {total_boxes}")
|
|
print(f" Labels directory : {labels_dir}")
|
|
print(f" Visual previews : {annotated_dir}")
|
|
print(f" Config file : {yaml_path}")
|
|
print("=" * 65)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|