wickedreg's picture
Update app.py
df374ff verified
raw
history blame
3.91 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 pathlib import Path
from ultralytics import YOLO
# ======================= МОДЕЛЬ ===================================
model_name = "yolov11m_best.pt"
model_path = Path(__file__).with_name(model_name)
model = YOLO(model_path)
# ================== ЧТЕНИЕ НАЗВАНИЙ И ЦЕН =======================
# with open('Fruit_Veggies_Price.json', 'r', encoding='utf-8') as file:
# fruits_data = json.load(file)
# =========================== ДЕТЕКЦИЯ ПЛОДА ============================
def detect_fruit(image, weight: float):
# Считываем изображение
# Предполагается, что 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
result_np_image = detections[0].plot()
result_np_image = cv2.cvtColor(result_np_image, cv2.COLOR_BGR2RGB)
# detected_fruit = None
# for det in detections.xyxy[0]: # Предполагается, что вы используете ультралайтики и доступ к детекциям
# label = model.names[int(det[5])] # название фрукта
# if label in fruits_data:
# detected_fruit = label
# break
# if not detected_fruit:
# return result_np_image, None, None, None
# fruit_name = fruits_data[detected_fruit]["name"]
# price_per_kg = fruits_data[detected_fruit]["price_per_kg"]
# total_price = round(price_per_kg * weight, 2)
return result_np_image
# return result_np_image, fruit_name, weight, total_price
# =========================== ЧЕК ============================
# def create_receipt(fruit_name, weight, total_price):
# receipt_img = Image.new("RGB", (300, 200), color="white")
# draw = ImageDraw.Draw(receipt_img)
# try:
# font = ImageFont.truetype("arial.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"Цена за кг: {fruits_data[fruit_name]['price_per_kg']} руб.", fill="black", font=font)
# draw.text((10, 140), f"Сумма: {total_price} руб.", fill="black", font=font)
# with BytesIO() as output:
# receipt_img.save(output, format="PNG")
# output.seek(0)
# return output.getvalue()
# ======================= ИНТЕРФЕЙС ============================
def gradio_interface(image):
# if weight <= 0:
# return image
detected_image = detect_fruit(image)
# if not fruit_name:
# return image
# return image, None # Вернуть пустой чек
return detected_image
# receipt = create_receipt(fruit_name, weight, total_price)
# return image, receipt
image_input = gr.Image(label="Изображение")
# weight_input = gr.Number(label="Вес (кг)")
image_output = gr.Image(label="Распознанный фрукт", type="numpy")
# receipt_output = gr.Image(label="Чек", type="numpy")
gr.Interface(
fn=gradio_interface,
inputs=[image_input],
outputs=[image_output],
# outputs=[image_output, receipt_output],
title="Определение плода и создание чека",
description="Загрузите изображение, введите вес и получите чек"
).launch()