| 48 | |
| 49 | |
| 50 | class TrainNetwork(DefaultTab): |
| 51 | def __init__(self, root, parent, h1_description): |
| 52 | super().__init__(root, parent, h1_description) |
| 53 | self._shuffle: ShuffleSpinBox = ShuffleSpinBox(root=self.root, parent=self) |
| 54 | self._shuffle_display = SelectedShuffleDisplay(self.root) |
| 55 | |
| 56 | self._attribute_layouts: dict[Engine, QtWidgets.QWidget] = {} |
| 57 | self._attribute_kwargs: dict[Engine, dict] = {} |
| 58 | self._rows_with_requirements: list = [] |
| 59 | self._set_page() |
| 60 | |
| 61 | self.root.engine_change.connect(self._on_engine_change) |
| 62 | self._shuffle_display.pose_cfg_signal.connect(self._pose_cfg_change) |
| 63 | |
| 64 | @Slot(Engine) |
| 65 | def _on_engine_change(self, engine: Engine) -> None: |
| 66 | for e, layout in self._attribute_layouts.items(): |
| 67 | if e == engine: |
| 68 | layout.show() |
| 69 | else: |
| 70 | layout.hide() |
| 71 | self._update_snapshot_selection_widgets_visibility() |
| 72 | |
| 73 | def _update_snapshot_selection_widgets_visibility(self): |
| 74 | if self.root.engine == Engine.PYTORCH: |
| 75 | self.resume_from_snapshot_label.show() |
| 76 | self.snapshot_selection_widget.show() |
| 77 | # Display detector snapshot selection widget only if in Top-Down mode |
| 78 | if self._shuffle_display.pose_cfg.get("method", "").lower() == "td": |
| 79 | self.detector_snapshot_selection_widget.show() |
| 80 | else: |
| 81 | self.detector_snapshot_selection_widget.hide() |
| 82 | else: |
| 83 | self.resume_from_snapshot_label.hide() |
| 84 | self.snapshot_selection_widget.hide() |
| 85 | self.detector_snapshot_selection_widget.hide() |
| 86 | |
| 87 | def _set_page(self): |
| 88 | self.main_layout.addWidget(_create_label_widget("Attributes", "font:bold")) |
| 89 | self._generate_layout_attributes() |
| 90 | |
| 91 | self.resume_from_snapshot_label = _create_label_widget( |
| 92 | "[Optional]: Select a snapshot to resume training from", "font:bold" |
| 93 | ) |
| 94 | self.resume_from_snapshot_label.setToolTip( |
| 95 | "<span style='font-weight:normal; white-space:nowrap;'>" |
| 96 | "If you've already trained a model on this shuffle, you can continue training it instead of starting " |
| 97 | "from scratch again. <br>When using top-down models, you can also choose a detector to resume training" |
| 98 | "from." |
| 99 | "</span>" |
| 100 | ) |
| 101 | self.main_layout.addWidget(self.resume_from_snapshot_label) |
| 102 | |
| 103 | self.snapshot_selection_widget = SnapshotSelectionWidget( |
| 104 | self.root, self, margins=(30, 0, 0, 0), select_button_text="Select snapshot" |
| 105 | ) |
| 106 | self.main_layout.addWidget(self.snapshot_selection_widget) |
| 107 | |