from transformers import pipeline
from PIL import Image, ImageDraw
import os

# 🔸 냉장고 이미지 파일 이름
image_path = "image.jpg"

# 🔸 탐지할 재료 라벨 (영문)
candidate_labels = [
    "salmon", "chicken breast", "broccoli", "lettuce", "mushroom",
    "bell pepper", "onion", "cherry tomato"
]


# 🔹 이미지 불러오기
if not os.path.exists(image_path):
    raise FileNotFoundError(f"{image_path} 파일이 존재하지 않습니다.")
image = Image.open(image_path).convert("RGB")

# 🔹 zero-shot 객체 탐지 파이프라인 만들기
detector = pipeline(
    task="zero-shot-object-detection",
    model="google/owlv2-base-patch16-ensemble"
)

# 🔹 탐지 실행
outputs = detector(image, candidate_labels=candidate_labels, threshold=0.2)

# 🔹 결과 출력
print("\n[탐지 결과]")
for pred in outputs:
    label = pred["label"]
    score = pred["score"]
    box = pred["box"]
    print(f"{label} ({score:.2%}) 위치: {box}")

# 🔹 결과 이미지에 박스 그리기
draw = ImageDraw.Draw(image)
for pred in outputs:
    box = pred["box"]
    label = pred["label"]
    score = pred["score"]
    draw.rectangle(
        [box["xmin"], box["ymin"], box["xmax"], box["ymax"]],
        outline="red", width=2
    )
    draw.text((box["xmin"], box["ymin"]), f"{label} {score:.2%}", fill="red")

# 🔹 시각화
image.show()