299 lines
10 KiB
Python
299 lines
10 KiB
Python
"""Lightweight image analysis for live capture feedback.
|
|
|
|
All stats are computed on a small downsampled copy so the demo loop stays
|
|
fast at full 1080p capture.
|
|
"""
|
|
|
|
from collections import Counter
|
|
|
|
from PIL import Image
|
|
|
|
ANALYSIS_SIZE = (320, 180)
|
|
_TOP_COLORS = 5
|
|
|
|
# Harvesting is detected from the single pixel at this window coordinate.
|
|
# Harvesting shows a light tan (~#f2c285, luma ~201); idle shows dark brown
|
|
# (~#664c33, luma ~81).
|
|
HARVEST_PIXEL = (1710, 1068)
|
|
HARVEST_LUMA_THRESHOLD = 140
|
|
|
|
# The small numeric readout just before each bar at the bottom of the panel.
|
|
# Regions are (x0, y0, x1, y1) in captured-window pixel coordinates. y0 starts
|
|
# below the panel's top border line (y=1032) which would confuse OCR.
|
|
NUMBER_REGIONS = {
|
|
"mana": (0, 1033, 62, 1048),
|
|
"food": (195, 1033, 226, 1048),
|
|
"health": (355, 1033, 389, 1048),
|
|
"inventory": (520, 1033, 552, 1048),
|
|
"action_points": (610, 1033, 714, 1048),
|
|
}
|
|
|
|
_OCR_SCALE = 3
|
|
|
|
# --- Crafting windows -------------------------------------------------------
|
|
# The Manufacturing and Inventory windows are semi-transparent panels the player
|
|
# can drag around, so nothing about them is fixed in screen space. Their title
|
|
# bars are located by OCR and everything else (status message, readouts) is
|
|
# offset from the title top-left. Calibrated from the frame sets in snapshots/.
|
|
TITLE_VERIFY_CROP = (14, 12, 150, 42) # (left, top, width, height) rel. to a candidate point
|
|
TITLE_VERIFY_THRESHOLD = 115
|
|
TITLE_SUBSTRING = {"manufacturing": "anufact", "inventory": "nven"}
|
|
|
|
# The craft status line sits at a fixed offset from the Manufacturing title,
|
|
# measured at 5 different window positions: always (title_x - 161, title_y + 121).
|
|
# A wider search region absorbs the few pixels of jitter from title detection.
|
|
MESSAGE_SEARCH = (175, 105, 600, 70) # (left, top, width, height) rel. to title
|
|
MESSAGE_COLOR_GREEN = (170, 120, 30) # (min green, min g-r, min g-b) for green text
|
|
MESSAGE_COLOR_RED = (170, 100, 60) # (min red, min r-g, min r-b) for red text
|
|
GREEN_OCR_THRESHOLD = 150
|
|
RED_OCR_THRESHOLD = 140
|
|
MESSAGE_OCR_PSM = 6
|
|
|
|
# Full-frame OCR is used to re-locate a window after it moves. Running it at two
|
|
# scales recovers some frames that fail at one scale on busy scenes. The word
|
|
# box OCR returns is usually the title itself (0..4px off) so the candidate
|
|
# title offsets only nudge that box.
|
|
FIND_OCR_SCALES = (1.0, 0.75)
|
|
# The OCR word box usually sits on the title text; -7 covers the Inventory
|
|
# window where OCR reads the glyph a few px below the verifying title bar.
|
|
FIND_TITLE_OFFSETS = ((0, 0), (0, 4), (0, -7))
|
|
|
|
try:
|
|
import pytesseract
|
|
except ImportError:
|
|
pytesseract = None
|
|
|
|
|
|
def _luma(pixel):
|
|
r, g, b = pixel
|
|
return round(0.299 * r + 0.587 * g + 0.114 * b)
|
|
|
|
|
|
def load_image(path):
|
|
"""Load a PNG as RGB; return None if the file is unreadable/corrupt."""
|
|
try:
|
|
return Image.open(path).convert("RGB")
|
|
except OSError:
|
|
return None
|
|
|
|
|
|
def is_harvesting(image):
|
|
"""Return whether the player is harvesting, based on the harvest pixel."""
|
|
rgb = image.convert("RGB")
|
|
color = rgb.getpixel(HARVEST_PIXEL)
|
|
lum = _luma(color)
|
|
return {
|
|
"is_harvesting": lum >= HARVEST_LUMA_THRESHOLD,
|
|
"color": "#%02x%02x%02x" % color,
|
|
"luma": lum,
|
|
}
|
|
|
|
|
|
def _ocr_number(image, region):
|
|
"""OCR a single number region; return int or None if it can't be read."""
|
|
if pytesseract is None:
|
|
return None
|
|
crop = image.convert("L").crop(region)
|
|
crop = crop.resize((crop.width * _OCR_SCALE, crop.height * _OCR_SCALE))
|
|
crop = crop.point(lambda p: 0 if p < 140 else 255)
|
|
try:
|
|
text = pytesseract.image_to_string(
|
|
crop,
|
|
config="--psm 7 -c tessedit_char_whitelist=0123456789",
|
|
)
|
|
except Exception:
|
|
return None
|
|
digits = "".join(ch for ch in text if ch.isdigit())
|
|
return int(digits) if digits else None
|
|
|
|
|
|
def read_bar_values(image):
|
|
"""Return the numeric readout before each bar, or None if unreadable."""
|
|
rgb = image.convert("RGB")
|
|
return {
|
|
name: _ocr_number(rgb, region)
|
|
for name, region in NUMBER_REGIONS.items()
|
|
}
|
|
|
|
|
|
def verify_window(image, x, y, window):
|
|
"""Return True if a window of the given kind has its title bar at (x, y)."""
|
|
if pytesseract is None:
|
|
return False
|
|
left, top, width, height = TITLE_VERIFY_CROP
|
|
crop = image.convert("L").crop((x - left, y - top, x - left + width, y - top + height))
|
|
crop = crop.resize((crop.width * _OCR_SCALE, crop.height * _OCR_SCALE))
|
|
crop = crop.point(lambda p: 0 if p < TITLE_VERIFY_THRESHOLD else 255)
|
|
try:
|
|
text = pytesseract.image_to_string(crop, config="--psm 7")
|
|
except Exception:
|
|
return False
|
|
return TITLE_SUBSTRING[window] in text.lower()
|
|
|
|
|
|
def _color_mask(crop, kind):
|
|
"""Return a binary mask isolating green or red text pixels in a crop."""
|
|
px = crop.load()
|
|
width, height = crop.size
|
|
if kind == "green":
|
|
lo, dr, db = MESSAGE_COLOR_GREEN
|
|
def hit(r, g, b):
|
|
return g > lo and g - r > dr and g - b > db
|
|
else:
|
|
lo, dg, db = MESSAGE_COLOR_RED
|
|
def hit(r, g, b):
|
|
return r > lo and r - g > dg and r - b > db
|
|
mask = Image.new("1", (width, height))
|
|
mask.putdata([1 if hit(*px[x, y]) else 0 for y in range(height) for x in range(width)])
|
|
return mask
|
|
|
|
|
|
def _ocr_message_channel(crop, box, kind):
|
|
"""OCR the green or red channel of a crop; return lowercased text or None."""
|
|
if pytesseract is None:
|
|
return None
|
|
channel = crop.split()[1] if kind == "green" else crop.split()[0]
|
|
channel = channel.crop(box)
|
|
threshold = GREEN_OCR_THRESHOLD if kind == "green" else RED_OCR_THRESHOLD
|
|
channel = channel.point(lambda p: 0 if p < threshold else 255)
|
|
channel = channel.resize((channel.width * _OCR_SCALE, channel.height * _OCR_SCALE))
|
|
try:
|
|
text = pytesseract.image_to_string(channel, config=f"--psm {MESSAGE_OCR_PSM}")
|
|
except Exception:
|
|
return None
|
|
text = " ".join(text.split()).lower()
|
|
return text or None
|
|
|
|
|
|
def classify_message(text, color):
|
|
"""Map an OCR'd craft status line to a stable category string."""
|
|
if not text:
|
|
return "none"
|
|
if color == "green":
|
|
if "started working" in text:
|
|
return "working"
|
|
if "successfully created" in text:
|
|
return "success"
|
|
if "stopped working" in text:
|
|
return "stopped"
|
|
if "lost the ingredients" in text:
|
|
return "failed_lost"
|
|
if "failed to create" in text or "failed to craft" in text:
|
|
return "failed"
|
|
if "hungry" in text:
|
|
return "hungry"
|
|
if "overloaded" in text:
|
|
return "overloaded"
|
|
return "unknown"
|
|
|
|
|
|
def read_manufacturing_message(image, title_x, title_y):
|
|
"""OCR the colored status line below a Manufacturing window title.
|
|
|
|
Returns {"text", "color", "category"} or None if OCR is unavailable.
|
|
"""
|
|
if pytesseract is None:
|
|
return None
|
|
left, top, width, height = MESSAGE_SEARCH
|
|
region = image.convert("RGB").crop(
|
|
(title_x - left, title_y + top, title_x - left + width, title_y + top + height)
|
|
)
|
|
for kind in ("green", "red"):
|
|
box = _color_mask(region, kind).getbbox()
|
|
if box is None:
|
|
continue
|
|
pad = 2
|
|
box = (max(0, box[0] - pad), max(0, box[1] - pad),
|
|
min(region.width, box[2] + pad), min(region.height, box[3] + pad))
|
|
text = _ocr_message_channel(region, box, kind)
|
|
if text:
|
|
return {"text": text, "color": kind, "category": classify_message(text, kind)}
|
|
return {"text": None, "color": None, "category": "none"}
|
|
|
|
|
|
def _find_title_words(image):
|
|
"""Return candidate title word boxes from full-frame OCR.
|
|
|
|
Each entry is a {"window": kind, "x": float, "y": float} in full-frame
|
|
pixel coordinates. Words containing the wrong substring (e.g. "manufacture"
|
|
inside the hungry status message) are filtered out by the caller's verify.
|
|
"""
|
|
if pytesseract is None:
|
|
return []
|
|
words = []
|
|
for scale in FIND_OCR_SCALES:
|
|
im = image.convert("L")
|
|
if scale != 1.0:
|
|
im = im.resize((int(im.width * scale), int(im.height * scale)), Image.BICUBIC)
|
|
try:
|
|
data = pytesseract.image_to_data(
|
|
im, config="--oem 1 --psm 11", output_type=pytesseract.Output.DICT
|
|
)
|
|
except Exception:
|
|
continue
|
|
for i, word in enumerate(data["text"]):
|
|
flat = word.lower().strip().replace(" ", "")
|
|
for kind, needle in TITLE_SUBSTRING.items():
|
|
if needle in flat:
|
|
words.append(
|
|
{"window": kind, "x": data["left"][i] / scale, "y": data["top"][i] / scale}
|
|
)
|
|
return words
|
|
|
|
|
|
def find_windows(image):
|
|
"""Locate Manufacturing and Inventory windows from a full frame scan.
|
|
|
|
Returns {"manufacturing": (x, y) or None, "inventory": (x, y) or None}.
|
|
"""
|
|
found = {"manufacturing": None, "inventory": None}
|
|
for word in _find_title_words(image):
|
|
kind = word["window"]
|
|
if found[kind] is not None:
|
|
continue
|
|
for dx, dy in FIND_TITLE_OFFSETS:
|
|
x, y = round(word["x"] + dx), round(word["y"] + dy)
|
|
if verify_window(image, x, y, kind):
|
|
found[kind] = (x, y)
|
|
break
|
|
return found
|
|
|
|
|
|
def _pixels(image):
|
|
px = image.load()
|
|
w, h = image.size
|
|
return [px[x, y] for y in range(h) for x in range(w)]
|
|
|
|
|
|
def summarize(image):
|
|
"""Return size, mean RGB, and dominant colors of a captured frame."""
|
|
small = image.convert("RGB").resize(ANALYSIS_SIZE)
|
|
pixels = _pixels(small)
|
|
n = len(pixels)
|
|
total = [0, 0, 0]
|
|
for r, g, b in pixels:
|
|
total[0] += r
|
|
total[1] += g
|
|
total[2] += b
|
|
mean = tuple(round(c / n) for c in total)
|
|
top = Counter(pixels).most_common(_TOP_COLORS)
|
|
dominant = [
|
|
{"hex": "#%02x%02x%02x" % color, "fraction": round(count / n, 3)}
|
|
for color, count in top
|
|
]
|
|
return {"size": (image.width, image.height), "mean_rgb": mean, "dominant": dominant}
|
|
|
|
|
|
def difference(a, b):
|
|
"""Sum of absolute channel differences between two same-size RGB images."""
|
|
if a.size != b.size:
|
|
return None
|
|
pa, pb = a.load(), b.load()
|
|
w, h = a.size
|
|
diff = 0
|
|
for y in range(h):
|
|
for x in range(w):
|
|
ca, cb = pa[x, y], pb[x, y]
|
|
diff += abs(ca[0] - cb[0]) + abs(ca[1] - cb[1]) + abs(ca[2] - cb[2])
|
|
return diff
|