支持單張圖片/圖片目錄批量預標注 默認使用cuda GPU
.env

HTTP_PROXY=http://192.168.2.109:10808
HTTPS_PROXY=http://192.168.2.109:10808

pyproject.toml

[project]
name = "yolo-test"
version = "0.1.0"
description = "Add your description here"
readme = "README.md"
requires-python = ">=3.11"
dependencies = ["opencv-python<=4.8.1.78","ultralytics>=8.3.198","torch==2.1.2","torchvision==0.16.2","torchaudio==2.1.2","python-dotenv>=1.1.1","pyinstaller>=6.15.0","pyarmor>=9.1.9","numpy<2",
]

main.py

import argparse
import json
import os
from dotenv import load_dotenv
from ultralytics import YOLO
from PIL import Image, UnidentifiedImageError
import torchclass YoloPreIdentify:""""""def __init__(self, image_path, proxy=False):if not os.path.exists(image_path):raise FileNotFoundError(image_path)if proxy:self.load_proxy()self.include_class = self.read_files()self.class_name_set = set()self.image_path = image_pathif torch.cuda.is_available():self.device = "0"else:self.device = "cpu"print(f"正在加載 YOLOv8 模型 ---- {self.device}...")self.model = YOLO('yolov8n.pt')print("YOLOv8 模型加載完成...")def load_proxy(self):""":return:"""load_dotenv()print("環境變量已加載")def get_file_from_path(self, path):""":param path::return:"""for root, dirs, files in os.walk(path):if files:for filename in files:abs_path = os.path.join(root, filename)yield filename, abs_pathdef check_images(self, file):""":param file::return:"""try:with Image.open(file) as img:img.verify()return Trueexcept (IOError, SyntaxError, UnidentifiedImageError) as e:print(f"文件 '{file}' 不是有效圖片，錯誤: {e}")return Falsedef read_files(self, file="include.txt"):""":param file::return:"""uninclude_set = set()if not os.path.exists(file):return uninclude_setwith open(file, "r", encoding="utf-8") as f:for i in f:line = i.strip()if line:uninclude_set.add(line)return uninclude_setdef save_results(self, data):""":param data::return:"""with open("result.json", "w", encoding="utf-8") as f:json.dump(data, f, ensure_ascii=False, indent=4)f.flush()def process_single(self, img_file):""":param img_file::return:"""tmp = []if not self.check_images(img_file):return tmpprint(f"開始檢測：{img_file}")results = self.model(img_file, device=self.device)print("檢測完成。")print("\n--- 檢測結果 ---")for result in results:boxes = result.boxes# 檢查是否檢測到了任何物體if len(boxes) == 0:print("圖片中未檢測到任何目標。")return tmpfor box in boxes:class_id = int(box.cls.item())class_name = self.model.names[class_id]self.class_name_set.add(class_name)if class_name not in self.include_class:continueconfidence = float(box.conf.item())coords = box.xyxy.tolist()[0]  # 轉換為 Python 列表x_min = int(coords[0])y_min = int(coords[1])x_max = int(coords[2])y_max = int(coords[3])# 打印格式化的結果print(f"檢測到目標: {class_name}")print(f"  - 置信度: {confidence:.2f}")  # 保留兩位小數print(f"  - 坐標 (左上角 x,y): ({x_min}, {y_min})")print(f"  - 坐標 (右下角 x,y): ({x_max}, {y_max})")print("-" * 20)  # 分隔符tmp.append({"class": class_name,"confidence": round(confidence, 2),"points": [{"x": x_min, "y": y_min},{"x": x_max, "y": y_max},]})return tmpdef process_batch(self):""":return:"""result_map = {}for filename, img_file in self.get_file_from_path(self.image_path):result_map.update({img_file: self.process_single(img_file)})self.save_results(result_map)print(f"文件夾中識別到標簽：{self.class_name_set}")return result_mapif __name__ == '__main__':parser = argparse.ArgumentParser()parser.add_argument('-i', "--input", dest='input', help='input', required=True)parser.add_argument('-p', "--proxy", dest='proxy', help='proxy', default=False)args = parser.parse_args()ypi = YoloPreIdentify(args.input, args.proxy)ypi.process_batch()

include.txt (預標注標注類型)

cat
dog
horse
person
bird

環境同步：

 uv sync --index-url https://pypi.tuna.tsinghua.edu.cn/simple

打包命令：

pyarmor gen -r --pack FC main.py && rm -rf .pyarmor* && rm main.spec && mv dist/main ./ 

執行命令：

./main -i "{image_folder}"