Update app.py

app.py

@@ -4,37 +4,35 @@ 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)
 
-
-# таблица с данными
-detections = []
-
-
 # =========================== ДЕТЕКЦИЯ ПЛОДА ============================
-
-def detect_image(image_path: str, weight: float):
+def detect_fruit(image, weight: float):
     # Считываем изображение
-    image = cv2.imread(image_path)
-
+    # Предполагается, что image - это объект PIL Image
+    image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+    
     # Выполняем детекцию
-    detections = model.predict(source=image_path, conf=0.5)
+    detections = model.predict(source=image_cv, conf=0.5)
+
+    # Проверка на наличие детекций
+    if len(detections) == 0:
+        return image, "Фрукт не обнаружен", None, 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:
+    for det in detections.xyxy[0]:  # Предполагается, что вы используете ультралайтики и доступ к детекциям
         label = model.names[int(det[5])]  # название фрукта
         if label in fruits_data:
             detected_fruit = label
@@ -49,34 +47,32 @@ def detect_image(image_path: str, weight: float):
     
     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[detected_fruit]['price_per_kg']} руб.", 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)
-    
-    return receipt_img
 
+    with BytesIO() as output:
+        receipt_img.save(output, format="PNG")
+        output.seek(0)
+        return output.getvalue()
 
 # ======================= ИНТЕРФЕЙС ============================
-
 def gradio_interface(image, weight):
+    if weight <= 0:
+        return image, "Вес должен быть больше 0", None
+
     image, fruit_name, weight, total_price = detect_fruit(image, weight)
     if not fruit_name:
         return image, None  # Вернуть пустой чек