Add search range for random frames

vspreview/plugins/builtins/slowpics_comp.ppy

@@ -30,3 +30,6 @@ class CompPlugin(AbstractPlugin, QTabWidget):
         self.addTab(self.settings_tab, 'Settings')
 
         self.settings.globals.settings = self.settings_tab
+
+    def on_current_output_changed(self, index: int, prev_index: int) -> None:
+        self.main_tab.on_current_output_changed(index, prev_index)

vspreview/plugins/builtins/slowpics_comp/main.py

@@ -18,7 +18,7 @@
 
 from vspreview.core import (
     CheckBox, ComboBox, ExtendedWidget, Frame, FrameEdit, HBoxLayout, LineEdit, ProgressBar, PushButton, VBoxLayout,
-    VideoOutput, main_window, try_load
+    VideoOutput, main_window, try_load, Stretch
 )
 from vspreview.models import GeneralModel
 
@@ -178,6 +178,9 @@ def setup_ui(self) -> None:
 
         self.random_frames_control = FrameEdit(self)
 
+        self.start_rando_frames_control = FrameEdit(self)
+        self.end_rando_frames_control = FrameEdit(self)
+
         self.manual_frames_lineedit = LineEdit('Manual frames: frame,frame,frame', self, )
 
         self.current_frame_checkbox = CheckBox('Current frame', self, checked=True)
@@ -297,8 +300,20 @@ def setup_ui(self) -> None:
         self.vlayout.addWidget(self.get_separator(True))
 
         HBoxLayout(self.vlayout, [
-            self.tmdb_id_lineedit,
-            self.tmdb_type_combox,
+            VBoxLayout([
+                QLabel('Start frame:'),
+                self.start_rando_frames_control
+            ]),
+            VBoxLayout([
+                QLabel('End frame:'),
+                self.end_rando_frames_control
+            ]),
+            self.get_separator(),
+            VBoxLayout([
+                self.tmdb_id_lineedit,
+                self.tmdb_type_combox,
+            ]),
+            Stretch()
         ])
 
         self.vlayout.addWidget(self.tag_separator)
@@ -406,6 +421,23 @@ def add_shortcuts(self) -> None:
     def update_tags(self) -> None:
         self.tag_list_combox.setModel(GeneralModel[str](sorted(self.tag_data.keys()), to_title=False))
 
+    def on_current_output_changed(self, index: int, prev_index: int) -> None:
+        assert self.main.outputs
+
+        if not self._thread_running:
+            old_start, old_end = self.start_rando_frames_control.value(), self.end_rando_frames_control.value()
+
+            self.start_rando_frames_control.setMaximum(self.main.current_output.total_frames - 1)
+            self.end_rando_frames_control.setMaximum(self.main.current_output.total_frames)
+
+            if (
+                (old_start, old_end) == (Frame(0), Frame(0))
+            ) or (
+                (old_start, old_end) == (Frame(0), self.main.outputs[prev_index].total_frames)
+            ):
+                self.start_rando_frames_control.setValue(Frame(0))
+                self.end_rando_frames_control.setValue(self.main.current_output.total_frames)
+
     def on_public_toggle(self, new_state: bool) -> None:
         if not new_state or self.tag_data:
             return
@@ -485,6 +517,9 @@ def on_end_upload(self, uuid: str, forced: bool = False) -> None:
 
         self._thread_running = False
 
+        self.start_rando_frames_control.setEnabled(True)
+        self.end_rando_frames_control.setEnabled(True)
+
         if forced:
             self.upload_progressbar.setValue(int())
             self.upload_status_label.setText('Stopped!')
@@ -669,8 +704,13 @@ def find_samples(self, uuid: str) -> bool:
                 if self.current_frame_checkbox.isChecked():
                     samples.append(self.main.current_output.last_showed_frame)
 
+                start_frame = int(self.start_rando_frames_control.value())
+                end_frame = int(self.end_rando_frames_control.value())
+
                 config = FindFramesWorkerConfiguration(
-                    uuid, self.main.current_output, self.outputs, self.main, lens_n, dark_num, light_num,
+                    uuid, self.main.current_output, self.outputs, self.main,
+                    start_frame, end_frame,
+                    min(lens_n, end_frame - start_frame), dark_num, light_num,
                     num, picture_types, samples
                 )
             except RuntimeError as e:
@@ -692,6 +732,9 @@ def find_samples(self, uuid: str) -> bool:
 
             self._thread_running = True
 
+            self.start_rando_frames_control.setEnabled(False)
+            self.end_rando_frames_control.setEnabled(False)
+
             return True
         except BaseException as e:
             self.main.show_message(str(e))

vspreview/plugins/builtins/slowpics_comp/workers.py

@@ -263,6 +263,8 @@ class FindFramesWorkerConfiguration(NamedTuple):
     current_output: VideoOutput
     outputs: list[VideoOutput]
     main: MainWindow
+    start_frame: int
+    end_frame: int
     num_frames: int
     dark_frames: int
     light_frames: int
@@ -309,7 +311,9 @@ def _select_samples_ptypes(self, conf: FindFramesWorkerConfiguration) -> list[Fr
                     logging.warning(f'There aren\'t enough of {conf.picture_types} in these clips')
                     raise StopIteration
 
-                rnum = rand_num_frames(_rnum_checked, partial(random.randrange, start=interval * num, stop=(interval * (num + 1)) - 1))
+                rnum = conf.start_frame + rand_num_frames(
+                    _rnum_checked, partial(random.randrange, start=interval * num, stop=(interval * (num + 1)) - 1)
+                )
                 _rnum_checked.add(rnum)
 
                 if all(
@@ -364,7 +368,9 @@ def _find_dark_light(self, conf: FindFramesWorkerConfiguration) -> list[Frame]:
                     logging.warning('There aren\'t enough of dark/light in these clips')
                     raise StopIteration
 
-                rnum = rand_num_frames(_rnum_checked, partial(random.randrange, start=interval * num, stop=(interval * (num + 1)) - 1))
+                rnum = conf.start_frame + rand_num_frames(
+                    _rnum_checked, partial(random.randrange, start=interval * num, stop=(interval * (num + 1)) - 1)
+                )
                 _rnum_checked.add(rnum)
 
                 avg = get_prop(stats.get_frame(rnum), "PlaneStatsAverage", float, None, 0)
@@ -402,10 +408,16 @@ def run(self, conf: FindFramesWorkerConfiguration) -> None:  # type: ignore
         try:
             if conf.ptype_num:
                 if conf.picture_types == {'I', 'P', 'B'}:
-                    interval = conf.num_frames // conf.ptype_num
-                    samples = list(map(
-                        Frame, list(random.randrange(interval * i, (interval * (i + 1)) - 1) for i in range(conf.ptype_num))
-                    ))
+                    interval = (conf.end_frame - conf.start_frame) // conf.ptype_num
+                    samples = list(
+                        map(
+                            Frame,
+                            list(
+                                conf.start_frame + random.randrange(interval * i, (interval * (i + 1)) - 1)
+                                for i in range(conf.ptype_num)
+                            )
+                        )
+                    )
                 else:
                     logging.info('Making samples according to specified picture types...')
                     samples = self._select_samples_ptypes(conf)