@@ -25,6 +25,7 @@ from __future__ import annotations
import contextlib
import os
import shutil
import time
from pathlib import Path
from PySide6 . QtCore import Qt , QThreadPool
@@ -32,9 +33,12 @@ from PySide6.QtGui import QAction, QBrush, QColor, QFont, QKeySequence, QShortcu
from PySide6 . QtWidgets import (
QAbstractItemView ,
QApplication ,
QComboBox ,
QDialog ,
QDialogButtonBox ,
QDoubleSpinBox ,
QFileDialog ,
QFormLayout ,
QHBoxLayout ,
QInputDialog ,
QLabel ,
@@ -46,6 +50,7 @@ from PySide6.QtWidgets import (
QPlainTextEdit ,
QProgressBar ,
QPushButton ,
QSizePolicy ,
QSplitter ,
QTableWidget ,
QTableWidgetItem ,
@@ -89,7 +94,9 @@ class MainWindow(QMainWindow):
self . _restorer_key = None
self . _restored : dict [ str , object ] = { } # path -> restored image (BGR ndarray)
self . _restored_count = 0 # frames with output in restored/ (for the progress summary)
self . _row_restored : set [ str ] = set ( ) # frame paths that have a restored version (row ✓)
self . _showing_restored = False
self . _filter_mode = " all " # file-list filter (see filter_combo)
self . _nav_sync = False # guard against slider<->list signal loops
self . _busy = False # a long operation is running
self . _cancel = False # the user asked to stop it
@@ -126,6 +133,8 @@ class MainWindow(QMainWindow):
file_menu . addAction ( " Движок восстановления… " , self . _open_restore_settings )
file_menu . addAction ( " Расцензурить все (дозапуск) " , lambda : self . _restore_all ( False ) )
file_menu . addAction ( " Расцензурить все заново " , lambda : self . _restore_all ( True ) )
file_menu . addAction ( " Расцензурить найденное (по детекции) " , lambda : self . _restore_all ( only_detected = True ) )
file_menu . addAction ( " Открыть папку результатов " , self . _open_restored_dir )
file_menu . addSeparator ( )
file_menu . addAction ( " В избранное" , self . _move_to_favorites ) . setShortcut ( " Ctrl+M " )
file_menu . addSeparator ( )
@@ -134,14 +143,19 @@ class MainWindow(QMainWindow):
def _build_toolbar ( self ) - > None :
tb = self . addToolBar ( " Главная " )
tb . setMovable ( False )
tb . setToolButtonStyle ( Qt . ToolButtonTextOnly )
tb . addAction ( QAction ( " Создать проект… " , self , triggered = self . _create_project ) )
tb . addAction ( QAction ( " Открыть п роект… " , self , triggered = self . _open_project_dialog ) )
from_video = QAction ( " Создать из ролика… " , self , triggered = self . _create_from_video )
from_video . setToolTip ( " Разложить видео на кадры в новый проект и открыть его " )
tb . addAction ( from_video )
# --- Проект: rarely-touched session actions collapsed into one dropdown.
project_btn = self . _dropdown_button ( " П роект ▾ " , " Действия с проектом " )
m = project_btn . menu ( )
m . addAction ( " Создать проект… " , self . _create_project )
m . addAction ( " Открыть проект… " , self . _open_project_dialog )
m . addAction ( " Создать из ролика… " , self . _create_from_video )
m . addAction ( " Импортировать папку как проект… " , self . _import_folder_as_project )
tb . addWidget ( project_btn )
tb . addSeparator ( )
# --- Детекторы: active YOLO models picker (unchanged).
tb . addWidget ( QLabel ( " Детекторы: " ) )
self . _models_menu = QMenu ( self )
self . models_button = QToolButton ( )
@@ -150,41 +164,62 @@ class MainWindow(QMainWindow):
self . models_button . setToolTip ( " Выбрать активные YOLO-модели (models/yolo/<категория>) " )
tb . addWidget ( self . models_button )
self . _rebuild_models_menu ( )
tb . addSeparator ( )
calc = QAction ( " Рассчитать кадр " , self , triggered = self . _recompute_current )
calc . setToolTip ( " Запустить детектор на выбранном кадре (Space / двойной клик по файлу) " )
tb . addAction ( calc )
detect_all = QAction ( " Детектировать все " , self , triggered = lambda : self . _detect_all ( False ) )
detect_all . setToolTip ( " Рассчитать все ещё не посчитанные кадры (дозапуск; кэш сохраняется) " )
tb . addAction ( detect_all )
regen = QAction ( " Все заново " , self , triggered = lambda : self . _detect_all ( True ) )
regen . setToolTip ( " Очистить кэш детекций и пересчитать всю папку заново " )
tb . addAction ( regen )
# --- Детекция: primary "Рассчитать кадр" + the bulk runs in its dropdown.
detect_btn = self . _split_button (
" Рассчитать кадр " , self . _recompute_current ,
" Запустить детектор на выбранном кадре (Space / двойной клик по файлу) " ,
)
dm = detect_btn . menu ( )
dm . addAction ( " Детектировать все (дозапуск) " , lambda : self . _detect_all ( False ) )
dm . addAction ( " Детектировать все заново " , lambda : self . _detect_all ( True ) )
tb . addWidget ( detect_btn )
# --- Расцензурить: primary "Расцензурить кадр" + the bulk/engine actions.
restore_btn = self . _split_button (
" Расцензурить кадр " , self . _restore_current ,
" Восстановить мозаику на текущем кадре (результат сохраняется в restored/) " ,
)
rm = restore_btn . menu ( )
rm . addAction ( " Расцензурить все (дозапуск) " , lambda : self . _restore_all ( False ) )
hits = rm . addAction ( " Расцензурить найденное (по детекции) " , lambda : self . _restore_all ( only_detected = True ) )
hits . setToolTip (
" Расцензурить только кадры с детекцией (быстро, пропускает чистые). \n "
" ВНИМАНИЕ: YOLO ловит не всю мозаику — может пропустить. "
)
rm . addAction ( " Расцензурить все заново " , lambda : self . _restore_all ( True ) )
rm . addSeparator ( )
rm . addAction ( " Движок восстановления… " , self . _open_restore_settings )
rm . addAction ( " Открыть папку результатов " , self . _open_restored_dir )
rm . addSeparator ( )
self . save_restored_action = QAction ( " Сохранить результат… " , self , triggered = self . _save_restored )
self . save_restored_action . setEnabled ( False )
self . save_restored_action . setToolTip ( " Экспортировать <имя>_restored.jpg рядом с кадром " )
rm . addAction ( self . save_restored_action )
tb . addWidget ( restore_btn )
# --- View toggle: kept visible (checkable) so the result/original state is obvious.
self . toggle_restored_action = QAction ( " Показать расцензуренное " , self , triggered = self . _toggle_restored )
self . toggle_restored_action . setCheckable ( True )
self . toggle_restored_action . setEnabled ( False )
self . toggle_restored_action . setShortcut ( " R " )
self . toggle_restored_action . setToolTip (
" Переключить просмотр оригинал ⇄ расцензуренное для всего проекта (R) "
)
tb . addAction ( self . toggle_restored_action )
tb . addSeparator ( )
# --- Stop: kept visible — must be reachable instantly during a long run.
self . stop_action = QAction ( " ■ Стоп " , self , triggered = self . _request_cancel )
self . stop_action . setToolTip ( " Отменить текущую операцию (Esc) " )
self . stop_action . setEnabled ( False )
tb . addAction ( self . stop_action )
tb . addSeparator ( )
restore = QAction ( " Расцензурить кадр " , self , triggered = self . _restore_current )
restore . setToolTip ( " Восстановить найденные области на текущем кадре " )
tb . addAction ( restore )
restore_all = QAction ( " Расцензурить все " , self , triggered = lambda : self . _restore_all ( False ) )
restore_all . setToolTip ( " Расцензурить все кадры в папку restored/ (дозапуск; видеодвижок — весь диапазон) " )
tb . addAction ( restore_all )
restore_regen = QAction ( " Все заново (расцензур) " , self , triggered = lambda : self . _restore_all ( True ) )
restore_regen . setToolTip ( " Перерасцензурить все кадры заново (перезапись restored/) " )
tb . addAction ( restore_regen )
self . toggle_restored_action = QAction ( " Показать оригинал " , self , triggered = self . _toggle_restored )
self . toggle_restored_action . setEnabled ( False )
tb . addAction ( self . toggle_restored_action )
self . save_restored_action = QAction ( " Сохранить результат " , self , triggered = self . _save_restored )
self . save_restored_action . setEnabled ( False )
tb . addAction ( self . save_restored_action )
tb . addSeparator ( )
# Push the threshold control to the right edge.
spacer = QWidget ( )
spacer . setSizePolicy ( QSizePolicy . Expanding , QSizePolicy . Preferred )
tb . addWidget ( spacer )
tb . addWidget ( QLabel ( " Порог: " ) )
self . threshold_spin = QDoubleSpinBox ( )
self . threshold_spin . setRange ( 0.0 , 1.0 )
@@ -193,12 +228,47 @@ class MainWindow(QMainWindow):
self . threshold_spin . valueChanged . connect ( self . _on_threshold_changed )
tb . addWidget ( self . threshold_spin )
def _dropdown_button ( self , text : str , tooltip : str ) - > QToolButton :
""" A toolbar button that just opens a menu (no default action). """
btn = QToolButton ( )
btn . setText ( text )
btn . setToolTip ( tooltip )
btn . setPopupMode ( QToolButton . InstantPopup )
btn . setMenu ( QMenu ( btn ) )
return btn
def _split_button ( self , text : str , slot , tooltip : str ) - > QToolButton :
""" A split button: click runs ``slot``; the arrow opens a menu of related actions. """
btn = QToolButton ( )
btn . setText ( text )
btn . setToolTip ( tooltip )
btn . setPopupMode ( QToolButton . MenuButtonPopup )
action = QAction ( text , btn , triggered = slot )
action . setToolTip ( tooltip )
btn . setDefaultAction ( action )
btn . setMenu ( QMenu ( btn ) )
return btn
def _build_central ( self ) - > None :
self . file_list = QListWidget ( )
self . file_list . setSelectionMode ( QAbstractItemView . ExtendedSelection ) # multi-select for moves
self . file_list . currentItemChanged . connect ( self . _on_file_selected )
self . file_list . itemDoubleClicked . connect ( self . _on_file_activated )
# Filter the (possibly huge) list down to the frames you care about.
self . filter_combo = QComboBox ( )
for label , mode in (
( " Все кадры " , " all " ) ,
( " С цензурой" , " hits " ) ,
( " Чистые " , " clean " ) ,
( " Не рассчитано " , " uncomputed " ) ,
( " Расцензуренные " , " restored " ) ,
( " Без расцензуривания " , " unrestored " ) ,
) :
self . filter_combo . addItem ( label , mode )
self . filter_combo . setToolTip ( " Показывать только кадры выбранной категории " )
self . filter_combo . currentIndexChanged . connect ( self . _on_filter_changed )
# One default collection ("Избранное"); the button acts on the list selection.
move_btn = QPushButton ( " ★ В избранное " )
move_btn . setToolTip ( " Переместить выбранные кадры в избранное проекта (Ctrl+M) " )
@@ -208,6 +278,7 @@ class MainWindow(QMainWindow):
left_layout = QVBoxLayout ( left )
left_layout . setContentsMargins ( 4 , 4 , 4 , 4 )
left_layout . setSpacing ( 4 )
left_layout . addWidget ( self . filter_combo )
left_layout . addWidget ( self . file_list , 1 )
left_layout . addWidget ( move_btn )
@@ -269,9 +340,14 @@ class MainWindow(QMainWindow):
self . frame_slider . setToolTip ( " Перемотка по кадрам (стрелки ← →); метки — кадры с детекцией " )
self . frame_slider . valueChanged . connect ( self . _on_slider )
self . pos_label = QLabel ( " 0 / 0 " )
# Clickable position readout — opens "go to frame N" (handy on huge sequences
# where the scrubber is too coarse, ~50 frames/px on 29k).
self . pos_label = QPushButton ( " 0 / 0 " )
self . pos_label . setFlat ( True )
self . pos_label . setMinimumWidth ( 90 )
self . pos_label . setAlignment ( Qt . AlignCenter )
self . pos_label . setCursor ( Qt . PointingHandCursor )
self . pos_label . setToolTip ( " Перейти к кадру по номеру " )
self . pos_label . clicked . connect ( self . _jump_to_frame )
self . prev_hit_btn = QPushButton ( " ◀ детекция " )
self . prev_hit_btn . setToolTip ( " Предыдущий кадр с детекцией ([) " )
@@ -297,8 +373,13 @@ class MainWindow(QMainWindow):
n = self . file_list . count ( )
if n == 0 :
return
row = max ( 0 , min ( n - 1 , self . file_list . currentRow ( ) + delta ) )
self . file_list . setCurrentRow ( row )
# Skip rows hidden by the filter so prev/next walk only the visible frames.
i = self . file_list . currentRow ( ) + delta
while 0 < = i < n :
if not self . file_list . item ( i ) . isHidden ( ) :
self . file_list . setCurrentRow ( i )
return
i + = delta
def _step_hit ( self , direction : int ) - > None :
""" Jump to the nearest frame (in `direction`) that has detections. """
@@ -318,6 +399,18 @@ class MainWindow(QMainWindow):
" (сначала «Детектировать все») "
)
def _jump_to_frame ( self ) - > None :
""" Ask for a 1-based frame number and select it (clamped to range). """
n = self . file_list . count ( )
if n == 0 :
return
cur = self . file_list . currentRow ( ) + 1
num , ok = QInputDialog . getInt (
self , " Перейти к кадру " , f " Номер кадра (1– { n } ): " , cur , 1 , n
)
if ok :
self . file_list . setCurrentRow ( num - 1 )
def _on_slider ( self , value : int ) - > None :
if self . _nav_sync :
return
@@ -469,6 +562,7 @@ class MainWindow(QMainWindow):
"""
self . _begin_busy ( total )
self . _tick_count = 0
self . _job_start = time . monotonic ( ) # for ETA in the progress messages
job = Job ( fn )
self . _job = job
if on_tick is not None :
@@ -483,7 +577,26 @@ class MainWindow(QMainWindow):
self . progress . setRange ( 0 , total )
self . progress . setValue ( done )
if message :
self . statusBar ( ) . showMessage ( message )
eta = self . _eta_suffix ( done , total )
self . statusBar ( ) . showMessage ( message + eta )
def _eta_suffix ( self , done : int , total : int ) - > str :
""" ' · осталось ~Xм Yс ' estimated from the average rate so far (or ' ' if N/A). """
start = getattr ( self , " _job_start " , None )
if not start or done < = 0 or total < = 0 or done > = total :
return " "
elapsed = time . monotonic ( ) - start
if elapsed < 0.5 :
return " "
remaining = elapsed / done * ( total - done )
secs = int ( remaining )
if secs > = 3600 :
text = f " { secs / / 3600 } ч { secs % 3600 / / 60 } м "
elif secs > = 60 :
text = f " { secs / / 60 } м { secs % 60 } с "
else :
text = f " { secs } с "
return f " · осталось ~ { text } "
def _finish_job ( self , result , on_done ) - > None :
cancelled = self . _job . cancelled if self . _job is not None else False
@@ -500,7 +613,10 @@ class MainWindow(QMainWindow):
# --------------------------------------------------------------- detector
def _make_detector ( self ) :
d = self . _cfg . detection
key = ( tuple ( sorted ( self . _cfg . detector_models ) ) , d . yolo_conf , d . yolo_imgsz )
key = (
tuple ( sorted ( self . _cfg . detector_models ) ) , d . yolo_conf , d . yolo_imgsz ,
self . _cfg . cross_model_nms , self . _cfg . nms_iou ,
)
if key != self . _detector_key :
self . _detector = build_detector ( self . _cfg ) # may raise ValueError / import / file errors
self . _detector_key = key
@@ -540,11 +656,74 @@ class MainWindow(QMainWindow):
act . setChecked ( e . path in selected )
act . toggled . connect ( lambda on , p = e . path : self . _on_model_toggled ( p , on ) )
self . _models_menu . addSeparator ( )
nms = self . _models_menu . addAction ( " Объединять пересечения (NMS) " )
nms . setCheckable ( True )
nms . setChecked ( self . _cfg . cross_model_nms )
nms . setToolTip (
" Убирать дублирующие рамки от перекрывающихся моделей (по IoU; остаётся \n "
" рамка с большей уверенностью). Меняет результат — кэш пересчитывается. "
)
nms . toggled . connect ( self . _on_nms_toggled )
self . _models_menu . addAction ( " Пороги по моделям… " , self . _edit_model_thresholds )
self . _models_menu . addSeparator ( )
self . _models_menu . addAction ( " Добавить модель… " , self . _add_model )
self . _models_menu . addAction ( " Открыть папку моделей " , self . _open_models_dir )
self . _models_menu . addAction ( " Обновить список " , self . _rebuild_models_menu )
self . _update_models_button ( )
def _on_nms_toggled ( self , on : bool ) - > None :
self . _cfg . cross_model_nms = on
self . _persist_settings ( )
self . _invalidate_results ( ) # merging changes detections => recompute
self . statusBar ( ) . showMessage (
" Объединение пересечений (NMS): " + ( " вкл " if on else " выкл " )
)
def _edit_model_thresholds ( self ) - > None :
""" Dialog: per-model overlay threshold overrides (display-only, not detection). """
models = [ m for m in self . _cfg . detector_models if Path ( m ) . is_file ( ) ]
if not models :
QMessageBox . information (
self , " Нет моделей " , " Сначала отметьте хотя бы одну модель. "
)
return
dlg = QDialog ( self )
dlg . setWindowTitle ( " Пороги отображения по моделям " )
layout = QVBoxLayout ( dlg )
layout . addWidget ( QLabel (
f " Порог отображения для каждой модели (по умолчанию { self . _cfg . default_threshold : .2f } ). \n "
" Влияет только на отрисовку/подсветку, не на саму детекцию. "
) )
form = QFormLayout ( )
spins : dict [ str , QDoubleSpinBox ] = { }
for m in models :
stem = Path ( m ) . stem
spin = QDoubleSpinBox ( )
spin . setRange ( 0.0 , 1.0 )
spin . setSingleStep ( 0.05 )
spin . setValue ( self . _cfg . model_thresholds . get ( stem , self . _cfg . default_threshold ) )
spins [ stem ] = spin
form . addRow ( stem , spin )
layout . addLayout ( form )
buttons = QDialogButtonBox ( QDialogButtonBox . Ok | QDialogButtonBox . Cancel )
buttons . accepted . connect ( dlg . accept )
buttons . rejected . connect ( dlg . reject )
layout . addWidget ( buttons )
if dlg . exec ( ) != QDialog . Accepted :
return
# Store only the overrides that differ from the global default — keeps it tidy.
thresholds = {
stem : round ( spin . value ( ) , 4 )
for stem , spin in spins . items ( )
if abs ( spin . value ( ) - self . _cfg . default_threshold ) > 1e-9
}
self . _cfg . model_thresholds = thresholds
self . view . set_model_thresholds ( thresholds )
self . _persist_settings ( )
self . statusBar ( ) . showMessage (
f " Пороги по моделям обновлены ( { len ( thresholds ) } переопределений) "
)
def _update_models_button ( self ) - > None :
n = len ( [ m for m in self . _cfg . detector_models if Path ( m ) . is_file ( ) ] )
self . models_button . setText ( f " Модели ( { n } ) ▾ " )
@@ -592,6 +771,16 @@ class MainWindow(QMainWindow):
with contextlib . suppress ( OSError , AttributeError ) :
os . startfile ( str ( root ) ) # noqa: S606 - Windows: open in Explorer
def _open_restored_dir ( self ) - > None :
""" Open the project ' s restored/ folder (where restored images are saved). """
if self . _project is None :
QMessageBox . information ( self , " Нет проекта " , " Сначала откройте проект. " )
return
d = self . _project . restored_dir
d . mkdir ( parents = True , exist_ok = True )
with contextlib . suppress ( OSError , AttributeError ) :
os . startfile ( str ( d ) ) # noqa: S606 - Windows: open in Explorer
def _invalidate_results ( self ) - > None :
""" Detector changed — drop the in-memory cache and refresh the current image.
@@ -775,6 +964,7 @@ class MainWindow(QMainWindow):
self . threshold_spin . setValue ( self . _cfg . default_threshold )
self . threshold_spin . blockSignals ( False )
self . view . set_threshold ( self . _cfg . default_threshold )
self . view . set_model_thresholds ( self . _cfg . model_thresholds )
def _persist_settings ( self ) - > None :
""" Save settings to the global defaults and (if open) into the project. """
@@ -795,7 +985,7 @@ class MainWindow(QMainWindow):
dialog = ExtractDialog ( self )
if dialog . exec ( ) != QDialog . Accepted :
return
keyframes_only , step , max_dim = dialog . options ( )
keyframes_only , step , max_dim , jpg_quality = dialog . options ( )
video = Path ( path )
root = video . parent / f " { video . stem } _frames "
if Project . is_project ( root ) :
@@ -821,7 +1011,7 @@ class MainWindow(QMainWindow):
try :
saved = extract_frames (
str ( video ) , str ( out ) , step = step , keyframes_only = keyframes_only ,
max_dim = max_dim , progress = cb ,
max_dim = max_dim , jpg_quality = jpg_quality , progress = cb ,
)
except Exception as exc :
QMessageBox . warning ( self , " Ошибка " , f " Не удалось извлечь кадры: \n { exc } " )
@@ -849,6 +1039,8 @@ class MainWindow(QMainWindow):
files = sorted ( p for p in folder . iterdir ( ) if p . suffix . lower ( ) in _IMAGE_EXTS )
self . _files = files
self . _results . clear ( )
self . _row_restored . clear ( )
self . _showing_restored = False # start a project in original-view mode
self . _current = None
self . file_list . blockSignals ( True )
@@ -926,7 +1118,7 @@ class MainWindow(QMainWindow):
return
key , dets = result
self . _results [ key ] = dets
self . _tag_file ( Path ( key ) , len ( dets ) )
self . _tag_file ( Path ( key ) )
self . _refresh_marks ( )
self . _save_results ( )
if then_show or self . _current == Path ( key ) :
@@ -937,13 +1129,22 @@ class MainWindow(QMainWindow):
self . _start_job ( fn , None , on_done = done )
def _show ( self , path : Path ) - > None :
""" Display the image with its cached detections (does not run the detector). """
""" Display the image (does not run the detector).
Honours the global " show restored " view mode (``_showing_restored``): when on and
a restored version exists (memory or ``restored/``), the restored image is shown
(no overlays); otherwise the original frame with its cached detections.
"""
self . _current = path
self . _showing_restored = False
img = imread_unicode ( str ( path ) )
dets = self . _results . get ( str ( path ) ) # None => not yet computed
self . view . set_image ( img , dets or [ ] )
self . _fill_detail_table ( path , img , dets )
restored = self . _restored_image_for ( path ) if self . _showing_restored else None
if restored is not None :
self . view . set_image ( restored , [ ] ) # restored: no overlays
self . _fill_detail_table ( path , restored , dets )
else :
img = imread_unicode ( str ( path ) )
self . view . set_image ( img , dets or [ ] )
self . _fill_detail_table ( path , img , dets )
self . _update_restore_actions ( )
def _recompute_current ( self ) - > None :
@@ -1000,7 +1201,7 @@ class MainWindow(QMainWindow):
""" GUI-thread handler for one streamed detect-all result. """
key , dets = payload
self . _results [ key ] = dets
self . _tag_file ( Path ( key ) , len ( dets ) )
self . _tag_file ( Path ( key ) )
# If the frame being viewed was just computed, show its overlay live.
if not self . _showing_restored and self . _current is not None and str ( self . _current ) == key :
self . _show ( self . _current )
@@ -1036,28 +1237,61 @@ class MainWindow(QMainWindow):
return
_ , k , restored , engine = result
self . _restored [ k ] = restored
# Persist to the project's restored/ folder (like the batch run), so a single
# restore is saved on disk and survives reopening — not just held in memory.
saved_to = " "
if self . _project is not None :
try :
self . _project . restored_dir . mkdir ( parents = True , exist_ok = True )
dst = self . _project . restored_dir / f " { Path ( k ) . stem } .jpg "
if imwrite_unicode ( str ( dst ) , restored ) :
saved_to = f " → { self . _project . restored_dir . name } / "
except OSError :
pass
self . _refresh_restored_marks ( )
if self . _current is not None and str ( self . _current ) == k :
self . _showing_restored = True
self . view . set_image ( restored , [ ] )
self . _update_restore_actions ( )
self . _refresh_restored_marks ( )
self . statusBar ( ) . showMessage ( f " Расцензурено ( { engine } ): { Path ( k ) . name } " )
self . _show ( self . _current )
self . statusBar ( ) . showMessage ( f " Расцензурено ( { engine } ): { Path ( k ) . name } { saved_to } " )
self . statusBar ( ) . showMessage ( f " Восстановление: { path . name } … " )
self . _start_job ( fn , None , on_done = done )
def _restore_all ( self , force : bool = False ) - > None :
""" Restore every frame on a background thread, writing results to ``restored/``.
def _restore_all ( self , force : bool = False , only_detected : bool = False ) - > None :
""" Restore frames on a background thread, writing results to ``restored/``.
DeepMosaics locates the mosaic itself, so no detection runs here. The per-frame
engine skips frames already restored (resume) unless ``force``. The temporal
engine (DeepMosaics-video) runs the whole contiguous sequence in order via
engine (DeepMosaics-video) runs a contiguous sequence in order via
``restore_sequence`` (its recurrence needs neighbours), so ``force`` is implied.
``only_detected`` uses the YOLO detection cache to skip frames known clean:
per-frame → restore just the frames with detections; temporal → restrict the run
to the contiguous span [first hit … last hit] (clean frames inside it still run,
for recurrence). NOTE: LADA misses some mosaic, so this can miss censorship YOLO
didn ' t flag — " Расцензурить все " stays the thorough option.
"""
if not self . _files or self . _project is None or self . _busy :
return
files = list ( self . _files ) # snapshot — favorites/move mutate self._files
total = len ( files )
is_temporal = self . _cfg . restorer == " deepmosaics_video "
hits = [ i for i , p in enumerate ( files ) if self . _results . get ( str ( p ) ) ]
if only_detected :
if not self . _results :
self . statusBar ( ) . showMessage (
" Детекция не посчитана — сначала «Детектировать все» (или «Расцензурить все») "
)
return
if not hits :
self . statusBar ( ) . showMessage ( " Цензура не найдена ни на одном кадре — нечего расцензуривать " )
return
span = ( hits [ 0 ] , hits [ - 1 ] ) if is_temporal else None
total = ( span [ 1 ] - span [ 0 ] + 1 ) if span else len ( hits )
else :
span = None
total = len ( files )
out_dir = self . _project . restored_dir
out_dir . mkdir ( parents = True , exist_ok = True )
@@ -1067,6 +1301,7 @@ class MainWindow(QMainWindow):
def fn ( job ) :
restorer = self . _make_restorer ( ) # built on the worker (may raise)
frame_cache : dict [ int , object ] = { } # small cache so the temporal window reuses reads
state = { " done " : 0 } # frames processed (for the progress bar)
def get_frame ( i ) :
img = frame_cache . get ( i )
@@ -1079,42 +1314,48 @@ class MainWindow(QMainWindow):
frame_cache [ i ] = img
return img
def emit ( i , restored ) :
imwrite_unicode ( str ( out_path ( files [ i ] ) ) , restored )
job . progress ( i + 1 , total , f " Расцензуривание { i + 1 } / { total } : { files [ i ] . name } " )
def emit ( i , restored , * , verb = " Расцензуривание " ) :
if restored is not None :
imwrite_unicode ( str ( out_path ( files [ i ] ) ) , restored )
state [ " done " ] + = 1
job . progress ( state [ " done " ] , total , f " { verb } { state [ ' done ' ] } / { total } : { files [ i ] . name } " )
if restorer . temporal :
start = span [ 0 ] if span else 0
end = span [ 1 ] if span else len ( files ) - 1
restorer . restore_sequence (
total , get_frame , lambda _i : [ ] , emit , should_cancel = lambda : job . cancelled
end - start + 1 ,
lambda li : get_frame ( start + li ) ,
lambda _li : [ ] ,
lambda li , res : emit ( start + li , res ) ,
should_cancel = lambda : job . cancelled ,
)
else :
for i , p in enumerate ( files ) :
indices = hits if only_detected else range ( len ( files ) )
for i in indices :
if job . cancelled :
break
p = files [ i ]
if not force and out_path ( p ) . is_file ( ) :
job . progress ( i + 1 , total , f " Пропуск { i + 1 } / { total } : { p . name } " )
emit ( i , None , verb = " Пропуск " ) # already restored — count, don't rewrite
continue
emit ( i , restorer . restore ( get_frame ( i ) , [ ] , should_cancel = lambda : job . cancelled ) )
frame_cache . pop ( i , None ) # per-frame: don't accumulate
return None
def done ( _result , cancelled ) :
if self . _current is not None : # live-preview the current frame's result, if any
rp = out_path ( self . _current )
if rp . is_file ( ) :
img = imread_unicode ( str ( rp ) )
if img is not None :
self . _restored [ str ( self . _current ) ] = img
self . _showing_restored = True
self . view . set_image ( img , [ ] )
self . _update_restore_actions ( )
self . _refresh_restored_marks ( )
if not cancelled and self . _row_restored :
self . _showing_restored = True # auto-switch to viewing the results
if self . _current is not None : # re-show current frame in the (new) mode
self . _show ( self . _current )
self . statusBar ( ) . showMessage (
" Расцензуривание отменено " if cancelled
else f " Готово: результаты в { out_dir . name } / ( { total } кадров) "
else f " Готово: результаты в { out_dir . name } / ( { total } кадров) — показываю расцензуренное "
)
self . statusBar ( ) . showMessage ( " Пакетное расцензуривание… " )
scope = " (по детекции) " if only_detected else " "
self . statusBar ( ) . showMessage ( f " Пакетное расцензуривание { scope } … " )
self . _start_job ( fn , total , on_done = done )
def _make_restorer ( self ) :
@@ -1133,30 +1374,58 @@ class MainWindow(QMainWindow):
self . _restorer_key = None # rebuild on next restore
self . statusBar ( ) . showMessage ( f " Движок восстановления: { self . _cfg . restorer } " )
# ---------------------------------------------------------- restored access
def _restored_disk_path ( self , path : Path ) - > Path | None :
""" Path of the restored output for ``path`` in ``restored/``, if it exists. """
if self . _project is None :
return None
rp = self . _project . restored_dir / f " { Path ( path ) . stem } .jpg "
return rp if rp . is_file ( ) else None
def _has_restored ( self , path : Path ) - > bool :
""" Cheap check (no decode): is there a restored version of ``path``? """
return str ( path ) in self . _restored or self . _restored_disk_path ( path ) is not None
def _restored_image_for ( self , path : Path ) :
""" Return the restored image for ``path`` (from memory or ``restored/``), or None. """
img = self . _restored . get ( str ( path ) )
if img is not None :
return img
rp = self . _restored_disk_path ( path )
return imread_unicode ( str ( rp ) ) if rp is not None else None
def _toggle_restored ( self ) - > None :
if self . _current is None or str ( self . _current ) not in self . _restored :
return
""" Flip the global view mode between original and restored, then re-show. """
self . _showing_restored = not self . _showing_restored
key = str ( self . _current )
if self . _showing_restored :
self . view . set_image ( self . _restored [ key ] , [ ] )
if self . _current is not None :
self . _show ( self . _current )
else :
self . view . set_image ( imread_unicode ( key ) , self . _results . get ( key ) or [ ] )
self . _update_restore_actions ( )
self . _update_restore_actions ( )
self . statusBar ( ) . showMessage (
" Показ: расцензуренное (где есть) " if self . _showing_restored else " Показ: оригинал "
)
def _update_restore_actions ( self ) - > None :
has = self . _current is not None and str ( self . _current ) in self . _restored
self . toggle_restored_action . setEnabled ( has )
any_restored = bool ( self . _row_restored ) or bool ( self . _restored )
self . toggle_restored_action . setEnabled ( any_restored )
# Keep the checkable state in sync with the mode (setChecked emits `toggled`,
# not `triggered`, so this never re-enters `_toggle_restored`).
self . toggle_restored_action . setChecked ( self . _showing_restored and any_restored )
self . toggle_restored_action . setText (
" Показать оригинал " if self . _showing_restored else " Показать р езультат "
" Показать оригинал " if self . _showing_restored else " Показать р асцензуренное "
)
self . save_restored_action . setEnabled (
self . _current is not None and self . _has_restored ( self . _current )
)
self . save_restored_action . setEnabled ( has )
def _save_restored ( self ) - > None :
if self . _current is None or str ( self . _current ) not in self . _restored :
if self . _current is None :
return
restored = self . _restored_image_for ( self . _current )
if restored is None :
return
out = self . _unique_dest ( self . _current . parent , f " { self . _current . stem } _restored.jpg " )
if imwrite_unicode ( str ( out ) , self . _restored [ str ( self . _current ) ] ) :
if imwrite_unicode ( str ( out ) , restored ) :
self . statusBar ( ) . showMessage ( f " Сохранено: { out } " )
else :
QMessageBox . warning ( self , " Ошибка " , " Не удалось сохранить файл. " )
@@ -1227,24 +1496,77 @@ class MainWindow(QMainWindow):
_TINT_HIT = QColor ( 200 , 80 , 80 , 70 )
_TINT_CLEAN = QColor ( 90 , 160 , 90 , 50 )
def _set_row_tag ( self , item : QListWidgetItem , count : int ) - > None :
base = item . data ( Qt . UserRole + 1 )
item . setText ( f " { base } · { count } " if count else f " { base } · — " )
item . setBackground ( self . _TINT_HIT if count else self . _TINT_CLEAN )
def _relabel_row ( self , item : QListWidgetItem ) - > None :
""" Set a row ' s text/tint/tooltip from its detection + restoration state.
def _tag_file ( self , path : Path , count : int ) - > None :
Text: ``name · <count|—> ✓`` — the count/— suffix appears once detected (red tint
= censorship, green = clean), and a trailing ✓ marks frames that have a restored
version in ``restored/``.
"""
base = item . data ( Qt . UserRole + 1 )
path = item . data ( Qt . UserRole )
restored = path in self . _row_restored
if path in self . _results : # `in`, not truthy: empty list = clean
dets = self . _results [ path ]
suffix = f " · { len ( dets ) } " if dets else " · — "
item . setBackground ( self . _TINT_HIT if dets else self . _TINT_CLEAN )
else :
suffix = " "
item . setBackground ( QBrush ( ) )
item . setText ( f " { base } { suffix } { ' ✓ ' if restored else ' ' } " )
item . setToolTip ( " Есть расцензуренная версия (restored/) " if restored else " " )
item . setHidden ( not self . _row_matches_filter ( path ) )
def _row_matches_filter ( self , path : str ) - > bool :
""" Whether a row should be visible under the current filter mode. """
mode = self . _filter_mode
if mode == " all " :
return True
if mode == " hits " :
return bool ( self . _results . get ( path ) )
if mode == " clean " :
return path in self . _results and not self . _results [ path ]
if mode == " uncomputed " :
return path not in self . _results
if mode == " restored " :
return path in self . _row_restored
if mode == " unrestored " :
return path not in self . _row_restored
return True
def _on_filter_changed ( self ) - > None :
self . _filter_mode = self . filter_combo . currentData ( ) or " all "
self . _relabel_all ( ) # re-applies hidden state per row
# If the current row got hidden, jump to the first visible one so the view isn't stale.
cur = self . file_list . currentItem ( )
if cur is not None and cur . isHidden ( ) :
for i in range ( self . file_list . count ( ) ) :
if not self . file_list . item ( i ) . isHidden ( ) :
self . file_list . setCurrentRow ( i )
break
n_vis = sum ( 1 for i in range ( self . file_list . count ( ) ) if not self . file_list . item ( i ) . isHidden ( ) )
self . statusBar ( ) . showMessage (
f " Фильтр: { self . filter_combo . currentText ( ) } — показано { n_vis } из { len ( self . _files ) } "
)
def _relabel_all ( self ) - > None :
for i in range ( self . file_list . count ( ) ) :
self . _relabel_row ( self . file_list . item ( i ) )
def _tag_file ( self , path : Path ) - > None :
for i in range ( self . file_list . count ( ) ) :
item = self . file_list . item ( i )
if item . data ( Qt . UserRole ) == str ( path ) :
self . _set_row_tag ( item , count )
self . _relabel_row ( item )
return
# ------------------------------------------------------------- result cache
def _results_key ( self ) - > dict :
""" Detector identity used to tag/validate the on-disk detection cache. """
d = self . _cfg . detection
nms_iou = self . _cfg . nms_iou if self . _cfg . cross_model_nms else None
return detection_cache . make_key (
self . _cfg . detector_models , d . yolo_conf , d . yolo_imgsz
self . _cfg . detector_models , d . yolo_conf , d . yolo_imgsz , nms_iou = nms_iou
)
def _save_results ( self ) - > None :
@@ -1265,20 +1587,16 @@ class MainWindow(QMainWindow):
if not cached :
return 0
self . _results = cached
for i in range ( self . file_list . count ( ) ) :
item = self . file_list . item ( i )
path = item . data ( Qt . UserRole )
if path in self . _results : # `in`, not truthy: empty list = checked-clean
self . _set_row_tag ( item , len ( self . _results [ path ] ) )
self . _relabel_all ( )
return len ( cached )
def _clear_results ( self ) - > None :
""" Drop all cached detections and reset row labels/tints (keeps the detector). """
""" Drop all cached detections and reset row labels/tints (keeps the detector).
Restored ✓ markers stay — restoration is independent of detection.
"""
self . _results . clear ( )
for i in range ( self . file_list . count ( ) ) :
item = self . file_list . item ( i )
item . setText ( item . data ( Qt . UserRole + 1 ) )
item . setBackground ( QBrush ( ) )
self . _relabel_all ( )
self . _refresh_marks ( )
def _refresh_marks ( self ) - > None :
@@ -1313,13 +1631,17 @@ class MainWindow(QMainWindow):
stems = { p . stem for p in self . _project . restored_dir . glob ( " *.jpg " ) }
mem = set ( self . _restored ) # single-frame restores held in memory (not on disk yet)
rows : set [ int ] = set ( )
paths : set [ str ] = set ( )
if stems or mem :
for i in range ( self . file_list . count ( ) ) :
fp = Path ( self . file_list . item ( i ) . data ( Qt . UserRole ) )
if fp . stem in stems or str ( fp ) in mem :
sp = self . file_list . item ( i ) . data ( Qt . UserRole )
if Path ( sp ) . stem in stems or sp in mem :
rows . add ( i )
paths . add ( sp )
self . _row_restored = paths
self . _restored_count = len ( rows )
self . frame_slider . set_restored_marks ( rows )
self . _relabel_all ( ) # show/refresh the ✓ markers in the file list
self . _update_counts_label ( )
def _fill_detail_table ( self , path : Path , img , dets : list [ Detection ] | None ) - > None :