jelly_classification.py

import os
from ultralytics import YOLO

def is_jellyfish(image_path):
    model_path = os.path.join(os.path.dirname(__file__), "models", "image_predict.pt")

    model = YOLO(model_path)

    results = model.predict(source=image_path, save=False)

    max_confidence = 0  # 해파리 신뢰도 중 최대값 저장
    detected_jellyfish = False  # 해파리 탐지 여부

    for result in results:
        for conf, cls in zip(result.boxes.conf, result.boxes.cls):
            class_name = result.names[int(cls)]  # 클래스명 가져오기
            confidence_percent = conf.item() * 100  # 신뢰도를 백분율로 변환

            if class_name.lower() == "jellyfish":  # 클래스명이 해파리인지 확인
                detected_jellyfish = True
                max_confidence = max(max_confidence, confidence_percent)  # 가장 높은 신뢰도 업데이트

    print(f"최대 해파리 신뢰도: {max_confidence:.2f}%")
    return 1 if detected_jellyfish else 0  # 해파리 여부 반환

if __name__ == '__main__':
    image_path = 'test_image.jpg'  # 테스트할 이미지 경로
    result = is_jellyfish(image_path)
    print(f"해파리 여부: {result}")  # 1이면 해파리, 0이면 아님