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()
|