Update app.py

app.py

@@ -30,78 +30,78 @@ def detect_fruit(image):
     result_np_image = detections[0].plot()
     result_np_image = cv2.cvtColor(result_np_image, cv2.COLOR_BGR2RGB)
 
-    # detected_fruit = None
+    detected_fruit = None
     
-    # for det in detections.xyxy[0]:  # Предполагается, что вы используете ультралайтики и доступ к детекциям
-    #     label = model.names[int(det[5])]  # название фрукта
-    #     if label in fruits_data:
-    #         detected_fruit = label
-    #         break
+    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
+    return result_np_image, detected_fruit
 
 # =========================== ЧЕК ============================
-# def create_receipt(fruit_name, weight, total_price):
-#     receipt_img = Image.new("RGB", (300, 200), color="white")
-#     draw = ImageDraw.Draw(receipt_img)
+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("arial.ttf", 18)
-#     except IOError:
-#         font = ImageFont.load_default()
+    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)
+    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)
 
-#     with BytesIO() as output:
-#         receipt_img.save(output, format="PNG")
-#         output.seek(0)
-#         return output.getvalue()
+    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
+def gradio_interface(image, weight):
+    if weight <= 0:
+        gr.Info('Укажите вес товара')
+        return None, None
 
-    detected_image = detect_fruit(image)
-    # if not fruit_name:
-    #     return image
-        # return image, None  # Вернуть пустой чек
+    result_np_image, detected_fruit = detect_fruit(image)
+    if not fruit_name:
+        gr.Info('Не удалось определить товар')
+        return result_np_image, None
         
-    return detected_image
-    # receipt = create_receipt(fruit_name, weight, total_price)
-    # return image, receipt
+    receipt = create_receipt(detected_fruit, weight)
+    return result_np_image, receipt
 
 image_input = gr.Image(
     label="Изображение",
     width=640,
     height=480
 )
-# weight_input = gr.Number(label="Вес (кг)")
+weight_input = gr.Number(label="Вес (кг)")
 image_output = gr.Image(
     label="Распознанный фрукт",
     type="numpy",
     width=640,
     height=480
 )
-# receipt_output = gr.Image(label="Чек", type="numpy")
+receipt_output = gr.Image(
+    label="Чек",
+    type="numpy",
+    width=640,
+    height=480
+)
 
 gr.Interface(
     fn=gradio_interface,
-    inputs=[image_input],
-    outputs=[image_output],
-    # outputs=[image_output, receipt_output],
+    inputs=[image_input, weight_input],
+    outputs=[image_output, receipt_output],
     title="Определение плода и создание чека",
     description="Загрузите изображение, введите вес и получите чек"
 ).launch()