stocksort.py

# stocksort.py
# Зависимости: customtkinter, rawpy, opencv-python, Pillow, numpy
# Сборка: pyinstaller --onefile --windowed stocksort.py

import customtkinter as ctk
import tkinter as tk
from tkinter import filedialog
import threading
import os
import shutil
import io

import cv2
import numpy as np
from PIL import Image

try:
    import rawpy
    _RAWPY_OK = True
except ImportError:
    _RAWPY_OK = False

# ── константы ────────────────────────────────────────────────────────────────

MODES = {
    0: ("Мягкий (Low)",     70),
    1: ("Стандарт (Medium)", 120),
    2: ("Строгий (Fine)",   180),
}
SCAN_EXTS = {".arw", ".jpg", ".jpeg", ".png"}


# ── класс приложения ─────────────────────────────────────────────────────────

class StockSortApp(ctk.CTk):

    def __init__(self):
        super().__init__()
        ctk.set_appearance_mode("dark")
        ctk.set_default_color_theme("blue")
        self.title("StockSort")
        self.geometry("360x240")
        self.resizable(False, False)
        self._build_ui()

    # ── построение интерфейса ─────────────────────────────────────────────

    def _build_ui(self):
        self.grid_columnconfigure(0, weight=1)

        # строка пути + кнопка «Обзор»
        pf = ctk.CTkFrame(self, fg_color="transparent")
        pf.grid(row=0, column=0, padx=10, pady=(10, 3), sticky="ew")
        pf.grid_columnconfigure(0, weight=1)

        self._path = tk.StringVar(value=".")
        self._entry = ctk.CTkEntry(pf, textvariable=self._path,
                                   height=26, font=("", 11))
        self._entry.grid(row=0, column=0, sticky="ew", padx=(0, 5))

        self._browse_btn = ctk.CTkButton(pf, text="Обзор", width=56, height=26,
                                          font=("", 11), command=self._browse)
        self._browse_btn.grid(row=0, column=1)

        # слайдер + метка режима
        sf = ctk.CTkFrame(self, fg_color="transparent")
        sf.grid(row=1, column=0, padx=10, pady=(2, 2), sticky="ew")
        sf.grid_columnconfigure(0, weight=1)

        self._slider = ctk.CTkSlider(sf, from_=0, to=2, number_of_steps=2,
                                      height=14, command=self._snap_slider)
        self._slider.set(1)
        self._slider.grid(row=0, column=0, sticky="ew", padx=(0, 6))

        self._mode_lbl = ctk.CTkLabel(sf, text=MODES[1][0], font=("", 11),
                                       text_color="gray", anchor="w", width=130)
        self._mode_lbl.grid(row=0, column=1)

        # прогресс-бар
        self._bar = ctk.CTkProgressBar(self, height=5)
        self._bar.set(0)
        self._bar.grid(row=2, column=0, padx=10, pady=(5, 5), sticky="ew")

        # кнопка «Отсортировать»
        self._sort_btn = ctk.CTkButton(self, text="Отсортировать", height=42,
                                        font=("", 14, "bold"),
                                        command=self._start)
        self._sort_btn.grid(row=3, column=0, padx=10, pady=(2, 4), sticky="ew")

        # строка статуса
        self._status = ctk.CTkLabel(self, text="Готово", font=("", 10),
                                     text_color="gray", anchor="center")
        self._status.grid(row=4, column=0, padx=10, pady=(0, 6), sticky="ew")

    # ── обработчики событий ───────────────────────────────────────────────

    def _browse(self):
        folder = filedialog.askdirectory(initialdir=self._path.get() or ".")
        if folder:
            self._path.set(folder)

    def _snap_slider(self, value):
        idx = round(float(value))
        self._slider.set(idx)
        self._mode_lbl.configure(text=MODES[idx][0])

    def _start(self):
        self._set_controls("disabled")
        self._bar.set(0)
        self._status.configure(text="Сканирование...")
        threading.Thread(target=self._worker, daemon=True).start()

    # ── рабочий поток (не блокирует GUI) ─────────────────────────────────

    def _worker(self):
        folder    = self._path.get()
        threshold = MODES[round(self._slider.get())][1]

        try:
            all_names = os.listdir(folder)
        except OSError as exc:
            self._ui(lambda: self._status.configure(text=f"Ошибка: {exc}"))
            self._ui(lambda: self._set_controls("normal"))
            return

        files = [n for n in all_names
                 if os.path.splitext(n)[1].lower() in SCAN_EXTS]

        if not files:
            self._ui(lambda: self._status.configure(text="Файлы не найдены"))
            self._ui(lambda: self._set_controls("normal"))
            return

        reject_dir = os.path.join(folder, "Брак")
        moved = 0
        total = len(files)

        for i, name in enumerate(files):
            path  = os.path.join(folder, name)
            score = _sharpness(path)

            if score is not None and score < threshold:
                os.makedirs(reject_dir, exist_ok=True)
                try:
                    shutil.move(path, os.path.join(reject_dir, name))
                    moved += 1
                except OSError:
                    pass

            prog = (i + 1) / total
            self._ui(lambda p=prog: self._bar.set(p))
            self._ui(lambda n=i + 1: self._status.configure(
                text=f"Обработка: {n}/{total}"))

        self._ui(lambda: self._status.configure(
            text=f"Очищено: {moved} кадров из {total}"))
        self._ui(lambda: self._set_controls("normal"))

    # ── вспомогательные методы ────────────────────────────────────────────

    def _ui(self, fn):
        """Безопасный вызов функции в главном потоке."""
        self.after(0, fn)

    def _set_controls(self, state):
        self._sort_btn.configure(state=state)
        self._browse_btn.configure(state=state)


# ── функция оценки резкости (модуль-уровень, не зависит от GUI) ──────────────

def _sharpness(filepath):
    """Возвращает дисперсию Лапласиана или None при ошибке."""
    ext = os.path.splitext(filepath)[1].lower()
    try:
        if ext == ".arw":
            if not _RAWPY_OK:
                return None
            with rawpy.imread(filepath) as raw:
                thumb = raw.extract_thumb()
            if thumb.format == rawpy.ThumbFormat.JPEG:
                img = Image.open(io.BytesIO(thumb.data))
            elif thumb.format == rawpy.ThumbFormat.BITMAP:
                img = Image.fromarray(thumb.data)
            else:
                return None
        else:
            img = Image.open(filepath)

        img  = img.convert("L").resize((640, 480), Image.LANCZOS)
        gray = np.array(img, dtype=np.uint8)
        return float(cv2.Laplacian(gray, cv2.CV_64F).var())
    except Exception:
        return None


# ── точка входа ───────────────────────────────────────────────────────────────

if __name__ == "__main__":
    app = StockSortApp()
    app.mainloop()