boki/poe2-bot
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

work on well of souls and yolo detection

Browse source
This commit is contained in:
Boki 2026-02-20 16:40:50 -05:00
parent 3456e0d62a
commit 40d30115bf
41 changed files with 3031 additions and 148 deletions
Download patch file Download diff file Expand all files Collapse all files

33
tools/python-detect/daemon.py
View file

@ -12,7 +12,7 @@ import sys
import json
import time
_model = None
_models = {}
def _redirect_stdout_to_stderr():
@ -26,36 +26,40 @@ def _restore_stdout(real_stdout):
sys.stdout = real_stdout
def load_model():
global _model
if _model is not None:
return _model
def load_model(name="enemy-v1"):
if name in _models:
return _models[name]
import os
from ultralytics import YOLO
model_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models")
model_path = os.path.join(model_dir, "enemy-v1.pt")
# Prefer TensorRT engine if available, fall back to .pt
engine_path = os.path.join(model_dir, f"{name}.engine")
pt_path = os.path.join(model_dir, f"{name}.pt")
model_path = engine_path if os.path.exists(engine_path) else pt_path
if not os.path.exists(model_path):
raise FileNotFoundError(f"Model not found: {model_path}")
raise FileNotFoundError(f"Model not found: {pt_path} (also checked {engine_path})")
sys.stderr.write(f"Loading YOLO model from {model_path}...\n")
sys.stderr.write(f"Loading YOLO model '{name}' from {model_path}...\n")
sys.stderr.flush()
real_stdout = _redirect_stdout_to_stderr()
try:
_model = YOLO(model_path)
model = YOLO(model_path)
# Warmup with dummy inference (triggers CUDA init)
import numpy as np
dummy = np.zeros((640, 640, 3), dtype=np.uint8)
_model.predict(dummy, verbose=False)
model.predict(dummy, verbose=False)
finally:
_restore_stdout(real_stdout)
sys.stderr.write("YOLO model loaded and warmed up.\n")
_models[name] = model
sys.stderr.write(f"YOLO model '{name}' loaded and warmed up.\n")
sys.stderr.flush()
return _model
return model
def handle_detect(req):
@ -70,12 +74,15 @@ def handle_detect(req):
img_bytes = base64.b64decode(image_base64)
img = np.array(Image.open(io.BytesIO(img_bytes)))
# PIL gives RGB, but ultralytics model.predict() assumes numpy arrays are BGR
img = img[:, :, ::-1]
conf = req.get("conf", 0.3)
iou = req.get("iou", 0.45)
imgsz = req.get("imgsz", 640)
model_name = req.get("model", "enemy-v1")
model = load_model()
model = load_model(model_name)
real_stdout = _redirect_stdout_to_stderr()
try: