Attributes: project_root: root directory path of the COCO project. model_config_path: path to the pytorch_config.yaml file train_json_filename: the name of the json file containing the train annotations test_json_filename: the name of the json file containing the
| 26 | |
| 27 | |
| 28 | class COCOLoader(Loader): |
| 29 | """ |
| 30 | Attributes: |
| 31 | project_root: root directory path of the COCO project. |
| 32 | model_config_path: path to the pytorch_config.yaml file |
| 33 | train_json_filename: the name of the json file containing the train annotations |
| 34 | test_json_filename: the name of the json file containing the train annotations. |
| 35 | None if there is no test set. |
| 36 | |
| 37 | Examples: |
| 38 | loader = COCOLoader( |
| 39 | project_root='/path/to/project/', |
| 40 | model_config_path='/path/to/project/experiments/train/pytorch_config.yaml' |
| 41 | train_json_filename="train.json", |
| 42 | test_json_filename="test.json", |
| 43 | ) |
| 44 | """ |
| 45 | |
| 46 | def __init__( |
| 47 | self, |
| 48 | project_root: str | Path, |
| 49 | model_config_path: str | Path, |
| 50 | train_json_filename: str = "train.json", |
| 51 | test_json_filename: str = "test.json", |
| 52 | ): |
| 53 | image_root = Path(project_root) / "images" |
| 54 | super().__init__(project_root, image_root, Path(model_config_path)) |
| 55 | self.train_json_filename = train_json_filename |
| 56 | self.test_json_filename = test_json_filename |
| 57 | self._dataset_parameters = None |
| 58 | |
| 59 | self.train_json = self.load_json(self.project_root, self.train_json_filename) |
| 60 | self.test_json = None |
| 61 | if self.test_json_filename: |
| 62 | self.test_json = self.load_json(self.project_root, self.test_json_filename) |
| 63 | |
| 64 | def get_dataset_parameters(self) -> PoseDatasetParameters: |
| 65 | """Retrieves dataset parameters based on the instance's configuration. |
| 66 | |
| 67 | Returns: |
| 68 | An instance of the PoseDatasetParameters with the parameters set. |
| 69 | """ |
| 70 | if self._dataset_parameters is None: |
| 71 | num_individuals, bodyparts = self.get_project_parameters(self.train_json) |
| 72 | |
| 73 | crop_cfg = self.model_cfg["data"]["train"].get("top_down_crop", {}) |
| 74 | crop_w, crop_h = crop_cfg.get("width", 256), crop_cfg.get("height", 256) |
| 75 | crop_margin = crop_cfg.get("margin", 0) |
| 76 | crop_with_context = crop_cfg.get("crop_with_context", True) |
| 77 | |
| 78 | self._dataset_parameters = PoseDatasetParameters( |
| 79 | bodyparts=bodyparts, |
| 80 | unique_bpts=[], |
| 81 | individuals=[f"individual{i}" for i in range(num_individuals)], |
| 82 | with_center_keypoints=self.model_cfg.get("with_center_keypoints", False), |
| 83 | color_mode=self.model_cfg.get("color_mode", "RGB"), |
| 84 | top_down_crop_size=(crop_w, crop_h), |
| 85 | top_down_crop_margin=crop_margin, |
no outgoing calls
no test coverage detected