subocr/models.py at main · ezatgz/subocr · 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
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
import os
from PySide6.QtCore import QAbstractTableModel, QModelIndex, Qt


class SubtitleItem:
    """字幕项数据模型"""

    def __init__(self, file_path, start_time, end_time, ocr_text="", output_directory=""):
        self.file_path = file_path
        self.start_time = start_time
        self.end_time = end_time
        self.ocr_text = ocr_text
        self.edited_text = ocr_text  # 用户编辑后的文本
        self.output_directory = output_directory  # 输出目录

    def get_filename(self):
        """获取文件名"""
        return os.path.basename(self.file_path)

    def get_txt_path(self):
        """获取对应txt文件路径"""
        txt_filename = os.path.splitext(self.get_filename())[0] + '.txt'
        if self.output_directory:
            # 如果指定了输出目录，则保存到输出目录
            return os.path.join(self.output_directory, txt_filename)
        else:
            # 否则保存到原文件相同目录
            return os.path.join(os.path.dirname(self.file_path), txt_filename)

    def save_text(self):
        """保存文本到txt文件"""
        try:
            txt_path = self.get_txt_path()
            # 确保输出目录存在
            os.makedirs(os.path.dirname(txt_path), exist_ok=True)
            with open(txt_path, 'w', encoding='utf-8') as f:
                f.write(self.edited_text)
            return True
        except Exception as e:
            return False

    def update_text(self, new_text):
        """更新文本内容"""
        self.edited_text = new_text


class SubtitleTableModel(QAbstractTableModel):
    def __init__(self, data=None):
        super().__init__()
        self._data = data or []
        self.headers = ["时间戳", "操作", "内容"]

    def rowCount(self, parent=QModelIndex()):
        return len(self._data)

    def columnCount(self, parent=QModelIndex()):
        return len(self.headers)

    def headerData(self, section, orientation, role):
        if role == Qt.DisplayRole and orientation == Qt.Horizontal:
            return self.headers[section]
        return None

    def data(self, index, role):
        if not index.isValid() or index.row() >= len(self._data):
            return None

        row = index.row()
        col = index.column()
        item = self._data[row]

        if role == Qt.DisplayRole:
            if col == 0:  # 时间戳列
                return f"{item.start_time}\n{item.end_time}"
            elif col == 2:  # 内容列
                return item.edited_text
        elif role == Qt.TextAlignmentRole:
            if col == 0:  # 时间戳列居中
                return Qt.AlignCenter
        return None

    def setData(self, index, value, role=Qt.EditRole):
        if not index.isValid() or index.row() >= len(self._data):
            return False

        row = index.row()
        col = index.column()
        item = self._data[row]

        if role == Qt.EditRole:
            if col == 2:  # 内容列
                item.update_text(value)
                self.dataChanged.emit(index, index)
                return True
        return False

    def flags(self, index):
        if not index.isValid():
            return Qt.NoItemFlags

        if index.column() == 2:  # 内容列可编辑
            return Qt.ItemIsEnabled | Qt.ItemIsSelectable | Qt.ItemIsEditable
        return Qt.ItemIsEnabled | Qt.ItemIsSelectable

    def add_item(self, item):
        self.beginInsertRows(QModelIndex(), len(self._data), len(self._data))
        self._data.append(item)
        self.endInsertRows()

    def remove_item(self, row):
        if 0 <= row < len(self._data):
            self.beginRemoveRows(QModelIndex(), row, row)
            del self._data[row]
            self.endRemoveRows()