cook_recommend / detect_text.py
jinu0804's picture
Initialize project without images
dfcc4e2
Raw
History Blame Contribute Delete
1.88 kB
import os
from transformers import pipeline
from PIL import Image
import gradio as gr
# ๐Ÿ”น ๊ฐ์ฒด ํƒ์ง€ ๋ชจ๋ธ ๋กœ๋”ฉ
detector = pipeline(
task="zero-shot-object-detection",
model="google/owlv2-base-patch16-ensemble"
)
# ๐Ÿ”น ํƒ์ง€ํ•  ์žฌ๋ฃŒ ๋ผ๋ฒจ ๋ฆฌ์ŠคํŠธ
candidate_labels = [
"salmon", "chicken breast", "broccoli", "lettuce", "mushroom",
"bell pepper", "onion", "cherry tomato", "egg", "milk", "cheese", "garlic"
]
# ๐Ÿ”น ์˜์–ด โ†’ ํ•œ๊ธ€ ๋ผ๋ฒจ ๋ณ€ํ™˜ ๋”•์…”๋„ˆ๋ฆฌ
label_ko = {
"salmon": "์—ฐ์–ด",
"chicken breast": "๋‹ญ๊ฐ€์Šด์‚ด",
"broccoli": "๋ธŒ๋กœ์ฝœ๋ฆฌ",
"lettuce": "์ƒ์ถ”",
"mushroom": "๋ฒ„์„ฏ",
"bell pepper": "ํ”ผ๋ง",
"onion": "์–‘ํŒŒ",
"cherry tomato": "๋ฐฉ์šธํ† ๋งˆํ† ",
"egg": "๊ณ„๋ž€",
"milk": "์šฐ์œ ",
"cheese": "์น˜์ฆˆ",
"garlic": "๋งˆ๋Š˜"
}
# ๐Ÿ”น ์žฌ๋ฃŒ ํƒ์ง€ ํ•จ์ˆ˜
def detect_ingredients(image):
if image is None:
return "์ด๋ฏธ์ง€๋ฅผ ๋จผ์ € ์—…๋กœ๋“œํ•ด์ฃผ์„ธ์š”."
outputs = detector(image, candidate_labels=candidate_labels, threshold=0.2)
detected_labels = list(set([o["label"] for o in outputs]))
if not detected_labels:
return "์žฌ๋ฃŒ๋ฅผ ์ธ์‹ํ•˜์ง€ ๋ชปํ–ˆ์Šต๋‹ˆ๋‹ค. ๋” ๋ช…ํ™•ํ•œ ์ด๋ฏธ์ง€๋ฅผ ์‚ฌ์šฉํ•ด์ฃผ์„ธ์š”."
# ์˜์–ด โ†’ ํ•œ๊ธ€ ๋ณ€ํ™˜
translated_labels = [label_ko.get(label, label) for label in detected_labels]
return ", ".join(translated_labels)
# ๐Ÿ”น Gradio UI ๊ตฌ์„ฑ
with gr.Blocks() as demo:
gr.Markdown("## ๐ŸงŠ ๋ƒ‰์žฅ๊ณ  ์žฌ๋ฃŒ ํƒ์ง€ ์‹œ์Šคํ…œ")
image_input = gr.Image(type="pil", label="๋ƒ‰์žฅ๊ณ  ์ด๋ฏธ์ง€ ์—…๋กœ๋“œ")
detect_button = gr.Button("์žฌ๋ฃŒ ํƒ์ง€ํ•˜๊ธฐ")
ingredient_output = gr.Textbox(label="๐Ÿ“Œ ํƒ์ง€๋œ ์žฌ๋ฃŒ (ํ•œ๊ธ€)")
detect_button.click(fn=detect_ingredients, inputs=image_input, outputs=ingredient_output)
if __name__ == "__main__":
demo.launch()