YOLOとEfficientNetで麻雀牌を推論する｜カメラ映像から二段階分類する

import cv2
import timm
import torch
from PIL import Image
from pathlib import Path
from ultralytics import YOLO
from datetime import datetime
from torchvision import transforms


# ==================================================
# 設定
# ==================================================
YOLO_MODEL_PATH = Path("./models/best.pt")
CLASSIFICATION_MODEL_DIR = Path("./models/20260729_220942")
RESULT_IMAGE_PATH = Path("./result.jpg")

SAVE_SAMPLES = True
SAMPLE_DIR = Path("./samples/uncheck")

YOLO_IMAGE_SIZE = 640
YOLO_CONFIDENCE = 0.5

STAGE2_FILENAMES = {
    "manzu": "manzu.pth",
    "pinzu": "pinzu.pth",
    "souzu": "souzu.pth",
    "jihai": "jihai.pth",
}

TILE_NAMES = {
    0: "1萬", 1: "2萬", 2: "3萬", 3: "4萬", 4: "5萬", 5: "6萬", 6: "7萬", 7: "8萬", 8: "9萬", 9: "赤5萬",
    10: "1筒", 11: "2筒", 12: "3筒", 13: "4筒", 14: "5筒", 15: "6筒", 16: "7筒", 17: "8筒", 18: "9筒", 19: "赤5筒",
    20: "1索", 21: "2索", 22: "3索", 23: "4索", 24: "5索", 25: "6索", 26: "7索", 27: "8索", 28: "9索", 29: "赤5索",
    30: "東", 31: "南", 32: "西", 33: "北", 34: "白", 35: "發", 36: "中",
}


def select_device():
    if torch.backends.mps.is_available():
        return "mps"
    if torch.cuda.is_available():
        return "cuda"
    return "cpu"


def load_checkpoint(path, device):
    return torch.load(path, map_location=device, weights_only=False)


class TorchImageClassifier:
    def __init__(self, checkpoint_path, device):
        checkpoint = load_checkpoint(checkpoint_path, device)
        self.class_names = [str(value) for value in checkpoint["class_names"]]
        self.model = timm.create_model("efficientnet_b0", pretrained=False, num_classes=len(self.class_names))
        self.model.load_state_dict(checkpoint["model_state_dict"])
        self.model = self.model.to(device)
        self.model.eval()
        self.device = device

    def predict(self, tensor):
        batch = tensor.unsqueeze(0).to(self.device)
        with torch.no_grad():
            outputs = self.model(batch)
            probabilities = torch.softmax(outputs, dim=1)
            confidence, predicted = torch.max(probabilities, dim=1)
        return self.class_names[predicted.item()], confidence.item()


class HierarchicalTileClassifier:
    def __init__(self, model_dir, device):
        self.stage1 = TorchImageClassifier(model_dir / "class4.pth", device)
        self.stage2 = {
            group: TorchImageClassifier(model_dir / filename, device)
            for group, filename in STAGE2_FILENAMES.items()
        }
        self.transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ])

    def classify(self, crop_image):
        rgb = cv2.cvtColor(crop_image, cv2.COLOR_BGR2RGB)
        tensor = self.transform(Image.fromarray(rgb))
        group, stage1_conf = self.stage1.predict(tensor)
        tile_name, stage2_conf = self.stage2[group].predict(tensor)
        tile_id = int(tile_name)
        confidence = stage1_conf * stage2_conf

        return {
            "tile_id": tile_id,
            "group": group,
            "stage1_conf": stage1_conf,
            "stage2_conf": stage2_conf,
            "confidence": confidence,
        }


def detect_tiles(image, yolo_model, classifier, device):
    results = yolo_model.predict(
        source=image,
        imgsz=YOLO_IMAGE_SIZE,
        conf=YOLO_CONFIDENCE,
        device=device,
        verbose=False,
    )

    detections = []
    boxes = results[0].boxes

    if boxes is not None:
        for box in boxes:
            x1, y1, x2, y2 = [int(value) for value in box.xyxy[0].tolist()]
            detections.append((x1, y1, x2, y2))

    detections.sort(key=lambda box: box[0])
    predictions = []

    for x1, y1, x2, y2 in detections:
        crop = image[y1:y2, x1:x2]
        if crop.size == 0:
            continue

        prediction = classifier.classify(crop)
        prediction["box"] = (x1, y1, x2, y2)
        prediction["crop"] = crop.copy()
        predictions.append(prediction)

    return predictions


def save_samples(predictions):
    SAMPLE_DIR.mkdir(parents=True, exist_ok=True)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")

    for index, prediction in enumerate(predictions, start=1):
        tile_id = prediction["tile_id"]
        filename = f"{tile_id}_{timestamp}_{index:02d}.jpg"
        cv2.imwrite(str(SAMPLE_DIR / filename), prediction["crop"])


def draw_result(image, predictions):
    result = image.copy()

    for prediction in predictions:
        x1, y1, x2, y2 = prediction["box"]
        tile_id = prediction["tile_id"]
        confidence = prediction["confidence"]

        cv2.rectangle(result, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(
            result,
            f"ID:{tile_id} {confidence:.2f}",
            (x1, max(20, y1 - 5)),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.6,
            (0, 255, 0),
            2,
            cv2.LINE_AA,
        )

    return result


def run_inference(frame, yolo_model, classifier, device):
    predictions = detect_tiles(frame, yolo_model, classifier, device)

    print(f"\n検出枚数: {len(predictions)}")
    for index, prediction in enumerate(predictions, start=1):
        tile_id = prediction["tile_id"]
        tile_name = TILE_NAMES[tile_id]
        print(
            f"{index:02d}: ID={tile_id:02d} {tile_name} / "
            f"Stage1={prediction['group']} ({prediction['stage1_conf']:.3f}) / "
            f"Stage2={prediction['stage2_conf']:.3f} / "
            f"Final={prediction['confidence']:.3f}"
        )

    if SAVE_SAMPLES:
        save_samples(predictions)
        print(f"サンプル保存: {SAMPLE_DIR}")

    result = draw_result(frame, predictions)
    cv2.imwrite(str(RESULT_IMAGE_PATH), result)
    return result


def main():
    device = select_device()
    print(f"device: {device}")

    yolo_model = YOLO(str(YOLO_MODEL_PATH))
    classifier = HierarchicalTileClassifier(CLASSIFICATION_MODEL_DIR, device)

    camera = cv2.VideoCapture(0)
    if not camera.isOpened():
        raise RuntimeError(f"カメラを開けません: {0}")

    print("Sキー: 推論")
    print("Qキー: 終了")

    try:
        while True:
            ret, frame = camera.read()
            if not ret:
                print("カメラ画像を取得できませんでした")
                break

            preview = frame.copy()
            cv2.putText(preview, "S: inference  Q: quit", (20, 35), cv2.FONT_HERSHEY_SIMPLEX, 0.8,(0, 0, 0),5,cv2.LINE_AA,)

            cv2.imshow("Camera", preview)
            key = cv2.waitKey(1) & 0xFF
            if key in (ord("q"), ord("Q")):
                break
            if key in (ord("s"), ord("S")):
                run_inference(frame.copy(), yolo_model, classifier, device)

    finally:
        camera.release()
        cv2.destroyAllWindows()


if __name__ == "__main__":
    main()