state.py 2.94 KB
import base64
import config
import numpy as np
import copy

from similarity import preprocess_image


class State(object):
    def __init__(self, board, index, state_check, class_names, browser, sim_func):
        self.class_names = class_names  # 클래스 이름들
        self.board = board  # 현재 상황 (set list)
        self.index = index  # 현재 위치
        self.id = self._convertStateToId()  # state를 식별할 id
        self.state_check = state_check
        self.allowedActions = self._available_action()
        self.browser = browser
        self._make_css(state_check)
        self.sim_func = sim_func
        self.now_image = preprocess_image(self.browser.capture())
        self.similarity = sim_func.similarity(self.now_image)
        # np.ones((len(self.board), config.NUM_ACTIONS), dtype=np.int)  # 실제 확인

    def _available_action(self):
        return np.nonzero(self.state_check[self.index])[0]

    def _convertStateToId(self):
        id = '{}'.format(self.index).join(map(lambda f: str(f) if len(f) != 0 else '{}', self.board))
        b64 = base64.b64encode(id.encode('utf-8'))
        return str(b64)

    def _checkForEndGame(self):
        if self.index >= len(self.board) - 2:
            return True
        return False

    def takeAction(self, act):
        """
        액션대로 상태를 변경
        :param act:
        :returns newState:
                 value:
                 done:
        """
        if act >= config.NUM_ACTIONS:
            raise IndexError

        newBoard = copy.deepcopy(self.board)
        newStateCheck = copy.deepcopy(self.state_check)
        value = 0
        done = 0
        index = self.index
        if act == config.LAST_ACTION:
            index += 1
        else:
            newBoard[index].add(act)
            newStateCheck[index][act] = 0
        newState = State(newBoard, index, newStateCheck, self.class_names, self.browser, self.sim_func)
        if newState._checkForEndGame():
            print('done')
            done = 1
        if newState.similarity >= 0.8:
            done = 1
            value = 1
            print('good')

        return newState, value, done

    def _make_css(self, state_check):
        """
        state_check로 css 작성
        :return: file_name
        """
        file_name = 'main.css'
        with open(file_name, 'w') as f:
            for class_index, actions in enumerate(state_check):
                class_index = int(class_index)
                f.write(".{}".format(self.class_names[class_index]))
                f.write("{")
                for action_index, action in enumerate(actions):
                    if action == 0:
                        f.write(" {}; ".format(config.ACTION_MEANING[action_index]))
                f.write("}\n")

        return file_name

    def reset(self):
        """
        상태 초기화
        :return:
        """
        self.index = 0
        self.state_check = np.ones((len(self.board), config.NUM_ACTIONS))