Main entry point to access all observation processor
| 606 | |
| 607 | |
| 608 | class ObservationHandler: |
| 609 | """Main entry point to access all observation processor""" |
| 610 | |
| 611 | def __init__( |
| 612 | self, |
| 613 | main_observation_type: str, |
| 614 | text_observation_type: str, |
| 615 | image_observation_type: str, |
| 616 | current_viewport_only: bool, |
| 617 | viewport_size: ViewportSize, |
| 618 | ) -> None: |
| 619 | self.main_observation_type = main_observation_type |
| 620 | self.text_processor = TextObervationProcessor( |
| 621 | text_observation_type, current_viewport_only, viewport_size |
| 622 | ) |
| 623 | self.image_processor = ImageObservationProcessor( |
| 624 | image_observation_type |
| 625 | ) |
| 626 | self.viewport_size = viewport_size |
| 627 | |
| 628 | @beartype |
| 629 | def get_observation_space(self) -> spaces.Dict: |
| 630 | text_space = spaces.Text( |
| 631 | min_length=0, |
| 632 | max_length=UTTERANCE_MAX_LENGTH, |
| 633 | charset=ASCII_CHARSET + FREQ_UNICODE_CHARSET, |
| 634 | ) |
| 635 | |
| 636 | image_space = spaces.Box( |
| 637 | # Each position stores the RGB values. Note the swapped axes (height first). |
| 638 | np.zeros( |
| 639 | (self.viewport_size["height"], self.viewport_size["width"], 3), |
| 640 | dtype=np.uint8, |
| 641 | ), |
| 642 | np.ones( |
| 643 | (self.viewport_size["height"], self.viewport_size["width"], 3), |
| 644 | dtype=np.uint8, |
| 645 | ) |
| 646 | * 255.0, |
| 647 | dtype=np.uint8, |
| 648 | ) |
| 649 | |
| 650 | return spaces.Dict({"text": text_space, "image": image_space}) |
| 651 | |
| 652 | @beartype |
| 653 | def get_observation( |
| 654 | self, page: Page, client: CDPSession |
| 655 | ) -> dict[str, Observation]: |
| 656 | text_obs = self.text_processor.process(page, client) |
| 657 | image_obs = self.image_processor.process(page, client) |
| 658 | return {"text": text_obs, "image": image_obs} |
| 659 | |
| 660 | @beartype |
| 661 | def get_observation_metadata(self) -> dict[str, ObservationMetadata]: |
| 662 | return { |
| 663 | "text": self.text_processor.meta_data, |
| 664 | "image": self.image_processor.meta_data, |
| 665 | } |