diff --git a/src/app.py b/src/app.py index 6b25eba..3b87df0 100644 --- a/src/app.py +++ b/src/app.py @@ -17,6 +17,7 @@ from discord_control.actions import setDiscordMute from ocr.ocr_actions import runOcrFromScreen from translate.translate_actions import saveTranslationText from translate.translate_actions import translateTextToJapanese +from vision.vision_actions import runVisionFromScreen from vrc_log.log_actions import collectVrchatLog from vrc_log.log_actions import startSelfMonitor from vrc_log.log_actions import startVrcLogMonitor @@ -27,6 +28,7 @@ PARAM_DISCORD_MUTE = "DiscordSend" PARAM_LOG_GREP = "vrc_log" PARAM_OCR = "ocrEnabled" PARAM_TRANSLATE_OCR = "translation" +PARAM_VISION = "visionEnabled" def create_gateway(): @@ -96,10 +98,20 @@ def create_gateway(): saved = saveTranslationText(stem, translated) gateway.log("INFO", f"translation saved={saved}") + def on_vision(address, *args): + if not args: + gateway.log("ERROR", "visionEnabled args empty") + return + + if gateway.is_rising_edge(address, args[0]): + gateway.log("INFO", f"received {address} args={args}") + runVisionFromScreen() + gateway.map(PARAM_DISCORD_MUTE, on_discord_mute) gateway.map(PARAM_LOG_GREP, on_log_grep) gateway.map(PARAM_OCR, on_ocr) gateway.map(PARAM_TRANSLATE_OCR, on_translate_ocr) + gateway.map(PARAM_VISION, on_vision) gateway.set_default_debug(enabled=False) return gateway diff --git a/src/common/config_loader.py b/src/common/config_loader.py index 364ea60..055b4f4 100644 --- a/src/common/config_loader.py +++ b/src/common/config_loader.py @@ -69,6 +69,23 @@ def loadOcrConfig(): } +def loadVisionConfig(): + config = loadTomlFile(CONFIG_FILE) + vision = config.get("vision", {}) + + return { + "model_name": str(vision.get("model_name", "google/gemma-3-4b-it")).strip(), + "prompt": str( + vision.get( + "prompt", + "Describe the image briefly and extract any readable text.", + ) + ).strip(), + "max_new_tokens": int(vision.get("max_new_tokens", 256)), + "temperature": float(vision.get("temperature", 0.2)), + } + + def loadVrcLogConfig(): config = loadTomlFile(CONFIG_FILE) vrc_log = config.get("vrc_log", {}) @@ -108,6 +125,14 @@ def loadVrcLogConfig(): notice_section = event_config.get("notice", {}) missing_count = notice_section.get("missing_count", 0) + log_patterns = ( + vrc_log.get("patterns") + or vrc_log.get("log_patterns") + or event_config.get("patterns") + or event_config.get("log_patterns") + or [] + ) + if not missing_count: missing_count = event_config.get("missing_count", 0) @@ -127,4 +152,5 @@ def loadVrcLogConfig(): "staff_names": [str(name).strip() for name in staff_names if str(name).strip()], "missing_count": int(missing_count), "self_name": str(self_name).strip(), + "log_patterns": [str(pattern).strip() for pattern in log_patterns if str(pattern).strip()], } diff --git a/src/vision/vision_actions.py b/src/vision/vision_actions.py new file mode 100644 index 0000000..a639f93 --- /dev/null +++ b/src/vision/vision_actions.py @@ -0,0 +1,102 @@ +from datetime import datetime + +from common.config_loader import loadVisionConfig +from common.runtime_log import appendRuntimeLog +from screen.screen_actions import captureOcrRegion + +_vision_pipeline = None + + +def log(level, message): + now = datetime.now().strftime("%H:%M:%S") + print(f"[{level}] {now} {message}", flush=True) + + +def getVisionPipeline(): + global _vision_pipeline + + if _vision_pipeline is not None: + return _vision_pipeline + + config = loadVisionConfig() + model_name = config["model_name"] + + if not model_name: + raise RuntimeError("vision.model_name is not set") + + try: + from transformers import AutoProcessor + from transformers import AutoModelForImageTextToText + import torch + except Exception as e: + raise RuntimeError(f"vision dependencies are missing detail={e}") + + log("INFO", f"Vision model initializing model={model_name}") + processor = AutoProcessor.from_pretrained(model_name) + model = AutoModelForImageTextToText.from_pretrained( + model_name, + torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, + device_map="auto", + ) + _vision_pipeline = (processor, model) + return _vision_pipeline + + +def saveVisionText(image_path, text): + return appendRuntimeLog( + "VISION TEXT", + f"image={image_path}\n{text}", + ) + + +def runVisionFromImage(image_path): + config = loadVisionConfig() + processor, model = getVisionPipeline() + + prompt = config["prompt"] + messages = [ + { + "role": "user", + "content": [ + {"type": "image"}, + {"type": "text", "text": prompt}, + ], + } + ] + + try: + import torch + from PIL import Image + except Exception as e: + log("ERROR", f"Vision runtime imports failed detail={e}") + return None + + try: + image = Image.open(image_path).convert("RGB") + inputs = processor(images=image, text=prompt, return_tensors="pt") + inputs = {key: value.to(model.device) for key, value in inputs.items()} + with torch.no_grad(): + output = model.generate( + **inputs, + max_new_tokens=config["max_new_tokens"], + temperature=config["temperature"], + ) + text = processor.batch_decode(output, skip_special_tokens=True)[0].strip() + except Exception as e: + log("ERROR", f"Vision inference failed detail={e}") + return None + + text_path = saveVisionText(image_path, text) + log("INFO", f"Vision text appended={text_path}") + print(text, flush=True) + return { + "image_path": image_path, + "text_path": text_path, + "text": text, + } + + +def runVisionFromScreen(): + log("ACTION", "runVisionFromScreen called") + image_path = captureOcrRegion() + return runVisionFromImage(image_path)