Dự án chặng 2: công cụ chụp và lưu ảnh kiểm tra OK/NG bằng OpenCV

"""Công cụ chụp ảnh dữ liệu OK/NG từ camera hoặc video.

Chạy: python chup_anh.py 0            (camera số 0)
      python chup_anh.py video.mp4    (file video)
Phím: r chọn ROI · o lưu OK · n lưu NG · q thoát
"""
import json
import sys
import time
from datetime import datetime
from pathlib import Path

import cv2
import numpy as np

FILE_CAU_HINH = Path("chup_anh.json")
THU_MUC_GOC = Path("dataset")
CUA_SO = "Chup anh kiem tra"


def mo_nguon(arg: str) -> cv2.VideoCapture:
    nguon = int(arg) if arg.isdigit() else arg
    cap = cv2.VideoCapture(nguon)
    if not cap.isOpened():
        raise SystemExit(f"Không mở được nguồn video: {arg}")
    return cap


def doc_roi():
    if FILE_CAU_HINH.exists():
        roi = json.loads(FILE_CAU_HINH.read_text(encoding="utf-8")).get("roi")
        if roi:
            return tuple(roi)
    return None


def luu_roi(roi) -> None:
    FILE_CAU_HINH.write_text(json.dumps({"roi": list(roi)}, indent=2), encoding="utf-8")


def cat_roi(khung, roi):
    if roi is None:
        return khung
    x, y, w, h = roi
    cao, rong = khung.shape[:2]
    x1, y1, x2, y2 = max(0, x), max(0, y), min(rong, x + w), min(cao, y + h)
    return khung[y1:y2, x1:x2] if x1 < x2 and y1 < y2 else khung


def chi_so(vung) -> tuple[float, float]:
    xam = cv2.cvtColor(vung, cv2.COLOR_BGR2GRAY)
    return float(xam.mean()), float(cv2.Laplacian(xam, cv2.CV_64F).var())


def luu_anh(vung, nhan: str) -> Path:
    bay_gio = datetime.now()
    thu_muc = THU_MUC_GOC / bay_gio.strftime("%Y-%m-%d") / nhan
    thu_muc.mkdir(parents=True, exist_ok=True)
    duong_dan = thu_muc / f"{bay_gio.strftime('%H%M%S_%f')[:-3]}_{nhan}.png"
    dem = 1
    while duong_dan.exists():                      # phòng khi bấm 2 lần trong cùng 1 ms
        duong_dan = duong_dan.with_name(f"{duong_dan.stem}_{dem}.png")
        dem += 1
    cv2.imwrite(str(duong_dan), vung)
    return duong_dan


def ve_thong_tin(khung, roi, fps, dem, sang, net):
    ve = khung.copy()
    if roi:
        x, y, w, h = roi
        cv2.rectangle(ve, (x, y), (x + w, y + h), (0, 220, 255), 2)
    dong = [f"FPS {fps:4.1f}   OK {dem['OK']}   NG {dem['NG']}",
            f"Sang {sang:5.1f}   Net {net:7.1f}",
            "r: chon ROI  o: luu OK  n: luu NG  q: thoat"]
    for i, chu in enumerate(dong):
        y = 28 + i * 26
        cv2.putText(ve, chu, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 4, cv2.LINE_AA)       # viền đen
        cv2.putText(ve, chu, (12, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1, cv2.LINE_AA)
    return ve


def main(arg: str = "0") -> dict:
    cap = mo_nguon(arg)
    roi = doc_roi()
    dem = {"OK": 0, "NG": 0}
    fps, t_truoc = 0.0, time.perf_counter()
    try:
        while True:
            co, khung = cap.read()
            if not co:
                print("Hết video hoặc mất kết nối camera")
                break
            bay_gio = time.perf_counter()
            fps = 0.9 * fps + 0.1 / max(bay_gio - t_truoc, 1e-6)      # FPS làm mượt
            t_truoc = bay_gio

            vung = cat_roi(khung, roi)
            sang, net = chi_so(vung)
            cv2.imshow(CUA_SO, ve_thong_tin(khung, roi, fps, dem, sang, net))

            phim = cv2.waitKey(1) & 0xFF
            if phim == ord("q"):
                break
            if phim == ord("r"):
                chon = cv2.selectROI(CUA_SO, khung, showCrosshair=True)
                if chon[2] > 0 and chon[3] > 0:
                    roi = tuple(int(v) for v in chon)
                    luu_roi(roi)
                    print("ROI mới:", roi)
            elif phim in (ord("o"), ord("n")):
                nhan = "OK" if phim == ord("o") else "NG"
                p = luu_anh(vung, nhan)          # lưu ảnh gốc, không phải ảnh đã vẽ
                dem[nhan] += 1
                print(f"Đã lưu {p}")
    finally:
        cap.release()
        cv2.destroyAllWindows()
    print(f"Tổng kết: {dem['OK']} ảnh OK, {dem['NG']} ảnh NG trong {THU_MUC_GOC}")
    return dem


if __name__ == "__main__":
    main(sys.argv[1] if len(sys.argv) > 1 else "0")