130 lines
4.3 KiB
Python
130 lines
4.3 KiB
Python
import os
|
||
import cv2
|
||
from ultralytics import YOLO
|
||
from pathlib import Path
|
||
from PIL import Image, ImageDraw, ImageFont
|
||
|
||
# ==================== 【配置区域】 ====================
|
||
class Config:
|
||
# 你的Plate检测模型路径
|
||
MODEL_PATH = r"best.pt"
|
||
|
||
CONF_THRESHOLD = 0.5
|
||
IOU_THRESHOLD = 0.45
|
||
|
||
BOX_THICKNESS = 2
|
||
FONT_SIZE = 12
|
||
|
||
# 输入输出路径
|
||
INPUT_PATH = r"D:\Comproject\image_classify\trian"
|
||
OUTPUT_DIR = r"D:\Comproject\image_classify\Copy"
|
||
|
||
BATCH_MODE = True
|
||
GENERATE_VISUALIZATION = True
|
||
|
||
# 只有 Plate 一个类别
|
||
LABEL_NAMES = {0: "Plate"}
|
||
COLORS = {'Plate': (0, 255, 0)}
|
||
|
||
IMAGE_EXTENSIONS = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']
|
||
|
||
# ========================================================
|
||
|
||
class YOLOv11PlateExtractor:
|
||
def __init__(self, config):
|
||
self.config = config
|
||
self.model = YOLO(config.MODEL_PATH)
|
||
print("✅ Plate 模型加载成功")
|
||
|
||
def predict(self, image_path):
|
||
return self.model.predict(
|
||
source=image_path,
|
||
conf=self.config.CONF_THRESHOLD,
|
||
iou=self.config.IOU_THRESHOLD,
|
||
verbose=False
|
||
)[0]
|
||
|
||
def crop_plate(self, image, results):
|
||
cropped_regions = []
|
||
if results.boxes is not None:
|
||
boxes = results.boxes.xyxy.cpu().numpy()
|
||
confs = results.boxes.conf.cpu().numpy()
|
||
|
||
h, w = image.shape[:2]
|
||
for box, conf in zip(boxes, confs):
|
||
x1, y1, x2, y2 = map(int, box)
|
||
x1 = max(0, x1)
|
||
y1 = max(0, y1)
|
||
x2 = min(w, x2)
|
||
y2 = min(h, y2)
|
||
if x2 > x1 and y2 > y1:
|
||
crop = image[y1:y2, x1:x2]
|
||
if crop.size > 0:
|
||
cropped_regions.append({"image": crop, "conf": conf})
|
||
return cropped_regions
|
||
|
||
def visualize_results(self, image, results, save_path):
|
||
img = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
|
||
draw = ImageDraw.Draw(img)
|
||
try:
|
||
font = ImageFont.truetype("simsun.ttc", self.config.FONT_SIZE)
|
||
except:
|
||
font = ImageFont.load_default()
|
||
|
||
if results.boxes is not None:
|
||
for box, conf in zip(results.boxes.xyxy, results.boxes.conf):
|
||
x1, y1, x2, y2 = map(int, box)
|
||
draw.rectangle([(x1, y1), (x2, y2)], outline=(0,255,0), width=self.config.BOX_THICKNESS)
|
||
draw.text((x1, y1-15), f"Plate {conf:.2f}", fill=(0,255,0), font=font)
|
||
img.save(save_path)
|
||
|
||
def process_single_image(self, image_path):
|
||
image = cv2.imread(image_path)
|
||
if image is None:
|
||
print(f"❌ 读取失败: {image_path}")
|
||
return
|
||
|
||
original_stem = Path(image_path).stem
|
||
results = self.predict(image_path)
|
||
cropped_regions = self.crop_plate(image, results)
|
||
|
||
# 所有裁剪图直接保存到 OUTPUT_DIR 根目录
|
||
os.makedirs(self.config.OUTPUT_DIR, exist_ok=True)
|
||
|
||
for i, region in enumerate(cropped_regions):
|
||
img_crop = region["image"]
|
||
# 命名：原图名_Plate_序号.png
|
||
filename = f"{original_stem}_Plate_{i+1}.png"
|
||
save_path = os.path.join(self.config.OUTPUT_DIR, filename)
|
||
cv2.imwrite(save_path, img_crop)
|
||
|
||
# 可视化
|
||
if self.config.GENERATE_VISUALIZATION:
|
||
vis_dir = os.path.join(self.config.OUTPUT_DIR, "visualization")
|
||
os.makedirs(vis_dir, exist_ok=True)
|
||
self.visualize_results(image, results, os.path.join(vis_dir, f"{original_stem}.jpg"))
|
||
|
||
print(f"✅ {original_stem} → 裁剪盘子数量：{len(cropped_regions)}")
|
||
|
||
def process_batch(self):
|
||
files = []
|
||
for ext in self.config.IMAGE_EXTENSIONS:
|
||
files.extend(Path(self.config.INPUT_PATH).glob(f"*{ext}"))
|
||
files = sorted(list({str(f) for f in files}))
|
||
|
||
print(f"\n📂 共找到 {len(files)} 张图片")
|
||
for f in files:
|
||
self.process_single_image(f)
|
||
|
||
def run(self):
|
||
print("\n========== Plate 检测 & 批量裁剪（平铺输出） ==========")
|
||
self.process_batch()
|
||
print("\n🎉 全部处理完成！")
|
||
|
||
def main():
|
||
cfg = Config()
|
||
app = YOLOv11PlateExtractor(cfg)
|
||
app.run()
|
||
|
||
if __name__ == "__main__":
|
||
main() |