diff --git a/ev3sim/attach_bot.py b/ev3sim/attach_bot.py index 47c17ed7..52791fda 100644 --- a/ev3sim/attach_bot.py +++ b/ev3sim/attach_bot.py @@ -73,6 +73,10 @@ def handle_recv(msg_type, msg): tick = msg["tick"] tick_rate = msg["tick_rate"] current_data = msg["data"] + if isinstance(current_data, str): + # Not pretty but it works. + e = Exception(current_data) + raise e for ev in msg["events"]: cur_events.put(ev) return msg_type, msg @@ -298,6 +302,8 @@ def seek(self, i): def write(self, value): send_q.put((DEVICE_WRITE, (f"{self.k2} {self.k3} {self.k4}", value.decode()))) + while self.k4 == "mode" and current_data[self.k2][self.k3][self.k4] != value.decode(): + wait_for_tick() def flush(self): pass diff --git a/ev3sim/devices/colour/base.py b/ev3sim/devices/colour/base.py index 0d2e72c4..d03277ea 100644 --- a/ev3sim/devices/colour/base.py +++ b/ev3sim/devices/colour/base.py @@ -5,6 +5,8 @@ class ColourSensorMixin: RGB_RAW = "RGB-RAW" + COL_REFLECT = "COL-REFLECT" + COL_COLOR = "COL-COLOR" device_type = "lego-sensor" mode = RGB_RAW @@ -44,6 +46,10 @@ def toObject(self): } if self.mode == self.RGB_RAW: data["value0"], data["value1"], data["value2"] = res + elif self.mode == self.COL_REFLECT: + data["value0"] = self.reflected_light_intensity() + elif self.mode == self.COL_COLOR: + data["value0"] = self.predict_color() else: raise ValueError(f"Unhandled mode {self.mode}") return data diff --git a/ev3sim/devices/colour/ev3.py b/ev3sim/devices/colour/ev3.py index 84f1e48d..f2da3633 100644 --- a/ev3sim/devices/colour/ev3.py +++ b/ev3sim/devices/colour/ev3.py @@ -31,6 +31,18 @@ class ColorSensor(ColourSensorMixin, Device): MIN_RGB_BIAS = 230 STARTING_CALIBRATION = 300 + PREDICTION_VECTORS = [ + (0, 0, 0), + (0, 0, 255), + (85, 166, 48), + (255, 255, 0), + (255, 0, 0), + (255, 255, 255), + (165, 42, 42), + ] + PREDICTION_INCREASE_REQUIREMENT = 60 + PREDICTION_LARGEST_DIFFERENCE = 110 + def generateBias(self): self.saved_raw = (0, 0, 0) if ScriptLoader.RANDOMISE_SENSORS: @@ -102,8 +114,19 @@ def rgb(self): ] def reflected_light_intensity(self): - """Not implemented""" - raise NotImplementedError("`reflected_light_intensity` is currently not implemented.") + r, g, b = self.rgb() + return int((r / 255 + g / 255 + b / 255) * 100 / 3) + + def predict_color(self): + r, g, b = self.rgb() + colors = sorted( + [(abs(r - c[0]) + abs(g - c[1]) + abs(b - c[2]), i) for i, c in enumerate(self.PREDICTION_VECTORS)] + ) + if colors[0][0] > self.PREDICTION_LARGEST_DIFFERENCE: + return 0 + if colors[1][0] - colors[0][0] < self.PREDICTION_INCREASE_REQUIREMENT: + return 0 + return colors[0][1] + 1 def reset(self): self.mode = self.RGB_RAW diff --git a/ev3sim/robot.py b/ev3sim/robot.py index 6b2e6d6b..a7d6cb77 100644 --- a/ev3sim/robot.py +++ b/ev3sim/robot.py @@ -123,7 +123,11 @@ def collectDeviceData(self): for port, device in self.devices.items(): if device.device_type not in res: res[device.device_type] = {} - res[device.device_type][device._getObjName(port)] = device.toObject() + try: + res[device.device_type][device._getObjName(port)] = device.toObject() + except Exception as e: + # Something has failed, so let the attached script know. + return str(e) return res def resetBot(self):