sample_tools/screen_caps.py at master · robotmitchum/sample_tools · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
# coding:utf-8
"""
    :module: screen_caps.py
    :description: Quick and dirty batch screen capture to update the documentation more easily
    :author: Michel 'Mitch' Pecqueur
    :date: 2025.04
"""

import importlib
import os
import sys
from functools import partial
from pathlib import Path

from PyQt5.QtCore import *
from PyQt5.QtWidgets import *

from dark_fusion_style import apply_dark_theme


class ScreenCap(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle('Capture UI')
        self.setGeometry(100, 100, 300, 200)

        self.setLocale(QLocale(QLocale.English, QLocale.UnitedStates))

        self.current_dir = Path(__file__).parent
        self.tools_path = self.current_dir / 'tools'

        self.caps_path = self.current_dir / 'screencaps'
        if not self.caps_path.exists():
            os.makedirs(self.caps_path, exist_ok=True)

        if self.tools_path.exists():
            print(f'tools sub-directory successfully found: {str(self.tools_path)}')
            sys.path.append(str(self.tools_path))

        self.scripts = ['sample_tools_UI.py', 'smp2ds_UI.py', 'drds_UI.py', 'split_tool_UI.py', 'rename_tool_UI.py',
                        'loop_tool_UI.py', 'st_tool_UI.py', 'mutate_tool_UI.py', 'upsample_tool_UI.py']

        self.script_modules = {}

        self.batch_capture()

        QTimer.singleShot(1000, partial(self.close))

    def batch_capture(self):
        for script in self.scripts:
            module_name = str(Path(script).stem)
            try:
                self.script_modules[module_name] = importlib.import_module(module_name)
                print(f'{module_name}: loaded')
                mod = self.script_modules[module_name]
                wid = mod.run(parent=self)
                filename = module_name.lower()

                # Clear output path if applicable
                if script not in ['sample_tools_UI.py', 'smp2ds_UI.py', 'drds_UI.py']:
                    try:
                        wid.output_path_l.setFullPath('')
                    except Exception as e:
                        print(e)

                QTimer.singleShot(100, partial(self.capture, wid, filename))
            except Exception as e:
                print(f'Failed to import {module_name}: {e}')

    def capture(self, widget, filename):
        pixmap = widget.grab()
        filepath = self.caps_path / f'{filename.lower()}.png'
        filepath = str(filepath.resolve())
        pixmap.save(filepath, 'PNG')


if __name__ == '__main__':
    app = QApplication(sys.argv)
    apply_dark_theme(app)

    font = app.font()
    font.setPointSize(11)
    app.setFont(font)

    window = ScreenCap()
    window.show()

    sys.exit(app.exec_())