Python Yolo8 物体识别
支持单张图片/图片目录批量预标注 默认使用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}"