1194 lines
46 KiB
Python
1194 lines
46 KiB
Python
# 재학습 페이지 — 이미지 로드·라벨링 UI + 학습 제어
|
|
import json
|
|
import os
|
|
import shutil
|
|
from datetime import datetime
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from PyQt5.QtCore import Qt, QPoint, QRect, pyqtSignal
|
|
from PyQt5.QtGui import QPixmap, QPainter, QPen, QColor, QFont, QImage, QCursor
|
|
from PyQt5.QtWidgets import (
|
|
QWidget, QHBoxLayout, QVBoxLayout, QGroupBox,
|
|
QPushButton, QLabel, QListWidget, QListWidgetItem,
|
|
QSpinBox, QLineEdit, QProgressBar, QTextEdit,
|
|
QFormLayout, QFileDialog, QSizePolicy, QScrollArea, QMessageBox,
|
|
)
|
|
|
|
from ai.trainer import Trainer, TrainWorker
|
|
from paths import resolve_path, to_project_relative
|
|
from logger import log_train, log_action
|
|
|
|
_CLASS_COLORS = {
|
|
"스크래치": "#854F0B",
|
|
"이물": "#185FA5",
|
|
"흑점": "#3C3489",
|
|
"변형": "#A32D2D",
|
|
}
|
|
_CLASS_NAMES = list(_CLASS_COLORS.keys())
|
|
|
|
_GRP = (
|
|
"QGroupBox {"
|
|
" background:#222222; border:1px solid #333333; border-radius:6px;"
|
|
" margin-top:14px; padding:12px 10px 10px 10px;"
|
|
"}"
|
|
"QGroupBox::title {"
|
|
" color:#aaaaaa; subcontrol-origin:margin; left:10px; padding:0 4px;"
|
|
"}"
|
|
)
|
|
|
|
|
|
# ============================================================ #
|
|
# 라벨링 캔버스
|
|
# ============================================================ #
|
|
|
|
class LabelingCanvas(QWidget):
|
|
box_added = pyqtSignal(dict)
|
|
boxes_changed = pyqtSignal()
|
|
selection_changed = pyqtSignal(int)
|
|
zoom_changed = pyqtSignal(int) # 정수 퍼센트 (100 = 100%)
|
|
|
|
CLASS_MAP = {
|
|
"스크래치": 0, "이물": 1, "흑점": 2, "변형": 3,
|
|
}
|
|
CLASS_COLORS = {
|
|
"스크래치": "#FF4444", "이물": "#44AAFF",
|
|
"흑점": "#AA44FF", "변형": "#FF8800",
|
|
}
|
|
|
|
_HANDLE_SIZE = 8
|
|
# 핸들 인덱스: 0=TL, 1=T, 2=TR, 3=R, 4=BR, 5=B, 6=BL, 7=L
|
|
_HANDLE_CURSORS = [
|
|
Qt.SizeFDiagCursor, # 0 TL
|
|
Qt.SizeVerCursor, # 1 T
|
|
Qt.SizeBDiagCursor, # 2 TR
|
|
Qt.SizeHorCursor, # 3 R
|
|
Qt.SizeFDiagCursor, # 4 BR
|
|
Qt.SizeVerCursor, # 5 B
|
|
Qt.SizeBDiagCursor, # 6 BL
|
|
Qt.SizeHorCursor, # 7 L
|
|
]
|
|
|
|
def __init__(self, parent=None):
|
|
super().__init__(parent)
|
|
|
|
# 이미지 & 박스 상태
|
|
self.image: np.ndarray = None
|
|
self.boxes: list = []
|
|
self.history: list = [] # Ctrl+Z 용 (새 박스 추가 직전 스냅샷)
|
|
|
|
# 현재 그릴 클래스
|
|
self.current_class_id: int = 0
|
|
self.current_class_name: str = "스크래치"
|
|
|
|
# 선택 & 인터랙션
|
|
self.selected_index: int = -1
|
|
self.drag_mode: str = "none" # none/new_box/move/resize/pan
|
|
self.resize_handle: int = -1
|
|
self.current_rect: QRect = None
|
|
|
|
# 줌 / 패닝
|
|
self.scale: float = 1.0
|
|
self.offset_x: float = 0.0
|
|
self.offset_y: float = 0.0
|
|
self.space_pressed: bool = False
|
|
self._need_fit: bool = False
|
|
|
|
# 드래그 임시 상태
|
|
self._drag_start_img: QPoint = None
|
|
self._drag_start_screen: QPoint = None
|
|
self._move_orig_rect: QRect = None
|
|
self._resize_orig_rect: QRect = None
|
|
self._pan_start: QPoint = None
|
|
|
|
self.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
|
|
self.setStyleSheet("background:#111111;")
|
|
self.setFocusPolicy(Qt.StrongFocus)
|
|
self.setMouseTracking(True)
|
|
self.setCursor(Qt.CrossCursor)
|
|
|
|
# ── 좌표 변환 ──────────────────────────────────────────────────── #
|
|
|
|
def screen_to_image(self, pos: QPoint):
|
|
if self.image is None or self.scale == 0:
|
|
return None
|
|
img_h, img_w = self.image.shape[:2]
|
|
x = (pos.x() - self.offset_x) / self.scale
|
|
y = (pos.y() - self.offset_y) / self.scale
|
|
return QPoint(int(max(0.0, min(x, img_w - 1))),
|
|
int(max(0.0, min(y, img_h - 1))))
|
|
|
|
def image_to_screen_rect(self, rect: QRect) -> QRect:
|
|
return QRect(
|
|
int(rect.x() * self.scale + self.offset_x),
|
|
int(rect.y() * self.scale + self.offset_y),
|
|
int(rect.width() * self.scale),
|
|
int(rect.height() * self.scale),
|
|
)
|
|
|
|
# ── 핸들 / 히트테스트 ──────────────────────────────────────────── #
|
|
|
|
def _handle_rects(self, sr: QRect) -> list:
|
|
hs = self._HANDLE_SIZE
|
|
hh = hs // 2
|
|
cx = sr.center().x()
|
|
cy = sr.center().y()
|
|
l, t, r, b = sr.left(), sr.top(), sr.right(), sr.bottom()
|
|
return [
|
|
QRect(l - hh, t - hh, hs, hs),
|
|
QRect(cx - hh, t - hh, hs, hs),
|
|
QRect(r - hh, t - hh, hs, hs),
|
|
QRect(r - hh, cy - hh, hs, hs),
|
|
QRect(r - hh, b - hh, hs, hs),
|
|
QRect(cx - hh, b - hh, hs, hs),
|
|
QRect(l - hh, b - hh, hs, hs),
|
|
QRect(l - hh, cy - hh, hs, hs),
|
|
]
|
|
|
|
def get_handle_at(self, pos: QPoint, sr: QRect) -> int:
|
|
for i, hr in enumerate(self._handle_rects(sr)):
|
|
if hr.adjusted(-2, -2, 2, 2).contains(pos):
|
|
return i
|
|
return -1
|
|
|
|
def get_box_at(self, pos: QPoint) -> int:
|
|
result = -1
|
|
for i, box in enumerate(self.boxes):
|
|
if self.image_to_screen_rect(box["rect"]).contains(pos):
|
|
result = i
|
|
return result
|
|
|
|
# ── 히스토리 / 복사 ────────────────────────────────────────────── #
|
|
|
|
def _copy_boxes(self) -> list:
|
|
return [
|
|
{
|
|
"class_id": b["class_id"],
|
|
"class_name": b["class_name"],
|
|
"rect": QRect(b["rect"].x(), b["rect"].y(),
|
|
b["rect"].width(), b["rect"].height()),
|
|
}
|
|
for b in self.boxes
|
|
]
|
|
|
|
def _save_history(self):
|
|
self.history.append(self._copy_boxes())
|
|
if len(self.history) > 50:
|
|
self.history.pop(0)
|
|
|
|
# ── 커서 업데이트 ──────────────────────────────────────────────── #
|
|
|
|
def _update_cursor(self, pos: QPoint):
|
|
if self.space_pressed:
|
|
self.setCursor(Qt.OpenHandCursor)
|
|
return
|
|
if 0 <= self.selected_index < len(self.boxes):
|
|
sr = self.image_to_screen_rect(self.boxes[self.selected_index]["rect"])
|
|
h = self.get_handle_at(pos, sr)
|
|
if h >= 0:
|
|
self.setCursor(self._HANDLE_CURSORS[h])
|
|
return
|
|
if sr.contains(pos):
|
|
self.setCursor(Qt.SizeAllCursor)
|
|
return
|
|
if self.get_box_at(pos) >= 0:
|
|
self.setCursor(Qt.SizeAllCursor)
|
|
return
|
|
self.setCursor(Qt.CrossCursor)
|
|
|
|
# ── 퍼블릭 API ─────────────────────────────────────────────────── #
|
|
|
|
def set_image(self, img: np.ndarray):
|
|
self.image = img
|
|
self.boxes = []
|
|
self.history = []
|
|
self.selected_index = -1
|
|
self.drag_mode = "none"
|
|
self.current_rect = None
|
|
self._need_fit = True
|
|
self.fit_to_window()
|
|
self.selection_changed.emit(-1)
|
|
|
|
def set_class(self, class_id: int, class_name: str):
|
|
self.current_class_id = class_id
|
|
self.current_class_name = class_name
|
|
|
|
def change_selected_class(self, class_id: int, class_name: str):
|
|
if 0 <= self.selected_index < len(self.boxes):
|
|
self._save_history()
|
|
self.boxes[self.selected_index]["class_id"] = class_id
|
|
self.boxes[self.selected_index]["class_name"] = class_name
|
|
self.boxes_changed.emit()
|
|
self.update()
|
|
|
|
def delete_selected_box(self, index: int = -1):
|
|
if index < 0:
|
|
index = self.selected_index
|
|
if not (0 <= index < len(self.boxes)):
|
|
return
|
|
self.boxes.pop(index)
|
|
if self.selected_index >= len(self.boxes):
|
|
self.selected_index = -1
|
|
elif self.selected_index == index:
|
|
self.selected_index = -1
|
|
self.selection_changed.emit(self.selected_index)
|
|
self.boxes_changed.emit()
|
|
self.update()
|
|
|
|
def clear_boxes(self):
|
|
self.boxes = []
|
|
self.selected_index = -1
|
|
self.selection_changed.emit(-1)
|
|
self.boxes_changed.emit()
|
|
self.update()
|
|
|
|
def fit_to_window(self):
|
|
if self.image is None or self.width() == 0 or self.height() == 0:
|
|
return
|
|
img_h, img_w = self.image.shape[:2]
|
|
self.scale = min(self.width() / img_w, self.height() / img_h)
|
|
self.offset_x = (self.width() - img_w * self.scale) / 2
|
|
self.offset_y = (self.height() - img_h * self.scale) / 2
|
|
self._need_fit = False
|
|
self.zoom_changed.emit(int(self.scale * 100))
|
|
self.update()
|
|
|
|
def get_yolo_labels(self) -> list:
|
|
if self.image is None:
|
|
return []
|
|
img_h, img_w = self.image.shape[:2]
|
|
result = []
|
|
for box in self.boxes:
|
|
r = box["rect"]
|
|
cx = (r.x() + r.width() / 2) / img_w
|
|
cy = (r.y() + r.height() / 2) / img_h
|
|
nw = r.width() / img_w
|
|
nh = r.height() / img_h
|
|
result.append(
|
|
f"{box['class_id']} {cx:.6f} {cy:.6f} {nw:.6f} {nh:.6f}"
|
|
)
|
|
return result
|
|
|
|
# ── 마우스 이벤트 ──────────────────────────────────────────────── #
|
|
|
|
def mousePressEvent(self, e):
|
|
if self.image is None:
|
|
return
|
|
pos = e.pos()
|
|
self.setFocus()
|
|
|
|
# 패닝: 스페이스+좌클릭 or 중간 버튼
|
|
if (e.button() == Qt.MiddleButton or
|
|
(e.button() == Qt.LeftButton and self.space_pressed)):
|
|
self.drag_mode = "pan"
|
|
self._pan_start = pos
|
|
self.setCursor(Qt.ClosedHandCursor)
|
|
return
|
|
|
|
if e.button() != Qt.LeftButton:
|
|
return
|
|
|
|
# 선택된 박스의 핸들 먼저 확인
|
|
if 0 <= self.selected_index < len(self.boxes):
|
|
sr = self.image_to_screen_rect(self.boxes[self.selected_index]["rect"])
|
|
h = self.get_handle_at(pos, sr)
|
|
if h >= 0:
|
|
self.drag_mode = "resize"
|
|
self.resize_handle = h
|
|
self._drag_start_screen = pos
|
|
self._resize_orig_rect = QRect(self.boxes[self.selected_index]["rect"])
|
|
return
|
|
if sr.contains(pos):
|
|
self.drag_mode = "move"
|
|
self._drag_start_screen = pos
|
|
self._move_orig_rect = QRect(self.boxes[self.selected_index]["rect"])
|
|
return
|
|
|
|
# 다른 박스 히트테스트
|
|
hit = self.get_box_at(pos)
|
|
if hit >= 0:
|
|
prev = self.selected_index
|
|
self.selected_index = hit
|
|
if hit != prev:
|
|
self.selection_changed.emit(hit)
|
|
sr = self.image_to_screen_rect(self.boxes[hit]["rect"])
|
|
h = self.get_handle_at(pos, sr)
|
|
if h >= 0:
|
|
self.drag_mode = "resize"
|
|
self.resize_handle = h
|
|
self._drag_start_screen = pos
|
|
self._resize_orig_rect = QRect(self.boxes[hit]["rect"])
|
|
else:
|
|
self.drag_mode = "move"
|
|
self._drag_start_screen = pos
|
|
self._move_orig_rect = QRect(self.boxes[hit]["rect"])
|
|
self.update()
|
|
return
|
|
|
|
# 빈 공간 클릭 → 선택 해제 + 새 박스 그리기 시작
|
|
if self.selected_index >= 0:
|
|
self.selected_index = -1
|
|
self.selection_changed.emit(-1)
|
|
img_pt = self.screen_to_image(pos)
|
|
if img_pt is not None:
|
|
self._drag_start_img = img_pt
|
|
self.drag_mode = "new_box"
|
|
self.current_rect = None
|
|
self.update()
|
|
|
|
def mouseMoveEvent(self, e):
|
|
pos = e.pos()
|
|
|
|
if self.drag_mode == "pan":
|
|
self.offset_x += pos.x() - self._pan_start.x()
|
|
self.offset_y += pos.y() - self._pan_start.y()
|
|
self._pan_start = pos
|
|
self.update()
|
|
return
|
|
|
|
if self.drag_mode == "new_box":
|
|
ip = self.screen_to_image(pos)
|
|
if ip is not None and self._drag_start_img is not None:
|
|
self.current_rect = QRect(self._drag_start_img, ip).normalized()
|
|
self.update()
|
|
return
|
|
|
|
if self.drag_mode == "move":
|
|
self._do_move(pos)
|
|
self.update()
|
|
return
|
|
|
|
if self.drag_mode == "resize":
|
|
self._do_resize(pos)
|
|
self.update()
|
|
return
|
|
|
|
self._update_cursor(pos)
|
|
|
|
def mouseReleaseEvent(self, e):
|
|
pos = e.pos()
|
|
|
|
if self.drag_mode == "pan" and e.button() in (Qt.LeftButton, Qt.MiddleButton):
|
|
self.drag_mode = "none"
|
|
self._update_cursor(pos)
|
|
return
|
|
|
|
if e.button() != Qt.LeftButton:
|
|
return
|
|
|
|
if self.drag_mode == "new_box":
|
|
if (self.current_rect is not None
|
|
and self.current_rect.width() >= 10
|
|
and self.current_rect.height() >= 10):
|
|
self._save_history()
|
|
box = {
|
|
"class_id": self.current_class_id,
|
|
"class_name": self.current_class_name,
|
|
"rect": self.current_rect,
|
|
}
|
|
self.boxes.append(box)
|
|
self.selected_index = len(self.boxes) - 1
|
|
self.selection_changed.emit(self.selected_index)
|
|
self.box_added.emit(box)
|
|
self.current_rect = None
|
|
self._drag_start_img = None
|
|
|
|
elif self.drag_mode in ("move", "resize"):
|
|
self.boxes_changed.emit()
|
|
|
|
self.drag_mode = "none"
|
|
self._update_cursor(pos)
|
|
self.update()
|
|
|
|
def mouseDoubleClickEvent(self, e):
|
|
if e.button() == Qt.LeftButton:
|
|
self.fit_to_window()
|
|
|
|
def wheelEvent(self, e):
|
|
if self.image is None:
|
|
return
|
|
pos = e.pos()
|
|
delta = e.angleDelta().y()
|
|
factor = 1.15 if delta > 0 else 1.0 / 1.15
|
|
new_scale = max(0.5, min(self.scale * factor, 5.0))
|
|
if new_scale == self.scale:
|
|
return
|
|
img_x = (pos.x() - self.offset_x) / self.scale
|
|
img_y = (pos.y() - self.offset_y) / self.scale
|
|
self.scale = new_scale
|
|
self.offset_x = pos.x() - img_x * self.scale
|
|
self.offset_y = pos.y() - img_y * self.scale
|
|
self.zoom_changed.emit(int(self.scale * 100))
|
|
self.update()
|
|
|
|
# ── 키보드 이벤트 ──────────────────────────────────────────────── #
|
|
|
|
def keyPressEvent(self, e):
|
|
key = e.key()
|
|
|
|
if key == Qt.Key_Space and not e.isAutoRepeat():
|
|
self.space_pressed = True
|
|
self.setCursor(Qt.OpenHandCursor)
|
|
return
|
|
|
|
if key in (Qt.Key_Delete, Qt.Key_Backspace):
|
|
if self.selected_index >= 0:
|
|
self.delete_selected_box(self.selected_index)
|
|
return
|
|
|
|
if key == Qt.Key_Z and (e.modifiers() & Qt.ControlModifier):
|
|
if self.history:
|
|
self.boxes = self.history.pop()
|
|
self.selected_index = min(self.selected_index, len(self.boxes) - 1)
|
|
if self.selected_index < 0:
|
|
self.selected_index = -1
|
|
self.selection_changed.emit(self.selected_index)
|
|
self.boxes_changed.emit()
|
|
self.update()
|
|
return
|
|
|
|
if key == Qt.Key_Escape:
|
|
self.selected_index = -1
|
|
self.drag_mode = "none"
|
|
self.current_rect = None
|
|
self._drag_start_img = None
|
|
self.selection_changed.emit(-1)
|
|
self._update_cursor(self.mapFromGlobal(QCursor.pos()))
|
|
self.update()
|
|
return
|
|
|
|
super().keyPressEvent(e)
|
|
|
|
def keyReleaseEvent(self, e):
|
|
if e.key() == Qt.Key_Space and not e.isAutoRepeat():
|
|
self.space_pressed = False
|
|
if self.drag_mode == "pan":
|
|
self.drag_mode = "none"
|
|
self._update_cursor(self.mapFromGlobal(QCursor.pos()))
|
|
super().keyReleaseEvent(e)
|
|
|
|
# ── paintEvent ─────────────────────────────────────────────────── #
|
|
|
|
def paintEvent(self, _e):
|
|
painter = QPainter(self)
|
|
painter.fillRect(self.rect(), QColor("#111111"))
|
|
|
|
if self.image is None:
|
|
painter.setPen(QColor("#555555"))
|
|
painter.setFont(QFont("Arial", 14))
|
|
painter.drawText(self.rect(), Qt.AlignCenter, "이미지를 선택하세요")
|
|
painter.end()
|
|
return
|
|
|
|
# numpy BGR → QPixmap
|
|
rgb = np.ascontiguousarray(cv2.cvtColor(self.image, cv2.COLOR_BGR2RGB))
|
|
ih, iw, ch = rgb.shape
|
|
qimg = QImage(rgb.data, iw, ih, ch * iw, QImage.Format_RGB888)
|
|
pix = QPixmap.fromImage(qimg)
|
|
sw = max(1, int(iw * self.scale))
|
|
sh = max(1, int(ih * self.scale))
|
|
scaled = pix.scaled(sw, sh, Qt.IgnoreAspectRatio, Qt.SmoothTransformation)
|
|
painter.drawPixmap(int(self.offset_x), int(self.offset_y), scaled)
|
|
|
|
# 박스 그리기
|
|
painter.setFont(QFont("Arial", 10, QFont.Bold))
|
|
for i, box in enumerate(self.boxes):
|
|
color = QColor(self.CLASS_COLORS.get(box["class_name"], "#ffffff"))
|
|
sr = self.image_to_screen_rect(box["rect"])
|
|
is_sel = (i == self.selected_index)
|
|
|
|
if is_sel:
|
|
fill = QColor(color); fill.setAlpha(40)
|
|
painter.fillRect(sr, fill)
|
|
painter.setPen(QPen(QColor("#ffffff"), 3))
|
|
painter.drawRect(sr)
|
|
# 8개 핸들
|
|
painter.setPen(QPen(QColor("#555555"), 1))
|
|
for hr in self._handle_rects(sr):
|
|
painter.fillRect(hr, QColor("#ffffff"))
|
|
painter.drawRect(hr)
|
|
lbl_bg = QColor("#ffffff")
|
|
lbl_txt = QColor("#222222")
|
|
else:
|
|
painter.setPen(QPen(color, 2))
|
|
painter.drawRect(sr)
|
|
lbl_bg = color
|
|
lbl_txt = QColor("#ffffff")
|
|
|
|
# 클래스명 레이블
|
|
text = box["class_name"]
|
|
lbl_w = max(len(text) * 11, 60)
|
|
lbl_y = sr.y() - 18 if sr.y() >= 18 else sr.y()
|
|
lbl_r = QRect(sr.x(), lbl_y, lbl_w, 18)
|
|
painter.fillRect(lbl_r, lbl_bg)
|
|
painter.setPen(lbl_txt)
|
|
painter.drawText(lbl_r.x() + 3, lbl_r.y() + 13, text)
|
|
|
|
# 드래그 중인 새 박스 (점선 흰색)
|
|
if self.drag_mode == "new_box" and self.current_rect is not None:
|
|
sr = self.image_to_screen_rect(self.current_rect)
|
|
painter.setPen(QPen(QColor("#ffffff"), 2, Qt.DashLine))
|
|
painter.setBrush(Qt.NoBrush)
|
|
painter.drawRect(sr)
|
|
|
|
painter.end()
|
|
|
|
def resizeEvent(self, e):
|
|
super().resizeEvent(e)
|
|
if self._need_fit and self.width() > 0 and self.height() > 0:
|
|
self.fit_to_window()
|
|
else:
|
|
self.update()
|
|
|
|
# ── 이동 / 리사이즈 헬퍼 ──────────────────────────────────────── #
|
|
|
|
def _do_move(self, screen_pos: QPoint):
|
|
if self.image is None or self.selected_index < 0:
|
|
return
|
|
img_h, img_w = self.image.shape[:2]
|
|
dx = (screen_pos.x() - self._drag_start_screen.x()) / self.scale
|
|
dy = (screen_pos.y() - self._drag_start_screen.y()) / self.scale
|
|
orig = self._move_orig_rect
|
|
nx = max(0, min(int(orig.x() + dx), img_w - orig.width()))
|
|
ny = max(0, min(int(orig.y() + dy), img_h - orig.height()))
|
|
self.boxes[self.selected_index]["rect"] = QRect(nx, ny, orig.width(), orig.height())
|
|
|
|
def _do_resize(self, screen_pos: QPoint):
|
|
if self.image is None or self.selected_index < 0:
|
|
return
|
|
img_h, img_w = self.image.shape[:2]
|
|
dx = (screen_pos.x() - self._drag_start_screen.x()) / self.scale
|
|
dy = (screen_pos.y() - self._drag_start_screen.y()) / self.scale
|
|
orig = self._resize_orig_rect
|
|
x1, y1 = orig.x(), orig.y()
|
|
x2, y2 = orig.x() + orig.width(), orig.y() + orig.height()
|
|
h = self.resize_handle
|
|
|
|
if h in (0, 6, 7): x1 = int(orig.x() + dx)
|
|
if h in (2, 3, 4): x2 = int(orig.x() + orig.width() + dx)
|
|
if h in (0, 1, 2): y1 = int(orig.y() + dy)
|
|
if h in (4, 5, 6): y2 = int(orig.y() + orig.height() + dy)
|
|
|
|
if x2 - x1 < 10:
|
|
if h in (0, 6, 7): x1 = x2 - 10
|
|
else: x2 = x1 + 10
|
|
if y2 - y1 < 10:
|
|
if h in (0, 1, 2): y1 = y2 - 10
|
|
else: y2 = y1 + 10
|
|
|
|
x1 = max(0, x1); y1 = max(0, y1)
|
|
x2 = min(img_w, x2); y2 = min(img_h, y2)
|
|
self.boxes[self.selected_index]["rect"] = QRect(x1, y1, x2 - x1, y2 - y1)
|
|
|
|
|
|
# ============================================================ #
|
|
# 재학습 페이지
|
|
# ============================================================ #
|
|
|
|
class RetrainPage(QWidget):
|
|
def __init__(self, parent=None):
|
|
super().__init__(parent)
|
|
self._img_dir = ""
|
|
self._img_files = []
|
|
self._cur_path = ""
|
|
self._trainer = Trainer()
|
|
self._worker = None
|
|
|
|
self._build_ui()
|
|
|
|
# ================================================================ #
|
|
# 레이아웃
|
|
# ================================================================ #
|
|
|
|
def _build_ui(self):
|
|
root = QHBoxLayout(self)
|
|
root.setContentsMargins(0, 0, 0, 0)
|
|
root.setSpacing(0)
|
|
root.addWidget(self._build_left(), stretch=2)
|
|
root.addWidget(self._build_right(), stretch=3)
|
|
|
|
def _build_left(self) -> QScrollArea:
|
|
scroll = QScrollArea()
|
|
scroll.setWidgetResizable(True)
|
|
scroll.setFrameShape(QScrollArea.NoFrame)
|
|
scroll.setStyleSheet("background:#1a1a1a;")
|
|
inner = QWidget()
|
|
inner.setStyleSheet("background:#1a1a1a;")
|
|
lay = QVBoxLayout(inner)
|
|
lay.setContentsMargins(12, 12, 8, 12)
|
|
lay.setSpacing(0)
|
|
lay.addWidget(self._build_img_load_section())
|
|
lay.addWidget(self._build_class_section())
|
|
lay.addWidget(self._build_dataset_section())
|
|
lay.addStretch()
|
|
scroll.setWidget(inner)
|
|
return scroll
|
|
|
|
# ── 섹션 1: 이미지 로드 ────────────────────────────────────────── #
|
|
|
|
def _build_img_load_section(self) -> QGroupBox:
|
|
g = QGroupBox("이미지 로드")
|
|
g.setStyleSheet(_GRP)
|
|
lay = QVBoxLayout(g)
|
|
lay.setSpacing(8)
|
|
|
|
btn = QPushButton("이미지 폴더 선택")
|
|
btn.setFixedHeight(56)
|
|
btn.setStyleSheet(_btn_style("#1a3a5c"))
|
|
btn.clicked.connect(self._on_select_folder)
|
|
|
|
self._folder_lbl = QLabel("폴더를 선택하세요")
|
|
self._folder_lbl.setStyleSheet("color:#777777; font-size:12px;")
|
|
self._folder_lbl.setWordWrap(True)
|
|
|
|
self._img_list = QListWidget()
|
|
self._img_list.setMinimumHeight(160)
|
|
self._img_list.setStyleSheet("""
|
|
QListWidget {
|
|
background:#1a1a1a; border:1px solid #333333;
|
|
border-radius:4px; font-size:13px; color:#cccccc;
|
|
}
|
|
QListWidget::item { padding:4px 8px; border-bottom:1px solid #2a2a2a; }
|
|
QListWidget::item:selected { background:#185FA5; color:#ffffff; }
|
|
""")
|
|
self._img_list.currentRowChanged.connect(self._on_img_selected)
|
|
|
|
lay.addWidget(btn)
|
|
lay.addWidget(self._folder_lbl)
|
|
lay.addWidget(self._img_list, stretch=1)
|
|
return g
|
|
|
|
# ── 섹션 2: 불량 클래스 선택 ───────────────────────────────────── #
|
|
|
|
def _build_class_section(self) -> QGroupBox:
|
|
g = QGroupBox("불량 클래스 선택")
|
|
g.setStyleSheet(_GRP)
|
|
lay = QVBoxLayout(g)
|
|
lay.setSpacing(6)
|
|
|
|
self._class_btns: dict[str, QPushButton] = {}
|
|
for cls_name, color in _CLASS_COLORS.items():
|
|
btn = QPushButton(cls_name)
|
|
btn.setFixedHeight(56)
|
|
btn.setCheckable(True)
|
|
btn.setStyleSheet(_cls_btn_style(color, checked=False))
|
|
btn.clicked.connect(lambda _, c=cls_name: self._on_class_select(c))
|
|
self._class_btns[cls_name] = btn
|
|
lay.addWidget(btn)
|
|
|
|
first = _CLASS_NAMES[0]
|
|
self._class_btns[first].setChecked(True)
|
|
self._class_btns[first].setStyleSheet(_cls_btn_style(_CLASS_COLORS[first], checked=True))
|
|
self._active_cls = first
|
|
return g
|
|
|
|
# ── 섹션 3: 데이터셋 설정 ──────────────────────────────────────── #
|
|
|
|
def _build_dataset_section(self) -> QGroupBox:
|
|
g = QGroupBox("데이터셋 설정")
|
|
g.setStyleSheet(_GRP)
|
|
lay = QVBoxLayout(g)
|
|
lay.setSpacing(8)
|
|
|
|
form = QFormLayout()
|
|
form.setHorizontalSpacing(12)
|
|
form.setVerticalSpacing(8)
|
|
|
|
self._epoch_spin = QSpinBox()
|
|
self._epoch_spin.setRange(1, 300)
|
|
self._epoch_spin.setValue(100)
|
|
self._epoch_spin.setStyleSheet(_spinbox_style())
|
|
|
|
self._batch_spin = QSpinBox()
|
|
self._batch_spin.setRange(1, 64)
|
|
self._batch_spin.setValue(16)
|
|
self._batch_spin.setStyleSheet(_spinbox_style())
|
|
|
|
form.addRow("Epoch", self._epoch_spin)
|
|
form.addRow("Batch", self._batch_spin)
|
|
lay.addLayout(form)
|
|
|
|
path_lbl = QLabel("모델 저장 경로")
|
|
path_lbl.setStyleSheet("color:#888888; font-size:13px;")
|
|
path_row = QHBoxLayout()
|
|
self._save_path_edit = QLineEdit("ai/models/best.pt")
|
|
self._save_path_edit.setFixedHeight(44)
|
|
self._save_path_edit.setStyleSheet(
|
|
"background:#2a2a2a; color:#ffffff; border:1px solid #555555;"
|
|
"border-radius:4px; padding:0 8px; font-size:13px;"
|
|
)
|
|
btn_browse = QPushButton("찾기")
|
|
btn_browse.setFixedSize(64, 44)
|
|
btn_browse.setStyleSheet(_btn_style("#333333", font_size=13))
|
|
btn_browse.clicked.connect(self._on_browse_save)
|
|
path_row.addWidget(self._save_path_edit, stretch=1)
|
|
path_row.addWidget(btn_browse)
|
|
|
|
lay.addWidget(path_lbl)
|
|
lay.addLayout(path_row)
|
|
return g
|
|
|
|
# ── 우측 패널 ──────────────────────────────────────────────────── #
|
|
|
|
def _build_right(self) -> QWidget:
|
|
w = QWidget()
|
|
w.setStyleSheet("background:#1a1a1a;")
|
|
lay = QVBoxLayout(w)
|
|
lay.setContentsMargins(8, 12, 12, 12)
|
|
lay.setSpacing(0)
|
|
lay.addWidget(self._build_labeling_section(), stretch=55)
|
|
lay.addWidget(self._build_train_section(), stretch=25)
|
|
lay.addWidget(self._build_save_section(), stretch=20)
|
|
return w
|
|
|
|
def _build_labeling_section(self) -> QGroupBox:
|
|
g = QGroupBox("이미지 표시 / 라벨링")
|
|
g.setStyleSheet(_GRP)
|
|
outer = QVBoxLayout(g)
|
|
outer.setSpacing(6)
|
|
|
|
# ── 툴바 (줌 표시 + 초기화 버튼) ──
|
|
toolbar = QHBoxLayout()
|
|
self._zoom_lbl = QLabel("100%")
|
|
self._zoom_lbl.setStyleSheet(
|
|
"color:#aaaaaa; font-size:12px; min-width:55px;"
|
|
)
|
|
hint = QLabel("더블클릭: fit | 휠: 줌 | Space+드래그: 패닝 | Del: 박스삭제 | Ctrl+Z: 실행취소")
|
|
hint.setStyleSheet("color:#555555; font-size:11px;")
|
|
btn_fit = QPushButton("초기화")
|
|
btn_fit.setFixedHeight(28)
|
|
btn_fit.setStyleSheet(_btn_style("#333333", font_size=12))
|
|
btn_fit.clicked.connect(lambda: self._canvas.fit_to_window())
|
|
toolbar.addWidget(self._zoom_lbl)
|
|
toolbar.addWidget(hint)
|
|
toolbar.addStretch()
|
|
toolbar.addWidget(btn_fit)
|
|
outer.addLayout(toolbar)
|
|
|
|
# ── 캔버스 + 박스 목록 ──
|
|
main_row = QHBoxLayout()
|
|
|
|
self._canvas = LabelingCanvas()
|
|
self._canvas.box_added.connect(lambda _: self._refresh_box_list())
|
|
self._canvas.boxes_changed.connect(self._refresh_box_list)
|
|
self._canvas.selection_changed.connect(self._on_canvas_selection_changed)
|
|
self._canvas.zoom_changed.connect(lambda pct: self._zoom_lbl.setText(f"{pct}%"))
|
|
main_row.addWidget(self._canvas, stretch=3)
|
|
|
|
side = QVBoxLayout()
|
|
side.setSpacing(6)
|
|
box_lbl = QLabel("박스 목록")
|
|
box_lbl.setStyleSheet("color:#888888; font-size:13px;")
|
|
self._box_list = QListWidget()
|
|
self._box_list.setStyleSheet("""
|
|
QListWidget {
|
|
background:#1a1a1a; border:1px solid #333333;
|
|
border-radius:4px; font-size:12px; color:#cccccc;
|
|
}
|
|
QListWidget::item { padding:3px 6px; }
|
|
QListWidget::item:selected { background:#3C3489; }
|
|
""")
|
|
self._box_list.currentRowChanged.connect(self._on_box_list_select)
|
|
|
|
btn_del = QPushButton("박스 삭제")
|
|
btn_del.setFixedHeight(44)
|
|
btn_del.setStyleSheet(_btn_style("#5c1a1a", font_size=13))
|
|
btn_del.clicked.connect(self._on_del_box)
|
|
side.addWidget(box_lbl)
|
|
side.addWidget(self._box_list, stretch=1)
|
|
side.addWidget(btn_del)
|
|
main_row.addLayout(side, stretch=1)
|
|
|
|
outer.addLayout(main_row, stretch=1)
|
|
|
|
btn_save = QPushButton("라벨 저장 (YOLO .txt)")
|
|
btn_save.setFixedHeight(50)
|
|
btn_save.setStyleSheet(_btn_style("#2e5c2e", font_size=14))
|
|
btn_save.clicked.connect(self._on_label_save)
|
|
outer.addWidget(btn_save)
|
|
return g
|
|
|
|
def _build_train_section(self) -> QGroupBox:
|
|
g = QGroupBox("학습 제어")
|
|
g.setStyleSheet(_GRP)
|
|
lay = QVBoxLayout(g)
|
|
lay.setSpacing(6)
|
|
|
|
btn_row = QHBoxLayout()
|
|
self._start_btn = QPushButton("학습 시작")
|
|
self._start_btn.setFixedHeight(70)
|
|
self._start_btn.setStyleSheet(_btn_style("#1D9E75", font_size=17, bold=True))
|
|
self._start_btn.clicked.connect(self._on_train_start)
|
|
|
|
self._stop_btn = QPushButton("학습 중지")
|
|
self._stop_btn.setFixedHeight(70)
|
|
self._stop_btn.setEnabled(False)
|
|
self._stop_btn.setStyleSheet(_btn_style("#A32D2D", font_size=17, bold=True))
|
|
self._stop_btn.clicked.connect(self._on_train_stop)
|
|
btn_row.addWidget(self._start_btn)
|
|
btn_row.addWidget(self._stop_btn)
|
|
|
|
self._progress_bar = QProgressBar()
|
|
self._progress_bar.setRange(0, 100)
|
|
self._progress_bar.setValue(0)
|
|
self._progress_bar.setFixedHeight(22)
|
|
self._progress_bar.setStyleSheet("""
|
|
QProgressBar {
|
|
background:#2a2a2a; border:1px solid #555555;
|
|
border-radius:4px; text-align:center; color:#ffffff; font-size:13px;
|
|
}
|
|
QProgressBar::chunk { background:#1D9E75; border-radius:4px; }
|
|
""")
|
|
|
|
self._status_lbl = QLabel("대기 중")
|
|
self._status_lbl.setStyleSheet("color:#888888; font-size:14px;")
|
|
self._status_lbl.setAlignment(Qt.AlignCenter)
|
|
|
|
self._log_box = QTextEdit()
|
|
self._log_box.setReadOnly(True)
|
|
self._log_box.setFixedHeight(80)
|
|
self._log_box.setStyleSheet(
|
|
"background:#111111; color:#aaaaaa; border:1px solid #333333;"
|
|
"border-radius:4px; font-size:12px; font-family:Consolas,monospace;"
|
|
)
|
|
|
|
lay.addLayout(btn_row)
|
|
lay.addWidget(self._progress_bar)
|
|
lay.addWidget(self._status_lbl)
|
|
lay.addWidget(self._log_box)
|
|
return g
|
|
|
|
def _build_save_section(self) -> QGroupBox:
|
|
g = QGroupBox("모델 저장")
|
|
g.setStyleSheet(_GRP)
|
|
lay = QVBoxLayout(g)
|
|
lay.setSpacing(8)
|
|
|
|
self._save_btn = QPushButton("모델 저장")
|
|
self._save_btn.setFixedHeight(56)
|
|
self._save_btn.setEnabled(False)
|
|
self._save_btn.setStyleSheet(_btn_style("#1a4d1a", font_size=15, bold=True))
|
|
self._save_btn.clicked.connect(self._on_model_save)
|
|
|
|
self._save_result_lbl = QLabel("")
|
|
self._save_result_lbl.setAlignment(Qt.AlignCenter)
|
|
self._save_result_lbl.setStyleSheet("color:#aaaaaa; font-size:13px;")
|
|
|
|
lay.addWidget(self._save_btn)
|
|
lay.addWidget(self._save_result_lbl)
|
|
return g
|
|
|
|
# ================================================================ #
|
|
# 슬롯
|
|
# ================================================================ #
|
|
|
|
def _on_select_folder(self):
|
|
folder = QFileDialog.getExistingDirectory(self, "이미지 폴더 선택", "")
|
|
if not folder:
|
|
return
|
|
self._img_dir = folder
|
|
self._folder_lbl.setText(folder)
|
|
exts = {".jpg", ".jpeg", ".png", ".bmp"}
|
|
self._img_files = sorted(
|
|
f for f in os.listdir(folder)
|
|
if os.path.splitext(f)[1].lower() in exts
|
|
)
|
|
self._img_list.clear()
|
|
for fname in self._img_files:
|
|
item = QListWidgetItem(fname)
|
|
item.setSizeHint(item.sizeHint().__class__(0, 44))
|
|
self._img_list.addItem(item)
|
|
print(f"[재학습] 폴더: {folder} ({len(self._img_files)}개)")
|
|
|
|
def _on_img_selected(self, row: int):
|
|
if row < 0 or row >= len(self._img_files):
|
|
return
|
|
path = os.path.join(self._img_dir, self._img_files[row])
|
|
self._cur_path = path
|
|
|
|
img = cv2.imdecode(np.fromfile(path, dtype=np.uint8), cv2.IMREAD_COLOR)
|
|
if img is None:
|
|
print(f"[재학습] 이미지 로드 실패: {path}")
|
|
return
|
|
|
|
self._canvas.set_image(img)
|
|
class_id = LabelingCanvas.CLASS_MAP.get(self._active_cls, 0)
|
|
self._canvas.set_class(class_id, self._active_cls)
|
|
|
|
txt_path = os.path.splitext(path)[0] + ".txt"
|
|
if os.path.exists(txt_path):
|
|
self._load_yolo_labels(txt_path, img.shape)
|
|
|
|
self._refresh_box_list()
|
|
|
|
def _load_yolo_labels(self, txt_path: str, img_shape):
|
|
img_h, img_w = img_shape[:2]
|
|
id_to_name = {v: k for k, v in LabelingCanvas.CLASS_MAP.items()}
|
|
try:
|
|
with open(txt_path, "r", encoding="utf-8") as f:
|
|
for line in f:
|
|
parts = line.strip().split()
|
|
if len(parts) != 5:
|
|
continue
|
|
cid = int(parts[0])
|
|
cx, cy, nw, nh = map(float, parts[1:])
|
|
x = int((cx - nw / 2) * img_w)
|
|
y = int((cy - nh / 2) * img_h)
|
|
w = int(nw * img_w)
|
|
h = int(nh * img_h)
|
|
self._canvas.boxes.append({
|
|
"class_id": cid,
|
|
"class_name": id_to_name.get(cid, "스크래치"),
|
|
"rect": QRect(x, y, w, h),
|
|
})
|
|
self._canvas.update()
|
|
except Exception as err:
|
|
print(f"[재학습] 라벨 로드 실패: {err}")
|
|
|
|
def _on_class_select(self, cls_name: str):
|
|
for name, btn in self._class_btns.items():
|
|
is_active = (name == cls_name)
|
|
btn.setChecked(is_active)
|
|
btn.setStyleSheet(_cls_btn_style(_CLASS_COLORS[name], checked=is_active))
|
|
self._active_cls = cls_name
|
|
class_id = LabelingCanvas.CLASS_MAP.get(cls_name, 0)
|
|
self._canvas.set_class(class_id, cls_name)
|
|
# 선택된 박스가 있으면 클래스 변경
|
|
if self._canvas.selected_index >= 0:
|
|
self._canvas.change_selected_class(class_id, cls_name)
|
|
|
|
# ── 캔버스 ↔ 박스 목록 동기화 ─────────────────────────────────── #
|
|
|
|
def _on_canvas_selection_changed(self, index: int):
|
|
"""캔버스 선택 변경 → 리스트 하이라이트 (루프 방지)."""
|
|
self._box_list.blockSignals(True)
|
|
self._box_list.setCurrentRow(index)
|
|
self._box_list.blockSignals(False)
|
|
|
|
def _on_box_list_select(self, row: int):
|
|
"""리스트 클릭 → 캔버스 선택 변경."""
|
|
if row != self._canvas.selected_index:
|
|
self._canvas.selected_index = row
|
|
self._canvas.update()
|
|
|
|
def _on_del_box(self):
|
|
idx = self._canvas.selected_index
|
|
if idx < 0:
|
|
idx = self._box_list.currentRow()
|
|
if idx >= 0:
|
|
self._canvas.delete_selected_box(idx)
|
|
# _refresh_box_list는 boxes_changed 시그널로 자동 호출
|
|
|
|
def _on_label_save(self):
|
|
if not self._cur_path:
|
|
QMessageBox.warning(self, "경고", "이미지를 먼저 선택하세요.")
|
|
return
|
|
labels = self._canvas.get_yolo_labels()
|
|
if not labels:
|
|
QMessageBox.warning(self, "경고", "박스를 먼저 그려주세요.")
|
|
return
|
|
txt_path = os.path.splitext(self._cur_path)[0] + ".txt"
|
|
try:
|
|
with open(txt_path, "w", encoding="utf-8") as f:
|
|
f.write("\n".join(labels))
|
|
QMessageBox.information(self, "저장 완료", f"라벨 저장 완료:\n{txt_path}")
|
|
self._log_append(f"라벨 저장: {os.path.basename(txt_path)}")
|
|
except Exception as err:
|
|
QMessageBox.critical(self, "저장 실패", str(err))
|
|
|
|
def _on_browse_save(self):
|
|
path, _ = QFileDialog.getSaveFileName(
|
|
self, "모델 저장 경로", "ai/models/best.pt", "PyTorch 모델 (*.pt)"
|
|
)
|
|
if path:
|
|
self._save_path_edit.setText(path)
|
|
|
|
def _on_train_start(self):
|
|
if not self._img_dir:
|
|
QMessageBox.warning(self, "경고", "이미지 폴더를 먼저 선택하세요.")
|
|
return
|
|
|
|
label_files = [
|
|
f for f in os.listdir(self._img_dir) if f.lower().endswith(".txt")
|
|
]
|
|
if not label_files:
|
|
QMessageBox.warning(
|
|
self, "경고", "라벨 파일이 없습니다. 먼저 라벨링해주세요."
|
|
)
|
|
return
|
|
|
|
self._progress_bar.setValue(0)
|
|
self._progress_bar.setFormat("0%")
|
|
self._log_box.clear()
|
|
self._status_lbl.setText("학습 중...")
|
|
self._status_lbl.setStyleSheet("color:#1D9E75; font-size:14px; font-weight:bold;")
|
|
self._start_btn.setEnabled(False)
|
|
self._stop_btn.setEnabled(True)
|
|
self._save_btn.setEnabled(False)
|
|
|
|
save_path = self._save_path_edit.text().strip() or "ai/models/best.pt"
|
|
epochs = self._epoch_spin.value()
|
|
batch = self._batch_spin.value()
|
|
log_action(
|
|
f"[재학습] 학습 시작 | epochs={epochs} | batch={batch} | 저장={save_path}"
|
|
)
|
|
log_train(
|
|
f"학습 시작 | 데이터셋={self._img_dir} | epochs={epochs} | "
|
|
f"batch={batch} | 저장={save_path} | 라벨파일={len(label_files)}개"
|
|
)
|
|
self._worker = TrainWorker(
|
|
self._trainer,
|
|
self._img_dir,
|
|
epochs,
|
|
batch,
|
|
save_path,
|
|
)
|
|
self._worker.log_signal.connect(self._on_log)
|
|
self._worker.progress_signal.connect(self._on_progress)
|
|
self._worker.finished_signal.connect(self._on_finished)
|
|
self._worker.start()
|
|
|
|
def _on_train_stop(self):
|
|
if self._worker and self._worker.isRunning():
|
|
self._worker.stop_subprocess()
|
|
self._worker.terminate()
|
|
self._trainer.stop()
|
|
self._log_append("학습 중지됨")
|
|
log_action("[재학습] 학습 중지 (사용자)")
|
|
log_train("학습 중지됨 (사용자 중지)")
|
|
self._status_lbl.setText("중지됨")
|
|
self._status_lbl.setStyleSheet("color:#A32D2D; font-size:14px; font-weight:bold;")
|
|
self._start_btn.setEnabled(True)
|
|
self._stop_btn.setEnabled(False)
|
|
|
|
def _on_log(self, message: str):
|
|
ts = datetime.now().strftime("%H:%M:%S")
|
|
self._log_box.append(f"[{ts}] {message}")
|
|
sb = self._log_box.verticalScrollBar()
|
|
sb.setValue(sb.maximum())
|
|
|
|
def _on_progress(self, value: int):
|
|
self._progress_bar.setValue(value)
|
|
self._progress_bar.setFormat(f"{value}%")
|
|
|
|
def _on_finished(self, success: bool):
|
|
self._start_btn.setEnabled(True)
|
|
self._stop_btn.setEnabled(False)
|
|
if success:
|
|
self._save_btn.setEnabled(True)
|
|
self._status_lbl.setText("학습 완료")
|
|
self._status_lbl.setStyleSheet("color:#22cc55; font-size:14px; font-weight:bold;")
|
|
log_train("학습 완료")
|
|
QMessageBox.information(self, "완료", "학습 완료!")
|
|
else:
|
|
self._status_lbl.setText("학습 실패")
|
|
self._status_lbl.setStyleSheet("color:#A32D2D; font-size:14px; font-weight:bold;")
|
|
log_train("학습 실패")
|
|
QMessageBox.critical(self, "실패", "학습 실패. 로그를 확인해주세요.")
|
|
|
|
def _on_model_save(self):
|
|
dest_input = self._save_path_edit.text().strip() or "ai/models/best.pt"
|
|
dest = resolve_path(dest_input)
|
|
src = resolve_path("ai/runs/train/weights/best.pt")
|
|
if not os.path.exists(src):
|
|
QMessageBox.warning(self, "경고", f"학습 결과 파일을 찾을 수 없습니다:\n{src}")
|
|
return
|
|
try:
|
|
os.makedirs(os.path.dirname(dest), exist_ok=True)
|
|
shutil.copy(src, dest)
|
|
saved_path = to_project_relative(dest)
|
|
# config.json ai.model_path 업데이트
|
|
config_path = resolve_path("config.json")
|
|
if os.path.exists(config_path):
|
|
with open(config_path, "r", encoding="utf-8") as fh:
|
|
cfg = json.load(fh)
|
|
cfg.setdefault("ai", {})["model_path"] = saved_path
|
|
with open(config_path, "w", encoding="utf-8") as fh:
|
|
json.dump(cfg, fh, ensure_ascii=False, indent=2)
|
|
self._save_result_lbl.setText(f"저장 완료: {saved_path}")
|
|
self._save_result_lbl.setStyleSheet("color:#22cc55; font-size:13px;")
|
|
log_action(f"[재학습] 모델 저장 완료 → {saved_path}")
|
|
log_train(f"모델 저장 완료 | {saved_path}")
|
|
QMessageBox.information(self, "완료", f"모델 저장 완료\n{dest}")
|
|
except Exception as e:
|
|
log_train(f"모델 저장 실패 | {e}")
|
|
QMessageBox.critical(self, "저장 실패", str(e))
|
|
|
|
# ================================================================ #
|
|
# 헬퍼
|
|
# ================================================================ #
|
|
|
|
def _refresh_box_list(self):
|
|
self._box_list.blockSignals(True)
|
|
self._box_list.clear()
|
|
for box in self._canvas.boxes:
|
|
r = box["rect"]
|
|
color = LabelingCanvas.CLASS_COLORS.get(box["class_name"], "#ffffff")
|
|
item = QListWidgetItem(
|
|
f"[{box['class_name']}] x:{r.x()} y:{r.y()} w:{r.width()} h:{r.height()}"
|
|
)
|
|
item.setForeground(QColor(color))
|
|
self._box_list.addItem(item)
|
|
self._box_list.setCurrentRow(self._canvas.selected_index)
|
|
self._box_list.blockSignals(False)
|
|
|
|
def _log_append(self, text: str):
|
|
ts = datetime.now().strftime("%H:%M:%S")
|
|
self._log_box.append(f"[{ts}] {text}")
|
|
sb = self._log_box.verticalScrollBar()
|
|
sb.setValue(sb.maximum())
|
|
|
|
|
|
# ============================================================ #
|
|
# 스타일 헬퍼
|
|
# ============================================================ #
|
|
|
|
def _btn_style(bg: str, font_size: int = 14, bold: bool = False) -> str:
|
|
weight = "bold" if bold else "normal"
|
|
return (
|
|
f"QPushButton {{"
|
|
f" background:{bg}; color:#ffffff; border:none; border-radius:4px;"
|
|
f" font-size:{font_size}px; font-weight:{weight}; min-height:28px;"
|
|
f"}}"
|
|
f"QPushButton:hover {{ background:{_lighten(bg)}; }}"
|
|
f"QPushButton:pressed {{ background:{_darken(bg)}; }}"
|
|
f"QPushButton:disabled {{ background:#3a3a3a; color:#666666; }}"
|
|
)
|
|
|
|
|
|
def _cls_btn_style(color: str, checked: bool) -> str:
|
|
bg = color if checked else "#333333"
|
|
border = f"border:2px solid {color};" if checked else "border:1px solid #555555;"
|
|
return (
|
|
f"QPushButton {{"
|
|
f" background:{bg}; color:#ffffff; {border} border-radius:4px;"
|
|
f" font-size:15px; font-weight:bold; min-height:56px;"
|
|
f"}}"
|
|
f"QPushButton:hover {{ background:{_lighten(bg)}; }}"
|
|
)
|
|
|
|
|
|
def _spinbox_style() -> str:
|
|
return (
|
|
"QSpinBox {"
|
|
" background:#2a2a2a; color:#ffffff; border:1px solid #555555;"
|
|
" border-radius:4px; padding:4px 8px; font-size:14px; min-height:38px;"
|
|
"}"
|
|
)
|
|
|
|
|
|
def _lighten(hex_color: str) -> str:
|
|
try:
|
|
r, g, b = int(hex_color[1:3], 16), int(hex_color[3:5], 16), int(hex_color[5:7], 16)
|
|
return f"#{min(r+30,255):02x}{min(g+30,255):02x}{min(b+30,255):02x}"
|
|
except Exception:
|
|
return hex_color
|
|
|
|
|
|
def _darken(hex_color: str) -> str:
|
|
try:
|
|
r, g, b = int(hex_color[1:3], 16), int(hex_color[3:5], 16), int(hex_color[5:7], 16)
|
|
return f"#{max(r-30,0):02x}{max(g-30,0):02x}{max(b-30,0):02x}"
|
|
except Exception:
|
|
return hex_color
|