app.py

import gradio as gr
import cv2
import os
import numpy as np
from ultralytics import YOLO


model = YOLO('best.pt')
# path  = [['lb0_1.jpg'],['lb1_1.jpg'],['lb2_1.jpg'],['lb3_1.jpg'],['lb4_1.jpg'],['lb5_1.jpg'],['lb6_1.jpg'],['lb7_1.jpg']]
path = []
video_path = []

label_colors = {
    0: (255, 0, 0),  # Màu đỏ cho nhãn 0
    1: (0, 255, 0),  # Màu xanh lá cho nhãn 1
    2: (0, 0, 255),  # Màu xanh dương cho nhãn 2
    3: (255, 255, 0),  # Màu vàng cho nhãn 3
    4: (255, 0, 255),  # Màu hồng cho nhãn 4
    5: (0, 255, 255),  # Màu cyan cho nhãn 5
    6: (128, 0, 128),  # Màu tím cho nhãn 6
    7: (0, 128, 128)   # Màu xanh dương đậm cho nhãn 7
}

def show_preds_image(image_path,threshold):
    image = cv2.imread(image_path)
    image_copy = image.copy()
    image_height, image_width = image.shape[:2]
    results = model(image)[0]
    
    label_counts = {}  # Đếm số lượng nhãn

    for result in results.boxes.data.tolist():
        x1, y1, x2, y2, score, class_id = result

        if score > threshold:
            color = label_colors.get(int(class_id), (255, 255, 255))  # Lấy màu từ bảng màu
            # cv2.rectangle(image_copy, (int(x1), int(y1)), (int(x2), int(y2)), color, 4)
            # cv2.putText(image_copy, results.names[int(class_id)].upper(), (int(x1), int(y1 - 10)),
            #             cv2.FONT_HERSHEY_SIMPLEX, 1.3, color, 3, cv2.LINE_AA)
            # cv2.putText(image_copy, str(score), (int(x1), int(y2 + 10)),
            #             cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 0, 255), 3, cv2.LINE_AA)
            cv2.rectangle(image_copy, (int(x1), int(y1)), (int(x2), int(y2)), color, 4)
            cv2.putText(image_copy, f"{results.names[int(class_id)].upper()} {score:.3f}", (int(x1), int(y1 - 10)),
                         cv2.FONT_HERSHEY_SIMPLEX, 1.3, color, 3, cv2.LINE_AA)
            # Đếm số lượng nhãn
            if class_id in label_counts:
                label_counts[class_id] += 1
            else:
                label_counts[class_id] = 1
    
    # Tạo hình ảnh với danh sách nhãn và màu sắc
    label_image = create_label_image(label_counts, results.names, image_width)
    
    combined_image = cv2.vconcat([cv2.cvtColor(image_copy, cv2.COLOR_BGR2RGB), label_image])
    return combined_image
# def show_preds_image(image_path):
#     image = cv2.imread(image_path)
#     image_copy = image.copy()
#     threshold = 0.1
#     results = model(image)[0]
#     for result in results.boxes.data.tolist():
#         x1, y1, x2, y2, score, class_id = result

#         if score > threshold:
#             cv2.rectangle(image_copy, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 4)
#             cv2.putText(image_copy, f"{results.names[int(class_id)].upper()} {score:.2f}", (int(x1), int(y1 - 10)),
#                         cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 255, 0), 3, cv2.LINE_AA)  # Hiển thị tên và điểm số
#     return cv2.cvtColor(image_copy, cv2.COLOR_BGR2RGB)
# Hàm tạo hình ảnh danh sách nhãn
def create_label_image(label_counts, names, image_width):
    line_height = 30
    height = line_height * (len(label_counts) + 1)
    label_image = 255 * np.ones(shape=[height, image_width, 3], dtype=np.uint8)  # Kích thước phù hợp với ảnh đầu vào
    
    y = line_height
    for class_id, count in label_counts.items():
        color = label_colors.get(int(class_id), (255, 255, 255))
        cv2.putText(label_image, f"{names[int(class_id)].upper()}: {count}", (10, y), 
                    cv2.FONT_HERSHEY_SIMPLEX, 1.3, color, 2, cv2.LINE_AA)
        y += line_height
    
    return label_image

inputs_image = [
    gr.components.Image(type="filepath", label="Input Image"),
    gr.components.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.1, label="Threshold")
]
outputs_image = [
    gr.components.Image(type="numpy", label="Output Image"),
]
interface_image = gr.Interface(
    fn=show_preds_image,
    inputs=inputs_image,
    outputs=outputs_image,
    title="Demo Wood",
    examples=path,
    cache_examples=False,
)

def show_preds_video(video_path):
    cap = cv2.VideoCapture(video_path)
    while(cap.isOpened()):
        ret, frame = cap.read()
        if ret:
            threshold = 0.1
            frame_copy = frame.copy()
            results = model(frame)[0]
            for result in results.boxes.data.tolist():
                x1, y1, x2, y2, score, class_id = result

                if score > threshold:
                    cv2.rectangle(frame_copy, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 4)
                    cv2.putText(frame_copy, results.names[int(class_id)].upper(), (int(x1), int(y1 - 10)),
                                cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 255, 0), 3, cv2.LINE_AA)
                    cv2.putText(frame_copy, str(score), (int(x1), int(y2 + 10)),
                                cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 0, 255), 3, cv2.LINE_AA)
            
            yield cv2.cvtColor(frame_copy, cv2.COLOR_BGR2RGB)
 
inputs_video = [
    gr.components.Video(type="filepath", label="Input Video"),
 
]
outputs_video = [
    gr.components.Image(type="numpy", label="Output Image"),
]
interface_video = gr.Interface(
    fn=show_preds_video,
    inputs=inputs_video,
    outputs=outputs_video,
    title="Demo Wood",
    examples=video_path,
    cache_examples=False,
)

gr.TabbedInterface(
    [interface_image, interface_video],
    tab_names=['Image inference', 'Video inference']
).queue().launch()