app.py

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()