import zipfile
import os
from pathlib import Path

def unzip_with_python(zip_path, dest_dir):
    """使用Python解压zip文件"""
    zip_path = Path(zip_path)
    dest_dir = Path(dest_dir)
    
    print(f"Unzipping {zip_path} to {dest_dir}...")
    
    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
        zip_ref.extractall(dest_dir)
    
    print(f"Done! Extracted {len(zip_ref.namelist())} files")

def generate_image_list(img_dir, output_file):
    """生成图片路径列表文件"""
    img_dir = Path(img_dir)
    output_file = Path(output_file)
    
    img_files = sorted(img_dir.glob("*.jpg"))
    
    print(f"Found {len(img_files)} images in {img_dir}")
    
    with open(output_file, 'w') as f:
        for img_path in img_files:
            f.write(str(img_path) + '\n')
    
    print(f"Generated {output_file} with {len(img_files)} entries")

if __name__ == "__main__":
    base_dir = Path(r"C:\Users\16337\PycharmProjects\Security\datasets\coco")
    
    # 解压 val2017.zip
    val_zip = base_dir / "images" / "val2017.zip"
    val_dir = base_dir / "images" / "val2017"
    
    if val_zip.exists() and not any(val_dir.glob("*.jpg")):
        unzip_with_python(val_zip, base_dir / "images")
    
    # 生成 val2017.txt
    generate_image_list(val_dir, base_dir / "val2017.txt")
    
    print("\nDone! Now you can run:")
    print('yolo val model=yolo11n_int8_b1_8.engine data=coco.yaml imgsz=640 rect=False')