arcaea-offline-ocr-model/ui/extends/samplesListWidget.py

70 lines
2.5 KiB
Python

import json
from pathlib import Path
from PySide6.QtCore import QByteArray, QMimeData, Qt
from PySide6.QtGui import QDrag, QPixmap
from PySide6.QtWidgets import QListWidget, QListWidgetItem, QMessageBox, QProgressDialog
class SamplesListWidget(QListWidget):
PathlibPathRole = Qt.ItemDataRole.UserRole
def __init__(self, parent=None):
super().__init__(parent)
self.setViewMode(QListWidget.ViewMode.IconMode)
self.setEditTriggers(QListWidget.EditTrigger.NoEditTriggers)
self.setMovement(QListWidget.Movement.Static)
self.setDragDropMode(QListWidget.DragDropMode.DragOnly)
self.setDragEnabled(True)
self.setSelectionMode(QListWidget.SelectionMode.MultiSelection)
def setSamples(self, samples: list[Path]):
self.clear()
samplesNum = len(samples)
progressDialog = QProgressDialog("", "Abort", 0, samplesNum, self)
progressDialog.setWindowModality(Qt.WindowModality.ApplicationModal)
for i, sample in enumerate(samples):
item = QListWidgetItem(QPixmap(str(sample)), f"{sample.stem[:3]}...", self)
item.setData(self.PathlibPathRole, sample)
self.addItem(item)
if samplesNum >= 1000:
updateInterval = 100
elif samplesNum >= 100:
updateInterval = 10
else:
updateInterval = 1
if i % updateInterval == 0:
progressDialog.setValue(i)
progressDialog.setLabelText(f"{i + 1}/{samplesNum}")
if progressDialog.wasCanceled():
break
progressDialog.setValue(samplesNum)
QMessageBox.information(
self, None, f"Loaded {self.model().rowCount()} samples."
)
def startDrag(self, supportedActions: Qt.DropAction):
drag = QDrag(self)
items = self.selectedItems()
paths = [str(item.data(self.PathlibPathRole).resolve()) for item in items]
mimeDataString = json.dumps(paths, ensure_ascii=False)
mimeData = QMimeData()
mimeData.setData(
"application/ao-ocr-model_sample",
QByteArray(mimeDataString.encode("utf-8")),
)
drag.setPixmap(items[0].icon().pixmap(items[0].icon().availableSizes()[0]))
drag.setMimeData(mimeData)
if drag.exec(Qt.DropAction.MoveAction) == Qt.DropAction.MoveAction:
for item in items:
index = self.indexFromItem(item)
self.model().removeRow(index.row())