中药图片分类数据集采集程序教程 简单

import PySimpleGUI as sg
import cv2
import os

# 图像尺寸设置
PREVIEW_SIZE = (160, 120)
SAVE_SIZE = (224, 224)
JPEG_QUALITY = 85

layout = [
    [sg.Image(filename='', key='-IMAGE-', size=PREVIEW_SIZE)],
    [sg.Text('保存路径：'), sg.Input(key='-PATH-'), sg.FolderBrowse()],
    [sg.Button('开始拍摄'), sg.Button('退出')],
    [sg.Text('状态：准备就绪', key='-STATUS-')]
]

window = sg.Window('数据集采集工具', layout)
cap = cv2.VideoCapture(0)

if not cap.isOpened():
    sg.popup_error('无法打开摄像头！')
    exit()

capturing = False
current_count = 0
round_count = 0

while True:
    event, values = window.read(timeout=20)
    
    if event in (None, '退出'):
        break
    
    if event == '开始拍摄':
        save_path = values['-PATH-']
        if not save_path:
            sg.popup_error('请先选择保存路径！')
            continue
        
        try:
            # 修复1：使用更安全的路径创建方式
            os.makedirs(save_path, exist_ok=True)
            # 修复2：验证路径可写性
            test_file = os.path.join(save_path, 'test.txt')
            with open(test_file, 'w') as f:
                f.write('test')
            os.remove(test_file)
        except Exception as e:
            sg.popup_error(f'路径无效或不可写：{str(e)}')
            continue
        
        capturing = True
        current_count = 0
        window['-STATUS-'].update(f'拍摄中：第 {chr(65 + round_count)} 轮，0/20')

    ret, frame = cap.read()
    if not ret:
        sg.popup_error('无法获取视频帧！')
        break

    preview_frame = cv2.resize(frame, PREVIEW_SIZE)
    img_bytes = cv2.imencode('.png', preview_frame)[1].tobytes()
    window['-IMAGE-'].update(data=img_bytes)

    if capturing and current_count < 20:
        save_frame = cv2.resize(frame, SAVE_SIZE)
        filename = f"{chr(65 + round_count)}{current_count + 1:02d}.jpg"
        filepath = os.path.join(save_path, filename)
        
        try:
            # 修复3：添加写入验证
            success = cv2.imwrite(filepath, save_frame, [cv2.IMWRITE_JPEG_QUALITY, JPEG_QUALITY])
            if not success:
                raise Exception("OpenCV写入失败")
            current_count += 1
            window['-STATUS-'].update(f'拍摄中：第 {chr(65 + round_count)} 轮，{current_count}/20')
        except Exception as e:
            sg.popup_error(f'保存失败：{str(e)}')
            capturing = False
            continue
        
        if current_count >= 20:
            capturing = False
            round_count += 1
            sg.popup(f'完成 {chr(65 + round_count - 1)} 轮拍摄！')
            window['-STATUS-'].update('状态：准备就绪')

cap.release()
window.close()