# 재학습 페이지 — 이미지 로드·라벨링 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