mower-ng/mower/solvers/infra/get_agent_from_room.py

127 lines
4.3 KiB
Python

from datetime import datetime
import cv2
from mower.solvers.infra.base_mixin import BaseMixin
from mower.solvers.infra.enter_room import EnterRoomSolver
from mower.utils import config
from mower.utils import typealias as tp
from mower.utils.graph import SceneGraphSolver
from mower.utils.image import cropimg, thres2
from mower.utils.log import logger
name_x = (1288, 1869)
name_y = [(135, 326), (344, 535), (553, 744), (532, 723), (741, 932)]
name_p = [tuple(zip(name_x, y)) for y in name_y]
mood_x = (1470, 1780)
mood_y = [(219, 220), (428, 429), (637, 638), (615, 616), (823, 824)]
mood_p = [tuple(zip(mood_x, y)) for y in mood_y]
time_x = (1650, 1780)
time_y = [(270, 305), (480, 515), (690, 725), (668, 703), (877, 912)]
time_p = [tuple(zip(time_x, y)) for y in time_y]
class GetAgentFromRoomSolver(SceneGraphSolver, BaseMixin):
solver_name = "读取干员信息"
def run(self, room: str, read_agent_time=False):
"""
Args:
room: 房间名
read_agent_time: 读取工作或休息时间
"""
self.room = room
self.read_agent_time = read_agent_time
self.index = 0
self.result = []
EnterRoomSolver().run(self.room)
self.length = self.number(((1867, 1020), (1885, 1045)), 19, 200)
self.need_swipe = True if self.length > 3 else False
super().run()
return self.result
def swipe_up(self):
self.swipe(
(config.recog.w * 0.8, config.recog.h * 0.5),
(0, config.recog.h * 0.45),
duration=500,
interval=1,
)
def swipe_down(self):
self.swipe(
(config.recog.w * 0.8, config.recog.h * 0.5),
(0, -config.recog.h * 0.45),
duration=500,
interval=1,
)
def number(self, scope: tp.Scope, height: int, thres: int) -> int:
"数字识别"
img = cropimg(config.recog.gray, scope)
default_height = 25
if height != default_height:
scale = 25 / height
img = cv2.resize(img, None, None, scale, scale)
img = thres2(img, thres)
img = cv2.bitwise_not(img)
return (
config.recog.num.number_int("riic_base", img=img, target_range=range(1, 6))
+ 1
)
@staticmethod
def is_power_station() -> bool:
img = cropimg(config.recog.img, ((568, 18), (957, 95)))
hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
mask = cv2.inRange(hsv, (35, 0, 0), (38, 255, 255))
if cv2.countNonZero(mask) > 1000:
return True
return False
def read_agent_data(self, index):
data = {}
if self.find("infra_no_operator", scope=name_p[index]):
data["agent"] = ""
data["mood"] = -1
else:
data["agent"] = self.read_screen(
cropimg(config.recog.gray, name_p[index]), type="name"
)
data["mood"] = self.read_accurate_mood(
cropimg(config.recog.gray, mood_p[index])
)
if self.read_agent_time:
if (
data["mood"] == 24
or self.room in ["central", "meeting", "factory"]
or self.is_power_station()
):
data["time"] = datetime.now()
else:
data["time"] = self.double_read_time(time_p[index])
self.result.append(data)
def transition(self) -> bool:
if self.detect_product_complete():
logger.info("检测到产物收取提示")
self.sleep(0.5)
elif self.find("room_detail"):
if self.index <= 1 and self.get_color((1800, 930))[0] < 51:
self.swipe_up()
return
if self.index >= 3 and self.need_swipe:
if self.get_color((1800, 930))[0] > 51:
self.swipe_down()
return
else:
self.need_swipe = False
elif self.index == self.length:
return True
elif self.index < self.length:
self.read_agent_data(self.index)
self.index += 1
elif self.scene() in self.waiting_scene:
self.waiting_solver()
else:
EnterRoomSolver().run(self.room)