194 lines
10 KiB
Python
194 lines
10 KiB
Python
|
# -*- coding: utf-8 -*-
|
|||
|
|
# 2026-10-01 PD 지시 — "우리 게임의 장비 아이콘 리소스를 기준으로 moba풍으로 다시 칠하도록 … 로라로 학습 가능한 수준으로 이미지를 뽑아내고
|
|||
|
|
# 합격 기준에 드는 것들을 통일감 있게 만들어. 그리고 나서 최종 결과물을 학습시켜."
|
|||
|
|
# 1단계(이 스크립트): FR 장비 아이콘 원본(형태)을 입력으로 넣고 mobaGameEquipment_v10 이 그 위에 MOBA 풍을 다시 칠한다(이미지 투 이미지).
|
|||
|
|
# 형태는 원본이 정하므로 글자 프롬프트만 쓸 때의 형태 실패가 없다. 다시 칠하는 세기(denoise) 2종 × 시드 3 = 아이콘당 6 후보.
|
|||
|
|
# 배경 제거 뒤 자동 채점: 원본 실루엣과의 겹침(IoU) · 덩어리 하나 · 가장자리 안 닿음. 아이콘마다 점수순으로 한 줄(원본 | 후보)을 만든다.
|
|||
|
|
# python comfy_moba_repaint.py <출력폴더> <round 이름> [시작] [끝] [세기들 "0.5,0.65"] [시드 수=3]
|
|||
|
|
import io, json, os, re, sys, time, uuid, urllib.request, urllib.parse
|
|||
|
|
from collections import deque
|
|||
|
|
from PIL import Image, ImageDraw
|
|||
|
|
|
|||
|
|
HOST = "http://127.0.0.1:8000"
|
|||
|
|
CKPT = "mobaGameEquipment_v10.safetensors"
|
|||
|
|
SRC = r"C:\BurningTimes\FirstRising\Assets\GameAssets\UI\EquipmentItemIcon" # FR 장비 아이콘 원본(1024 · 투명 배경)
|
|||
|
|
PREFIX = "(masterpiece:1.2), (best quality:1.2), moba game icon"
|
|||
|
|
NEG = ("lowres, text, error, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, "
|
|||
|
|
"1girl, 1boy, person, human, character, face, "
|
|||
|
|
"circle, round frame, border, badge, emblem, aura, halo, light rays, glowing background, explosion, multiple objects, duplicate")
|
|||
|
|
RES = 512
|
|||
|
|
BG = (16, 16, 20) # 입력 바탕 — 어두운 단색(배경 제거가 깔끔하고, 모델이 원형 바탕을 덜 만든다)
|
|||
|
|
FIT = 0.80 # 캔버스에서 아이콘이 차지하는 비율(여백 확보 — 합격 기준 "전체가 여백 안에")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def load_captions():
|
|||
|
|
"""앞선 학습 스크립트(comfy_lora_train.py)의 CAP 표(아이콘 번호 → 내용 설명)를 그대로 읽는다 — 표를 두 벌로 두지 않는다."""
|
|||
|
|
text = io.open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "comfy_lora_train.py"), encoding="utf-8").read()
|
|||
|
|
m = re.search(r"CAP = \{(.*?)\n\}", text, re.S)
|
|||
|
|
return eval("{" + m.group(1) + "}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def post(path, data):
|
|||
|
|
req = urllib.request.Request(HOST + path, data=json.dumps(data).encode("utf-8"), headers={"Content-Type": "application/json"})
|
|||
|
|
return json.loads(urllib.request.urlopen(req, timeout=60).read())
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get(path):
|
|||
|
|
return urllib.request.urlopen(HOST + path, timeout=300).read()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def upload(image, name):
|
|||
|
|
"""ComfyUI 입력 폴더로 그림을 올린다(multipart) — LoadImage 가 이 이름으로 읽는다."""
|
|||
|
|
buf = io.BytesIO(); image.save(buf, "PNG"); data = buf.getvalue()
|
|||
|
|
boundary = "----bt" + uuid.uuid4().hex
|
|||
|
|
body = b"".join([
|
|||
|
|
("--%s\r\nContent-Disposition: form-data; name=\"image\"; filename=\"%s\"\r\nContent-Type: image/png\r\n\r\n" % (boundary, name)).encode(), data, b"\r\n",
|
|||
|
|
("--%s\r\nContent-Disposition: form-data; name=\"overwrite\"\r\n\r\ntrue\r\n" % boundary).encode(),
|
|||
|
|
("--%s--\r\n" % boundary).encode()])
|
|||
|
|
req = urllib.request.Request(HOST + "/upload/image", data=body, headers={"Content-Type": "multipart/form-data; boundary=" + boundary})
|
|||
|
|
return json.loads(urllib.request.urlopen(req, timeout=60).read())["name"]
|
|||
|
|
|
|||
|
|
|
|||
|
|
def workflow(image_name, prompt, seed, denoise):
|
|||
|
|
return {
|
|||
|
|
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": CKPT}},
|
|||
|
|
"3": {"class_type": "CLIPSetLastLayer", "inputs": {"clip": ["1", 1], "stop_at_clip_layer": -2}},
|
|||
|
|
"4": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["3", 0], "text": prompt}},
|
|||
|
|
"5": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["3", 0], "text": NEG}},
|
|||
|
|
"20": {"class_type": "LoadImage", "inputs": {"image": image_name}},
|
|||
|
|
"21": {"class_type": "VAEEncode", "inputs": {"pixels": ["20", 0], "vae": ["1", 2]}},
|
|||
|
|
"7": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "positive": ["4", 0], "negative": ["5", 0], "latent_image": ["21", 0], "seed": seed, "steps": 30, "cfg": 7.0, "sampler_name": "dpmpp_2m", "scheduler": "karras", "denoise": denoise}},
|
|||
|
|
"8": {"class_type": "VAEDecode", "inputs": {"samples": ["7", 0], "vae": ["1", 2]}},
|
|||
|
|
"12": {"class_type": "InspyrenetRembg", "inputs": {"image": ["8", 0], "torchscript_jit": "default"}},
|
|||
|
|
"9": {"class_type": "SaveImage", "inputs": {"images": ["12", 0], "filename_prefix": "moba_repaint"}},
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
|
|||
|
|
def run(image_name, prompt, seed, denoise):
|
|||
|
|
pid = post("/prompt", {"prompt": workflow(image_name, prompt, seed, denoise), "client_id": str(uuid.uuid4())})["prompt_id"]
|
|||
|
|
t0 = time.time()
|
|||
|
|
while time.time() - t0 < 600:
|
|||
|
|
hist = json.loads(get("/history/" + pid))
|
|||
|
|
if pid in hist:
|
|||
|
|
st = hist[pid].get("status", {})
|
|||
|
|
if st.get("status_str") == "error":
|
|||
|
|
raise RuntimeError(json.dumps(st)[:900])
|
|||
|
|
for node in hist[pid].get("outputs", {}).values():
|
|||
|
|
for im in node.get("images", []):
|
|||
|
|
q = urllib.parse.urlencode({"filename": im["filename"], "subfolder": im.get("subfolder", ""), "type": im.get("type", "output")})
|
|||
|
|
return Image.open(io.BytesIO(get("/view?" + q))).convert("RGBA")
|
|||
|
|
time.sleep(1.0)
|
|||
|
|
raise TimeoutError("comfy timeout")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def prepare(path):
|
|||
|
|
"""원본 아이콘을 캔버스 가운데에 여백을 두고 놓는다 — (어두운 바탕에 합성한 RGB, 같은 자리의 RGBA)"""
|
|||
|
|
im = Image.open(path).convert("RGBA")
|
|||
|
|
box = im.getchannel("A").getbbox()
|
|||
|
|
if box:
|
|||
|
|
im = im.crop(box)
|
|||
|
|
scale = RES * FIT / float(max(im.size))
|
|||
|
|
im = im.resize((max(1, int(im.width * scale)), max(1, int(im.height * scale))), Image.LANCZOS)
|
|||
|
|
rgba = Image.new("RGBA", (RES, RES), (0, 0, 0, 0))
|
|||
|
|
rgba.paste(im, ((RES - im.width) // 2, (RES - im.height) // 2), im)
|
|||
|
|
rgb = Image.new("RGB", (RES, RES), BG)
|
|||
|
|
rgb.paste(rgba, (0, 0), rgba)
|
|||
|
|
return rgb, rgba
|
|||
|
|
|
|||
|
|
|
|||
|
|
def mask(im, n=96, th=40):
|
|||
|
|
a = im.getchannel("A").resize((n, n), Image.BILINEAR)
|
|||
|
|
px = a.load()
|
|||
|
|
return [[px[x, y] > th for x in range(n)] for y in range(n)]
|
|||
|
|
|
|||
|
|
|
|||
|
|
def score(result, source_mask):
|
|||
|
|
"""(점수, 겹침 IoU, 가장 큰 덩어리 비율, 가장자리 닿음) — 점수가 높을수록 원본 형태를 지킨 깔끔한 한 덩어리"""
|
|||
|
|
n = len(source_mask)
|
|||
|
|
m = mask(result, n)
|
|||
|
|
inter = union = total = 0
|
|||
|
|
for y in range(n):
|
|||
|
|
for x in range(n):
|
|||
|
|
a, b = m[y][x], source_mask[y][x]
|
|||
|
|
if a: total += 1
|
|||
|
|
if a and b: inter += 1
|
|||
|
|
if a or b: union += 1
|
|||
|
|
if total == 0 or union == 0:
|
|||
|
|
return 0.0, 0.0, 0.0, 1.0
|
|||
|
|
iou = inter / float(union)
|
|||
|
|
edge = (sum(1 for i in range(n) if m[0][i]) + sum(1 for i in range(n) if m[n - 1][i]) + sum(1 for i in range(n) if m[i][0]) + sum(1 for i in range(n) if m[i][n - 1])) / float(4 * n)
|
|||
|
|
seen = [[False] * n for _ in range(n)]
|
|||
|
|
best = 0
|
|||
|
|
for y in range(n):
|
|||
|
|
for x in range(n):
|
|||
|
|
if m[y][x] and not seen[y][x]:
|
|||
|
|
q = deque([(x, y)]); seen[y][x] = True; size = 0
|
|||
|
|
while q:
|
|||
|
|
cx, cy = q.popleft(); size += 1
|
|||
|
|
for dx, dy in ((1, 0), (-1, 0), (0, 1), (0, -1)):
|
|||
|
|
nx, ny = cx + dx, cy + dy
|
|||
|
|
if 0 <= nx < n and 0 <= ny < n and m[ny][nx] and not seen[ny][nx]:
|
|||
|
|
seen[ny][nx] = True; q.append((nx, ny))
|
|||
|
|
best = max(best, size)
|
|||
|
|
main = best / float(total)
|
|||
|
|
s = iou * 0.7 + main * 0.3 - edge * 2.0
|
|||
|
|
return s, iou, main, edge
|
|||
|
|
|
|||
|
|
|
|||
|
|
def checker(size, n=16):
|
|||
|
|
im = Image.new("RGB", (size, size), (58, 58, 58))
|
|||
|
|
d = ImageDraw.Draw(im)
|
|||
|
|
for y in range(0, size, n):
|
|||
|
|
for x in range(0, size, n):
|
|||
|
|
if (x // n + y // n) % 2:
|
|||
|
|
d.rectangle([x, y, x + n - 1, y + n - 1], fill=(74, 74, 74))
|
|||
|
|
return im
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
out, tag = sys.argv[1], sys.argv[2]
|
|||
|
|
lo = int(sys.argv[3]) if len(sys.argv) > 3 else 0
|
|||
|
|
hi = int(sys.argv[4]) if len(sys.argv) > 4 else 999
|
|||
|
|
denoises = [float(v) for v in (sys.argv[5] if len(sys.argv) > 5 else "0.5,0.65").split(",")]
|
|||
|
|
seeds = int(sys.argv[6]) if len(sys.argv) > 6 else 3
|
|||
|
|
os.makedirs(out, exist_ok=True)
|
|||
|
|
cap = load_captions()
|
|||
|
|
ids = sorted(cap.keys())[lo:hi]
|
|||
|
|
show = len(denoises) * seeds
|
|||
|
|
cell, pad = 150, 4
|
|||
|
|
sheet = Image.new("RGB", (pad + (show + 1) * (cell + pad), len(ids) * (cell + pad) + pad), (40, 40, 40))
|
|||
|
|
draw = ImageDraw.Draw(sheet)
|
|||
|
|
scores = {}
|
|||
|
|
for i, gid in enumerate(ids):
|
|||
|
|
rgb, rgba = prepare(os.path.join(SRC, "%d.png" % gid))
|
|||
|
|
name = upload(rgb, "bt_moba_src_%d.png" % gid)
|
|||
|
|
src_mask = mask(rgba)
|
|||
|
|
prompt = "%s, (%s:1.2), no humans, single item, isolated, dark background" % (PREFIX, cap[gid])
|
|||
|
|
results = []
|
|||
|
|
t0 = time.time()
|
|||
|
|
for d in denoises:
|
|||
|
|
for j in range(seeds):
|
|||
|
|
im = run(name, prompt, 20000 + gid * 7 + j * 911 + int(d * 100), d)
|
|||
|
|
s, iou, mainc, edge = score(im, src_mask)
|
|||
|
|
fn = "%s_%d_d%02d_%d.png" % (tag, gid, int(d * 100), j)
|
|||
|
|
im.save(os.path.join(out, fn))
|
|||
|
|
results.append((s, fn, im, iou, mainc, edge, d))
|
|||
|
|
results.sort(key=lambda t: -t[0])
|
|||
|
|
scores[gid] = [dict(file=r[1], score=round(r[0], 3), iou=round(r[3], 3), main=round(r[4], 3), edge=round(r[5], 3), denoise=r[6]) for r in results]
|
|||
|
|
print("%s %d 최고 IoU %.2f(세기 %.2f) · %.0f초" % (tag, gid, results[0][3], results[0][6], time.time() - t0)); sys.stdout.flush()
|
|||
|
|
y = pad + i * (cell + pad)
|
|||
|
|
bg = checker(cell); small = rgba.resize((cell, cell), Image.LANCZOS); bg.paste(small, (0, 0), small); sheet.paste(bg, (pad, y))
|
|||
|
|
draw.text((pad + 3, y + 2), "%d 원본" % gid, fill=(255, 255, 255))
|
|||
|
|
for k, r in enumerate(results[:show]):
|
|||
|
|
bg = checker(cell); small = r[2].resize((cell, cell), Image.LANCZOS); bg.paste(small, (0, 0), small)
|
|||
|
|
sheet.paste(bg, (pad + (k + 1) * (cell + pad), y))
|
|||
|
|
draw.text((pad + (k + 1) * (cell + pad) + 3, y + 2), "d%.2f #%s IoU%.2f" % (r[6], r[1][-5], r[3]), fill=(255, 255, 255))
|
|||
|
|
json.dump(scores, io.open(os.path.join(out, "%s_scores_%d_%d.json" % (tag, lo, hi)), "w", encoding="utf-8"), ensure_ascii=False, indent=1)
|
|||
|
|
path = os.path.join(out, "%s_sheet_%d_%d.png" % (tag, lo, hi))
|
|||
|
|
sheet.save(path)
|
|||
|
|
print("sheet", path)
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|