wickedreg's picture
Update app.py
81d99cc verified
raw
history blame
3.45 kB
import numpy as np
import pandas as pd
import gradio as gr
from PIL import Image, ImageDraw, ImageFont
from io import BytesIO
import json
import cv2
from ultralytics import YOLO
# ======================= МОДЕЛЬ ===================================
model = YOLO("yolov11m_best.pt")
# ================== ЧТЕНИЕ НАЗВАНИЙ И ЦЕН =======================
with open('Fruit_Veggies_Price.json', 'r', encoding='utf-8') as file:
fruits_data = json.load(file)
# =========================== ДЕТЕКЦИЯ ПЛОДА ============================
def detect_fruit(image):
# считываем изображение
# предполагается, что image - это объект PIL Image
image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
# детекция
detections = model.predict(source=image_cv, conf=0.5)
# проверка на наличие детекций
if len(detections) == 0:
return image, None
result_np_image = detections[0].plot()
result_np_image = cv2.cvtColor(result_np_image, cv2.COLOR_BGR2RGB)
detected_fruit = None
for det in detections:
label = model.names[int(det.boxes.cls[0])] # название фрукта
if label in fruits_data:
detected_fruit = label
break
return result_np_image, detected_fruit
# =========================== ЧЕК ============================
def create_receipt(detected_fruit, weight):
data = fruits_data[detected_fruit]
fruit_name = data['name']
price = data['price_per_kg']
total_price = round(price * weight, 2)
receipt_img = Image.new("RGB", (300, 200), color="white")
draw = ImageDraw.Draw(receipt_img)
try:
font = ImageFont.truetype("DejaVuSans.ttf", 18)
except IOError:
font = ImageFont.load_default()
draw.text((10, 10), "Чек", fill="black", font=font)
draw.text((10, 50), f"Продукт: {fruit_name}", fill="black", font=font)
draw.text((10, 80), f"Вес: {weight} кг", fill="black", font=font)
draw.text((10, 110), f"Цена за кг: {price} руб.", fill="black", font=font)
draw.text((10, 140), f"Сумма: {total_price} руб.", fill="black", font=font)
return receipt_img
# ======================= ИНТЕРФЕЙС ============================
def gradio_interface(image, weight):
if weight <= 0:
gr.Info('Укажите вес товара')
return None, None
result_np_image, detected_fruit = detect_fruit(image)
if not detected_fruit:
gr.Info('Не удалось определить товар')
return image, None
receipt = create_receipt(detected_fruit, weight)
return result_np_image, receipt
image_input = gr.Image(
label="Изображение",
width=640,
height=380
)
weight_input = gr.Number(label="Вес (кг)")
image_output = gr.Image(
label="Распознанный товар",
type="numpy",
width=640,
height=380
)
receipt_output = gr.Image(
label="Чек",
type="pil",
width=400,
height=400
)
gr.Interface(
fn=gradio_interface,
inputs=[image_input, weight_input],
outputs=[image_output, receipt_output],
title="Определение товара и создание чека",
description="Загрузите изображение, введите вес и получите чек"
).launch()