220 lines
6.5 KiB
Python
220 lines
6.5 KiB
Python
import os
|
|
import sys
|
|
import cv2
|
|
import torch
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
# =========================
|
|
# 需要确保本脚本能找到 YOLOv5 源码
|
|
# 如果脚本不在 yolov5 根目录，可手动添加：
|
|
# sys.path.append("/path/to/yolov5")
|
|
# =========================
|
|
|
|
from models.common import DetectMultiBackend
|
|
from utils.general import (
|
|
check_img_size,
|
|
non_max_suppression,
|
|
scale_boxes
|
|
)
|
|
from utils.torch_utils import select_device
|
|
from utils.augmentations import letterbox
|
|
|
|
|
|
def is_image_file(filename):
|
|
image_exts = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".webp"}
|
|
return Path(filename).suffix.lower() in image_exts
|
|
|
|
|
|
def load_image(img_path, img_size=640, stride=32):
|
|
"""
|
|
读取图像，并进行与YOLOv5推理一致的预处理
|
|
返回：
|
|
im0: 原图(BGR)
|
|
im: 预处理后的输入张量所对应的numpy图像(CHW, RGB)
|
|
"""
|
|
im0 = cv2.imread(str(img_path))
|
|
if im0 is None:
|
|
raise ValueError(f"无法读取图片: {img_path}")
|
|
|
|
# letterbox缩放
|
|
im = letterbox(im0, img_size, stride=stride, auto=True)[0]
|
|
|
|
# BGR -> RGB, HWC -> CHW
|
|
im = im[:, :, ::-1].transpose(2, 0, 1)
|
|
im = im.copy()
|
|
|
|
return im0, im
|
|
|
|
|
|
def save_crop(img, xyxy, save_path):
|
|
"""
|
|
根据检测框裁剪并保存
|
|
xyxy: [x1, y1, x2, y2]
|
|
"""
|
|
x1, y1, x2, y2 = map(int, xyxy)
|
|
|
|
h, w = img.shape[:2]
|
|
x1 = max(0, min(x1, w - 1))
|
|
x2 = max(0, min(x2, w - 1))
|
|
y1 = max(0, min(y1, h - 1))
|
|
y2 = max(0, min(y2, h - 1))
|
|
|
|
if x2 <= x1 or y2 <= y1:
|
|
return False
|
|
|
|
crop = img[y1:y2, x1:x2]
|
|
if crop.size == 0:
|
|
return False
|
|
|
|
cv2.imwrite(str(save_path), crop)
|
|
return True
|
|
|
|
|
|
def run(
|
|
weights,
|
|
source,
|
|
output_dir,
|
|
imgsz=640,
|
|
conf_thres=0.25,
|
|
iou_thres=0.45,
|
|
device='',
|
|
max_det=1000,
|
|
classes=None,
|
|
class_num=5
|
|
):
|
|
# 设备
|
|
device = select_device(device)
|
|
|
|
# 加载模型
|
|
model = DetectMultiBackend(weights, device=device, dnn=False, fp16=False)
|
|
stride, names, pt = model.stride, model.names, model.pt
|
|
imgsz = check_img_size(imgsz, s=stride)
|
|
|
|
# 创建输出目录和类别子目录
|
|
output_dir = Path(output_dir)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
for i in range(class_num):
|
|
(output_dir / str(i)).mkdir(parents=True, exist_ok=True)
|
|
|
|
# 收集图片
|
|
source = Path(source)
|
|
if not source.exists():
|
|
raise FileNotFoundError(f"输入文件夹不存在: {source}")
|
|
|
|
image_files = [p for p in source.iterdir() if p.is_file() and is_image_file(p.name)]
|
|
image_files.sort()
|
|
|
|
if len(image_files) == 0:
|
|
print(f"在文件夹中未找到图片: {source}")
|
|
return
|
|
|
|
print(f"共找到 {len(image_files)} 张图片")
|
|
print(f"类别名称: {names}")
|
|
|
|
total_crops = 0
|
|
|
|
for img_idx, img_path in enumerate(image_files, 1):
|
|
try:
|
|
im0, im = load_image(img_path, img_size=imgsz, stride=stride)
|
|
|
|
# 转tensor
|
|
im = torch.from_numpy(im).to(device)
|
|
im = im.float() / 255.0
|
|
if im.ndimension() == 3:
|
|
im = im.unsqueeze(0)
|
|
|
|
# 推理
|
|
pred = model(im)
|
|
|
|
# NMS
|
|
pred = non_max_suppression(
|
|
pred,
|
|
conf_thres,
|
|
iou_thres,
|
|
classes=classes,
|
|
agnostic=False,
|
|
max_det=max_det
|
|
)
|
|
|
|
det = pred[0]
|
|
|
|
crop_count_this_image = 0
|
|
|
|
if len(det):
|
|
# 将框映射回原图尺寸
|
|
det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()
|
|
|
|
# 遍历每个检测目标
|
|
for det_idx, (*xyxy, conf, cls) in enumerate(det):
|
|
cls_id = int(cls.item())
|
|
|
|
# 只保存 0~4 这5类
|
|
if cls_id < 0 or cls_id >= class_num:
|
|
continue
|
|
|
|
# 保存文件名：原图名_序号_类别_置信度.jpg
|
|
stem = img_path.stem
|
|
conf_val = float(conf.item())
|
|
save_name = f"{stem}_{det_idx:03d}_cls{cls_id}_{conf_val:.4f}.jpg"
|
|
save_path = output_dir / str(cls_id) / save_name
|
|
|
|
ok = save_crop(im0, xyxy, save_path)
|
|
if ok:
|
|
crop_count_this_image += 1
|
|
total_crops += 1
|
|
|
|
print(f"[{img_idx}/{len(image_files)}] {img_path.name} -> 裁剪保存 {crop_count_this_image} 个目标")
|
|
|
|
except Exception as e:
|
|
print(f"处理失败: {img_path}，原因: {e}")
|
|
|
|
print("\n处理完成")
|
|
print(f"总共保存裁剪目标数: {total_crops}")
|
|
print(f"输出目录: {output_dir.resolve()}")
|
|
|
|
|
|
def parse_opt():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--weights', type=str, required=True, help='训练好的YOLOv5 pt权重路径')
|
|
parser.add_argument('--source', type=str, required=True, help='待检测图片文件夹路径')
|
|
parser.add_argument('--output-dir', type=str, required=True, help='裁剪结果输出文件夹')
|
|
parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='推理尺寸')
|
|
parser.add_argument('--conf-thres', type=float, default=0.25, help='置信度阈值')
|
|
parser.add_argument('--iou-thres', type=float, default=0.45, help='NMS IoU阈值')
|
|
parser.add_argument('--device', default='', help='cuda device，例如 0 或 0,1,2,3 或 cpu')
|
|
parser.add_argument('--max-det', type=int, default=1000, help='每张图最多检测目标数')
|
|
parser.add_argument('--class-num', type=int, default=5, help='类别总数，这里默认是5')
|
|
parser.add_argument('--classes', nargs='+', type=int, default=None, help='只保留某些类别，例如 --classes 0 1 2')
|
|
opt = parser.parse_args()
|
|
|
|
# imgsz兼容处理
|
|
if len(opt.imgsz) == 1:
|
|
opt.imgsz = opt.imgsz[0]
|
|
elif len(opt.imgsz) == 2:
|
|
# yolov5这里一般接受单值，若传两个值则取最大，避免报错
|
|
opt.imgsz = max(opt.imgsz)
|
|
else:
|
|
raise ValueError("imgsz 参数格式错误")
|
|
|
|
return opt
|
|
|
|
|
|
def main(opt):
|
|
run(
|
|
weights=opt.weights,
|
|
source=opt.source,
|
|
output_dir=opt.output_dir,
|
|
imgsz=opt.imgsz,
|
|
conf_thres=opt.conf_thres,
|
|
iou_thres=opt.iou_thres,
|
|
device=opt.device,
|
|
max_det=opt.max_det,
|
|
classes=opt.classes,
|
|
class_num=opt.class_num
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
opt = parse_opt()
|
|
main(opt) |