app.py

import cv2
from mmdeploy_runtime import Detector, Segmentor, Classifier
import numpy as np
import gradio as gr
import math
import os


# Load models globally to avoid redundancy

helmet_detector = Detector(model_path='/mnt/e/AI/mmdeploy/output/helmet', device_name='cuda', device_id=0)
red_tree_segmentor = Segmentor(model_path='/mnt/e/AI/mmdeploy/output/red_tree', device_name='cuda', device_id=0)
vest_detector = Detector(model_path='/mnt/e/AI/mmdeploy/output/vest_detection', device_name='cuda', device_id=0)
car_detector = Detector(model_path='/mnt/e/AI/mmdeploy/output/car_calculation', device_name='cuda', device_id=0)
crack_classifier = Classifier(model_path='/mnt/e/AI/mmdeploy/output/crack_classification', device_name='cuda', device_id=0)
disease_object_detector = Detector(model_path='/mnt/e/AI/mmdeploy/output/disease_object_detection', device_name='cuda', device_id=0)
crack_segmentor = Segmentor(model_path='/mnt/e/AI/mmdeploy/output/crack_detection2', device_name='cuda', device_id=0)
leaf_disease_segmentor = Segmentor(model_path='/mnt/e/AI/mmdeploy/output/disease_leaf', device_name='cuda', device_id=0)
single_label_disease_segmentor = Segmentor(model_path='/mnt/e/AI/mmdeploy/output/disease_detection', device_name='cuda', device_id=0)
fall_detector = Detector(model_path='/mnt/e/AI/mmdeploy/output/fall_detection_fastercnn', device_name='cuda', device_id=0)
mask_detector = Detector(model_path='/mnt/e/AI/mmdeploy/output/mask_detection', device_name='cuda', device_id=0)
smoker_detector_object = Detector(model_path='/mnt/e/AI/mmdeploy/output/smoker_nonsmoker', device_name='cuda', device_id=0)

def smoker_detector(frame, confidence_threshold=0.3):
    SMOKE_LABELS = ['smoker', 'nonsmoker']  # 新的标签列表
    bboxes, labels, masks = smoker_detector_object(frame)  # 修改检测器名字

    # 获取有效的bbox索引
    valid_indices = [(i, SMOKE_LABELS[label]) for i, label in enumerate(labels) if SMOKE_LABELS[label] == 'smoker' and bboxes[i][4] >= confidence_threshold]

    smoker_count = 0

    for i, label_name in valid_indices:
        bbox = bboxes[i]
        [left, top, right, bottom], score = bbox[0:4].astype(int), bbox[4]

        if label_name == 'smoker':
            color = (255, 0, 0)  # 绿色用于'smoker'
            smoker_count += 1
        
        line_thickness = 2
        font_scale = 0.8
        cv2.rectangle(frame, (left, top), (right, bottom), color, thickness=line_thickness)
        label_text = f"{label_name} ({score:.2f})"
        cv2.putText(frame, label_text, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, font_scale, color, line_thickness)

        if masks and masks[i].size:
            mask = masks[i]
            blue, green, red = cv2.split(frame)
            if mask.shape == frame.shape[:2]:
                mask_img = blue
            else:
                x0 = int(max(math.floor(bbox[0]) - 1, 0))
                y0 = int(max(math.floor(bbox[1]) - 1, 0))
                mask_img = blue[y0:y0 + mask.shape[0], x0:x0 + mask.shape[1]]
            cv2.bitwise_or(mask, mask_img, mask_img)
            frame = cv2.merge([blue, green, red])

    # 显示smoker的数量
    frame_height, frame_width = frame.shape[:2]
    summary_text = f"Smokers: {smoker_count}"
    cv2.putText(frame, summary_text, (frame_width - 200, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)

    return frame, smoker_count



def crack_classification(frame, confidence_threshold=0.5):
    # 定义标签
    labels_dict = {0: 'Negative', 1: 'Positive'}
    
    # 使用裂缝分类器进行预测
    result = crack_classifier(frame)
    
    # 获取最大置信度的标签ID
    label_id, score = max(result, key=lambda x: x[1])

    if label_id == 1 and score > confidence_threshold:  # 如果检测到有裂缝，并且置信度超过阈值
        seg = crack_segmentor(frame)
        crack_pixel_count = np.sum(seg == 1)
        current_palette = [(255, 255, 255), (255, 0, 0)]  # 背景为白色，裂缝为红色
        color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)
        for label, color in enumerate(current_palette):
            color_seg[seg == label, :] = color
        frame = frame * 0.5 + color_seg * 0.5
        frame = frame.astype(np.uint8)
    elif label_id == 0 and score <= confidence_threshold:
        crack_pixel_count = None
    else:
        crack_pixel_count = None
        label_id = 0  # 这里我默认设置为0，即"Negative"，但你可以根据实际情况进行调整

    # 在图像右上角显示预测结果和置信度
    label_text = labels_dict[label_id] + f" ({score:.2f})"
    color = (255, 0, 0) if label_id == 1 else (0, 255, 0)  # 裂缝为红色，否则为绿色
    font_scale = 0.8
    line_thickness = 2
    text_size = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, line_thickness)[0]
    cv2.putText(frame, label_text, (frame.shape[1] - text_size[0] - 10, text_size[1] + 10), cv2.FONT_HERSHEY_SIMPLEX, font_scale, color, line_thickness)

    return frame, labels_dict[label_id], crack_pixel_count


def crack_detection(frame):
    # 使用裂缝检测器进行检测
    seg = crack_segmentor(frame)
    crack_pixel_count = np.sum(seg == 1)

    # 如果检测到裂缝，进行可视化处理
    if crack_pixel_count > 0:
        current_palette = [(255, 255, 255), (255, 0, 0)]  # 背景为白色，裂缝为红色
        color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)
        for label, color in enumerate(current_palette):
            color_seg[seg == label, :] = color
        frame = frame * 0.5 + color_seg * 0.5
        frame = frame.astype(np.uint8)

    # 在图像右上角显示检测到的裂缝像素数量
    label_text = f"Crack Pixels: {crack_pixel_count}"
    color = (255, 0, 0) if crack_pixel_count > 0 else (0, 255, 0)  # 如果有裂缝则为红色，否则为绿色
    font_scale = 0.8
    line_thickness = 2
    text_size = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, line_thickness)[0]
    cv2.putText(frame, label_text, (frame.shape[1] - text_size[0] - 10, text_size[1] + 10), cv2.FONT_HERSHEY_SIMPLEX, font_scale, color, line_thickness)

    return frame, crack_pixel_count


def car_calculation(frame, confidence_threshold=0.7):
    CAR_LABEL = 'car'  # 这里只有一个车辆标签
    bboxes, labels, masks = car_detector(frame)
    valid_indices = [i for i, label in enumerate(labels) if bboxes[i][4] >= confidence_threshold]
    
    car_count = 0
    
    for i in valid_indices:
        bbox = bboxes[i]
        [left, top, right, bottom], score = bbox[0:4].astype(int), bbox[4]
        
        color = (0, 255, 0)  # 使用绿色标记车辆
        line_thickness = 2
        font_scale = 0.8

        cv2.rectangle(frame, (left, top), (right, bottom), color, thickness=line_thickness)
        label_text = CAR_LABEL + f" ({score:.2f})"
        cv2.putText(frame, label_text, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, font_scale, color, line_thickness)
        
        if masks and masks[i].size:
            mask = masks[i]
            blue, green, red = cv2.split(frame)
            if mask.shape == frame.shape[:2]:
                mask_img = blue
            else:
                x0 = int(max(math.floor(bbox[0]) - 1, 0))
                y0 = int(max(math.floor(bbox[1]) - 1, 0))
                mask_img = blue[y0:y0 + mask.shape[0], x0:x0 + mask.shape[1]]
            cv2.bitwise_or(mask, mask_img, mask_img)
            frame = cv2.merge([blue, green, red])
        
        car_count += 1

    return frame, car_count



def vest_detection(frame, confidence_threshold=0.3):
    VEST_LABELS = ['other_clothes', 'vest']  # 新的标签列表
    bboxes, labels, masks = vest_detector(frame)
    
    # 获取有效的bbox索引
    valid_indices = [(i, VEST_LABELS[label]) for i, label in enumerate(labels) if VEST_LABELS[label] in ['vest', 'other_clothes'] and bboxes[i][4] >= confidence_threshold]

    vest_count = 0
    other_clothes_count = 0

    for i, label_name in valid_indices:
        bbox = bboxes[i]
        [left, top, right, bottom], score = bbox[0:4].astype(int), bbox[4]

        if label_name == 'vest':
            color = (0, 255, 255)  # 黄色用于'vest'
            vest_count += 1
        else:
            color = (255, 0, 0)  # 蓝色用于'other_clothes'
            other_clothes_count += 1
        
        line_thickness = 2
        font_scale = 0.8
        cv2.rectangle(frame, (left, top), (right, bottom), color, thickness=line_thickness)
        label_text = f"{label_name} ({score:.2f})"
        cv2.putText(frame, label_text, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, font_scale, color, line_thickness)

        if masks and masks[i].size:
            mask = masks[i]
            blue, green, red = cv2.split(frame)
            if mask.shape == frame.shape[:2]:
                mask_img = blue
            else:
                x0 = int(max(math.floor(bbox[0]) - 1, 0))
                y0 = int(max(math.floor(bbox[1]) - 1, 0))
                mask_img = blue[y0:y0 + mask.shape[0], x0:x0 + mask.shape[1]]
            cv2.bitwise_or(mask, mask_img, mask_img)
            frame = cv2.merge([blue, green, red])

    # 显示vest和other_clothes的数量和置信度
    frame_height, frame_width = frame.shape[:2]
    summary_text = f"Vests: {vest_count}, Other Clothes: {other_clothes_count}"
    cv2.putText(frame, summary_text, (frame_width - 300, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)

    return frame, vest_count, other_clothes_count

def detect_falls(frame, confidence_threshold=0.5):
    # 假设输入图像是RGB格式，转换为BGR
    frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)

    LABELS = ['fall', 'person']
    # 初始化摔倒计数器
    fall_count = 0

    # 使用模型进行检测
    bboxes, labels, masks = fall_detector(frame)
    
    for bbox, label_id in zip(bboxes, labels):
        [left, top, right, bottom], score = bbox[0:4].astype(int), bbox[4]
        if score < confidence_threshold:
            continue
        if LABELS[label_id] == 'fall':  # 仅显示摔倒的标注框
            cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255), 2)
            label_text = f"{LABELS[label_id]}: {int(score*100)}%"
            cv2.putText(frame, label_text, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
            # 递增摔倒计数器
            fall_count += 1

    # 转换图像回RGB格式
    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    # 返回处理后的图像和摔倒的数量
    return frame, fall_count

def leaf_disease_detection(frame, confidence_threshold=0.3):
    # 假设输入图像是RGB格式，转换为BGR
    frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)

    LABELS = ['disease']
    # 初始化病害计数器
    disease_count = 0
    bboxes, labels, masks = disease_object_detector(frame)
    indices = [i for i in range(len(bboxes))]
    for index, bbox, label_id in zip(indices, bboxes, labels):
        [left, top, right, bottom], score = bbox[0:4].astype(int), bbox[4]
        if score < confidence_threshold:
            continue
        cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255), 1)
        label_text = f"{LABELS[label_id]}: {int(score*100)}%"
        cv2.putText(frame, label_text, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)

        if masks[index].size:
            mask = masks[index]
            blue, green, red = cv2.split(frame)
            if mask.shape == frame.shape[:2]:
                mask_img = blue
            else:
                x0 = int(max(math.floor(bbox[0]) - 1, 0))
                y0 = int(max(math.floor(bbox[1]) - 1, 0))
                mask_img = blue[y0:y0 + mask.shape[0], x0:x0 + mask.shape[1]]
            cv2.bitwise_or(mask, mask_img, mask_img)
            frame = cv2.merge([blue, green, red])
        # 递增病害计数器
        disease_count += 1

    # 转换图像回RGB格式
    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    # 返回处理后的图像、病害计数和保存的图像路径
    return frame, disease_count

def detect_masks(frame, confidence_threshold=0.5):
    # 假设输入图像是RGB格式，转换为BGR
    frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)

    LABELS = ['unfit', 'mask', 'nomask']
    # 初始化三个标签的计数器
    mask_count, nomask_count, unfit_count = 0, 0, 0

    # 使用模型进行检测
    bboxes, labels, masks = mask_detector(frame)
    
    for bbox, label_id in zip(bboxes, labels):
        [left, top, right, bottom], score = bbox[0:4].astype(int), bbox[4]
        if score < confidence_threshold:
            continue

        # 根据标签ID判断类别，并进行相应的计数
        if LABELS[label_id] == 'mask':
            mask_count += 1
            color = (0, 255, 0)
        elif LABELS[label_id] == 'nomask':
            nomask_count += 1
            color = (0, 0, 255)
        elif LABELS[label_id] == 'unfit':
            unfit_count += 1
            color = (255, 0, 0)
        
        cv2.rectangle(frame, (left, top), (right, bottom), color, 2)
        label_text = f"{LABELS[label_id]}: {int(score*100)}%"
        cv2.putText(frame, label_text, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)

    # 转换图像回RGB格式
    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    # 返回处理后的图像和每个标签的数量
    return frame, mask_count, nomask_count, unfit_count

def helmet_detection(frame, confidence_threshold=0.3):

    HEL_LABELS = ['head', 'helmet']
    bboxes, labels, masks = helmet_detector(frame)
    valid_indices = [i for i, bbox in enumerate(bboxes) if bbox[4] >= confidence_threshold]
    
    helmet_count = 0
    head_count = 0
    
    for i in valid_indices:
        bbox = bboxes[i]
        label_id = labels[i]
        [left, top, right, bottom], score = bbox[0:4].astype(int), bbox[4]
        
        if HEL_LABELS[label_id] == 'helmet':
            color = (0, 255, 0)  # Green color for 'helmet'
            line_thickness = 1
            font_scale = 0.5
        elif HEL_LABELS[label_id] == 'head':
            color = (255, 0, 0)  # Red color for 'head'
            line_thickness = 1  # Increased line thickness for 'head' boxes
            font_scale = 0.5  # Decreased font size for 'head' labels

        cv2.rectangle(frame, (left, top), (right, bottom), color, thickness=line_thickness)
        label_text = HEL_LABELS[label_id] + f" ({score:.2f})"
        cv2.putText(frame, label_text, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, font_scale, color, line_thickness)

        if HEL_LABELS[label_id] == 'helmet':
            helmet_count += 1
        elif HEL_LABELS[label_id] == 'head':
            head_count += 1

    return frame, helmet_count, head_count



def human_calculation(frame, confidence_threshold=0.3):
    """
    Process the given image to count the number of humans.
    """
    HEL_LABELS = ['head', 'helmet']
    bboxes, labels, masks = helmet_detector(frame)
    
    human_count = 0  # Initialize human count
    
    for i in range(len(bboxes)):
        bbox = bboxes[i]
        label_id = labels[i]
        score = bbox[4]
        
        # Check if the label is 'head' or 'helmet' and the score is greater than confidence_threshold
        if HEL_LABELS[label_id] in ['head', 'helmet'] and score > confidence_threshold:
            human_count += 1
            [left, top, right, bottom] = bbox[0:4].astype(int)
            cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255), thickness=1)  # Red color for boxes
            label_text = f"human ({score:.2f})"  # Include confidence score in label_text
            cv2.putText(frame, label_text, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)


    return frame, human_count


def red_tree(img):
    img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    def get_palette(num_classes=2):
        return [(255, 255, 255), (255, 0, 0)]
    seg = red_tree_segmentor(img_bgr)
    red_tree_pixel_count = np.sum(seg == 1)
    current_palette = get_palette()
    color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)
    for label, color in enumerate(current_palette):
        color_seg[seg == label, :] = color
    color_seg_bgr = color_seg[..., ::-1]

    img_bgr = img_bgr * 0.5 + color_seg_bgr * 0.5
    img_bgr = img_bgr.astype(np.uint8)
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)

    return img_rgb, red_tree_pixel_count


def leaf_disease(img):
    img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    
    def get_palette(num_classes=3):
        return [(255, 255, 255), (0, 255, 0), (255, 0, 0)]

    seg = leaf_disease_segmentor(img_bgr)
    
    leaf_pixel_count = np.sum(seg == 1)
    disease_pixel_count = np.sum(seg == 2)
    
    current_palette = get_palette()
    color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)
    
    for label, color in enumerate(current_palette):
        color_seg[seg == label, :] = color
    
    color_seg_bgr = color_seg[..., ::-1]
    img_bgr = img_bgr * 0.5 + color_seg_bgr * 0.5
    img_bgr = img_bgr.astype(np.uint8)
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    
    return img_rgb, leaf_pixel_count, disease_pixel_count

def single_label_disease(img):
    img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    
    def get_palette(num_classes=2):
        return [(255, 255, 255), (255, 0, 0)]

    seg = single_label_disease_segmentor(img_bgr)
    
    disease_pixel_count = np.sum(seg == 1)
    
    current_palette = get_palette()
    color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)
    
    for label, color in enumerate(current_palette):
        color_seg[seg == label, :] = color
    
    color_seg_bgr = color_seg[..., ::-1]
    img_bgr = img_bgr * 0.5 + color_seg_bgr * 0.5
    img_bgr = img_bgr.astype(np.uint8)
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)

    
    return img_rgb, disease_pixel_count



def get_image_examples():
    image_dir = "/mnt/e/AI/mmdeploy/gradio/photo"
    image_files = [f for f in os.listdir(image_dir) if f.endswith(('.png', '.jpg', '.jpeg'))]
    image_files.sort(key=lambda f: int(''.join(filter(str.isdigit, f))))  # 按数字排序
    example_choices = [
        '红树林识别', '红树林识别', '红树林识别', 
        '安全帽检测', '安全帽检测', '安全帽检测',
        '人数统计', '人数统计', '人数统计',
        '反光衣检测','反光衣检测','反光衣检测',
        '道路车辆统计', '道路车辆统计', '道路车辆统计',
        '裂缝识别', '裂缝识别', '裂缝识别',
        '吸烟检测','吸烟检测','吸烟检测',
        '树叶病害识别1','树叶病害识别1','树叶病害识别1',
        '树叶病害识别2','树叶病害识别2','树叶病害识别2',
        '树叶病害检测3','树叶病害检测3','树叶病害检测3',
        '摔倒检测','摔倒检测','摔倒检测',
        '口罩佩戴检测','口罩佩戴检测','口罩佩戴检测',
    ]
    
    confidence_thresholds = [
    0, 0, 0, 
    0.7, 0.8, 0.6, 
    0.3, 0.8, 0.5, 
    0.8, 0.7, 0.8, 
    0.5, 0.2, 0.7, 
    0, 0, 0, 
    0.6, 0.9, 0.5, 
    0, 0, 0, 
    0, 0, 0, 
    0.4, 0.4, 0.5, 
    0.9, 0.9, 0.5, 
    0.8, 0.6, 0.5
]
    examples = [[example_choices[i], f"{image_dir}/{image_file}", confidence_thresholds[i]] for i, image_file in enumerate(image_files)]
    return examples

  


model_choices = ['红树林识别','裂缝识别','树叶病害识别1','树叶病害识别2','树叶病害检测3', '安全帽检测','反光衣检测', '吸烟检测','摔倒检测', '口罩佩戴检测','人数统计','道路车辆统计']


def create_blank_image(width=640, height=480, color=(255, 255, 255)):
    blank_image = np.zeros((height, width, 3), np.uint8)
    blank_image[:, :] = color
    return blank_image

def process_image(model_choice, image_array=None, confidence_threshold=0.3):
    output_text = '当前未有图片输入，请上传图片后再次点击运行。'
    
    if image_array is None:
        img = create_blank_image()
    else:
        if model_choice not in model_choices:
            model_choice = "安全帽检测"
        # 以下是模型选择和执行逻辑
        if model_choice == "红树林识别":
            img, red_tree_pixel_count = red_tree(image_array)  # 语义分割模型
            output_text = f"红树林的像素点有 {red_tree_pixel_count} 个。"
        elif model_choice == "安全帽检测":
            img, helmet_count, head_count = helmet_detection(image_array, confidence_threshold)
            output_text = f"佩戴安全帽的人数为：{helmet_count}，未佩戴安全帽的人数为：{head_count}。"
        elif model_choice == "人数统计":
            img, human_count = human_calculation(image_array, confidence_threshold)
            output_text = f"该图片人员总人数为: {human_count}。"
        elif model_choice == "反光衣检测":
            img, vest_count, other_clothes_count= vest_detection(image_array, confidence_threshold)
            output_text = f"该图片中总有 {vest_count} 人配备了反光衣，{other_clothes_count} 人没有配备反光衣。"
        elif model_choice == "道路车辆统计":
            img, car_count = car_calculation(image_array, confidence_threshold)
            output_text = f"该道路上目前共有 {car_count} 台车辆。"
        elif model_choice == "裂缝识别":
            img, crack_result, crack_pixel_count = crack_classification(image_array, confidence_threshold)
            if crack_result == "Positive":
                output_text = f"该图片内存在裂缝，裂缝的像素点有 {crack_pixel_count} 个。"
            else:
                output_text = "该图片不存在裂缝。"
        elif model_choice == "树叶病害检测3":  
            img, disease_count = leaf_disease_detection(image_array, confidence_threshold)
            if disease_count > 0:
                output_text = f"共检测到 {disease_count} 处病害。"
            else:
                output_text = "并未检测到病害。"
        elif model_choice == "吸烟检测":
            img, smoker_count = smoker_detector(image_array, confidence_threshold)
            output_text = f"当前图片有 {smoker_count} 人在吸烟。"
        elif model_choice == "树叶病害识别1":
            img, leaf_pixel_count, disease_pixel_count = leaf_disease(image_array)  # 语义分割模型
            if disease_pixel_count == 0:
                output_text = "该树叶并未出现病害。"
            else:
                output_text = f"病害的像素点有 {disease_pixel_count} 个。"
        elif model_choice == "树叶病害识别2":  
            img, disease_pixel_count = single_label_disease(image_array)  # 语义分割模型
            output_text = f"病害的像素点有 {disease_pixel_count} 个。"
        elif model_choice == "摔倒检测":  # 您可以根据实际情况调整模型选择的名称
            img, fall_count = detect_falls(image_array,confidence_threshold)
            output_text = f"图像中摔倒的人数为 {fall_count} 人。"
        elif model_choice == "口罩佩戴检测":  # 您可以根据实际情况调整模型选择的名称
            img, mask_count, nomask_count, unfit_count = detect_masks(image_array,confidence_threshold)
            output_text = f"当前佩戴口罩的人数为 {mask_count}，未正确佩戴口罩的人数为 {unfit_count}，没有佩戴口罩的人数为 {nomask_count}。"
            
    return img, output_text

def process_video(model_choice, video=None, confidence_threshold=0.3):

    # 内部函数：创建空白视频
    def create_blank_video(filename, duration=5, fps=30, width=640, height=480, color=(255, 255, 255)):
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 使用mp4v编解码器
        out = cv2.VideoWriter(filename, fourcc, fps, (width, height))
        blank_image = np.zeros((height, width, 3), np.uint8)
        blank_image[:, :] = color
        for _ in range(int(fps * duration)):
            out.write(blank_image)
        out.release()

    # 检查视频是否存在
    if video is None:
        video_output_path = '/mnt/e/AI/mmdeploy/gradio/video/none.mp4'
        create_blank_video(video_output_path)
        output_text2 = '当前未有视频输入，请上传视频后再次点击运行。'
        return video_output_path, output_text2
    else:
        video_output_path = '/mnt/e/AI/mmdeploy/gradio/video/output_video.mp4'
        cap = cv2.VideoCapture(video)
        if not cap.isOpened():
            raise ValueError("无法打开视频文件")
        fps = int(cap.get(cv2.CAP_PROP_FPS))
        num_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        # 获取输入视频的分辨率
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        # 使用帧采样的逻辑，但考虑到所有帧都需要处理，我们使用间隔为1的采样。
        clip_len, frame_interval, num_clips = 1, 1, num_frames
        avg_interval = (num_frames - clip_len * frame_interval + 1) / float(num_clips)
        frame_inds = []
        for i in range(num_clips):
            clip_offset = int(i * avg_interval + avg_interval / 2.0)
            for j in range(clip_len):
                ind = (j * frame_interval + clip_offset) % num_frames
                if num_frames <= clip_len * frame_interval - 1:
                    ind = j % num_frames
                frame_inds.append(ind)

        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        processed_frames = []
        for frame_id in sorted(frame_inds): 
            cap.set(cv2.CAP_PROP_POS_FRAMES, frame_id)  # 设置读取特定的帧
            ret, frame = cap.read()
            if not ret:
                break
            # 将帧率添加到视频的左上角
            cv2.putText(frame, "FPS: {}".format(fps), (10, 30), 
                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2, cv2.LINE_AA)
            
            if model_choice == "红树林识别":
                # 在此处调用红树林模型处理帧
                processed_frame, red_tree_pixel_count = red_tree(frame)
                    # 在处理后的帧的右上角添加文字
                cv2.putText(processed_frame, "Number of pixels in this frame: {}".format(red_tree_pixel_count), 
                            (processed_frame.shape[1] - 300, 30), cv2.FONT_HERSHEY_SIMPLEX, 
                            0.7, (255, 255, 255), 2)


            elif model_choice == "安全帽检测":
                # 在此处调用安全帽检测模型处理帧
                processed_frame, helmet_count, head_count = helmet_detection(frame, confidence_threshold)
                
                cv2.putText(processed_frame, "Number of people wearing helmets: {}".format(helmet_count), 
                (processed_frame.shape[1] - 400, 30), cv2.FONT_HERSHEY_SIMPLEX, 
                0.7, (255, 255, 255), 2)
    
                # 在上一行文字下方添加表示未佩戴安全帽的人数的文字
                cv2.putText(processed_frame, "Number of people without helmets: {}".format(head_count - helmet_count), 
                            (processed_frame.shape[1] - 450, 60), cv2.FONT_HERSHEY_SIMPLEX, 
                            0.7, (255, 255, 255), 2)


            elif model_choice == "人数统计":
                # 在此处调用人数统计模型处理帧
                processed_frame, human_count = human_calculation(frame, confidence_threshold)
                cv2.putText(processed_frame, "Current number of people: {}".format(human_count), 
                (processed_frame.shape[1] - 300, 30), cv2.FONT_HERSHEY_SIMPLEX, 
                0.7, (255, 255, 255), 2)

            elif model_choice == "反光衣检测":
                # 在此处调用反光衣检测模型处理帧
                processed_frame, vest_count, other_clothes_count= vest_detection(image_array, confidence_threshold)
                cv2.putText(processed_frame, "Number of reflective vests: {}".format(vest_count), 
                (processed_frame.shape[1] - 350, 30), cv2.FONT_HERSHEY_SIMPLEX, 
                0.7, (255, 255, 255), 2)
                cv2.putText(processed_frame, "Number without reflective vests: {}".format(other_clothes_count), 
                        (processed_frame.shape[1] - 450, 60), cv2.FONT_HERSHEY_SIMPLEX, 
                        0.7, (255, 255, 255), 2)

            elif model_choice == "道路车辆统计":
                # 在此处调用道路车辆统计模型处理帧
                processed_frame, car_count = car_calculation(frame, confidence_threshold)
                cv2.putText(processed_frame, "Number of vehicles: {}".format(car_count), 
                (processed_frame.shape[1] - 250, 30), cv2.FONT_HERSHEY_SIMPLEX, 
                0.7, (255, 255, 255), 2)               

            elif model_choice == "裂缝识别":
                # 在此处调用裂缝识别模型处理帧
                processed_frame, crack_pixel_count= crack_detection(frame)

            elif model_choice == "树叶病害检测3":
                # 在此处调用树叶病害检测模型处理帧
                processed_frame, disease_count= leaf_disease_detection(frame, confidence_threshold)
                # 在图像右上角显示叶片的病害数量
                label_text = f"Leaf Disease Count: {disease_count}"
                color = (0, 0, 255)  # 红色
                font_scale = 0.8
                line_thickness = 2
                text_size = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, line_thickness)[0]
                cv2.putText(processed_frame, label_text, (processed_frame.shape[1] - text_size[0] - 10, text_size[1] + 10), cv2.FONT_HERSHEY_SIMPLEX, font_scale, color, line_thickness)



            elif model_choice == "吸烟检测":
                # 在此处调用吸烟检测模型处理帧
                processed_frame, smoker_count = smoker_detector(frame, confidence_threshold)

                # 准备要显示的文本
                text = f"吸烟者数量: {smoker_count}"

                # 获取文本大小
                text_size = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]

                # 计算文本的位置，以便它出现在帧的右上角
                text_position = (processed_frame.shape[1] - text_size[0] - 10, text_size[1] + 10)

                # 将文本绘制到帧上
                cv2.putText(processed_frame, text, text_position, cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)
                
            elif model_choice == "树叶病害识别1":
                # 在此处调用树叶病害识别模型处理帧
                processed_frame, _, disease_pixel_count = leaf_disease(frame)
                text = f"Current disease pixel count on the leaf: {disease_pixel_count}"
                text_size = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]
                cv2.putText(processed_frame, text, 
                            (processed_frame.shape[1] - text_size[0] - 10, text_size[1] + 10), 
                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)


            elif model_choice == "树叶病害识别2":  
                # 在此处调用树叶病害识别模型处理帧
                processed_frame, disease_pixel_count= single_label_disease(frame)
                text = f"Current disease pixel count on the leaf: {disease_pixel_count}"
                text_size = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]
                cv2.putText(processed_frame, text, 
                            (processed_frame.shape[1] - text_size[0] - 10, text_size[1] + 10), 
                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)
                
            elif model_choice == "口罩佩戴检测":  # 您可以根据实际情况调整模型选择的名称
                processed_frame, mask_count, nomask_count, unfit_count = detect_masks(frame,confidence_threshold)
                cv2.putText(processed_frame, "Number wearing masks: {}".format(mask_count), 
                            (processed_frame.shape[1] - 350, 30), cv2.FONT_HERSHEY_SIMPLEX, 
                            0.7, (255, 255, 255), 2)

                cv2.putText(processed_frame, "Number not wearing masks: {}".format(nomask_count), 
                            (processed_frame.shape[1] - 400, 60), cv2.FONT_HERSHEY_SIMPLEX, 
                            0.7, (255, 255, 255), 2)

                cv2.putText(processed_frame, "Number wearing masks incorrectly: {}".format(unfit_count), 
                            (processed_frame.shape[1] - 500, 90), cv2.FONT_HERSHEY_SIMPLEX, 
                            0.7, (255, 255, 255), 2)
                
            elif model_choice == "摔倒检测":  
                
                processed_frame, fall_count= detect_falls(frame,confidence_threshold)
                cv2.putText(processed_frame, "Number of people who fell: {}".format(fall_count), 
            (processed_frame.shape[1] - 350, 30), cv2.FONT_HERSHEY_SIMPLEX, 
            0.7, (255, 255, 255), 2)

            processed_frames.append(processed_frame)
        out = cv2.VideoWriter(video_output_path, fourcc, fps, (width,height))
        for frame in processed_frames:
            out.write(frame)
        out.release()
        cap.release()
        output_text2 = '请点击蓝色按钮下载视频。'
    return video_output_path, output_text2

with gr.Blocks() as demo:
    gr.Markdown("# <center>启云科技AI识别示例样板V1.12</center>")
    gr.Markdown("请上传图像或视频进行预测")
    with gr.Tab("AI图像处理"):
        with gr.Row():
            image_input2 = gr.Image(label="上传图像", type="numpy")
            with gr.Column():
                image_input1 = gr.Dropdown(choices=model_choices, label="模型选择")
                image_input3 = gr.Slider(minimum=0, maximum=1, step=0.1, label="置信度阈值")
        with gr.Row():
            image_output1 = gr.Image(label="处理后的图像", type="numpy")
            with gr.Column():
                image_output2 = gr.Textbox(label="图像输出信息")
        image_button = gr.Button('请点击按钮进行图像预测')
        gr.Examples(get_image_examples(),inputs=[image_input1, image_input2, image_input3],outputs=[image_output1, image_output2], fn=process_image ,examples_per_page=6 ,cache_examples=True)
    with gr.Tab("AI视频处理"):

        with gr.Row():
            video_input2 = gr.Video(label = '上传视频', format='mp4',interactive = True)
            with gr.Column():
                video_input1 = gr.Dropdown(choices=model_choices, label="模型选择")
                video_input3 = gr.Slider(minimum=0, maximum=1, 
                                         step=0.1, label="置信度阈值")
        with gr.Row():
            video_output1 = gr.File(label='处理后的视频', type='file')
            with gr.Column():
                video_output2 = gr.Textbox(label = '视频输出信息')   
        video_button = gr.Button('请点击按钮进行视频预测')
    with gr.Accordion("平台简介"):
        gr.Markdown("红树林识别模型、裂缝识别模型、树叶病害识别模型、安全帽检测模型、反光衣检测模型、吸烟检测模型、口罩佩戴检测、摔倒检测、人数统计模型及道路车辆统计模型展示平台。")
    image_button.click(process_image, inputs = [image_input1, image_input2, image_input3], outputs=[image_output1, image_output2])
    video_button.click(process_video, inputs=[video_input1,video_input2, video_input3], outputs=[video_output1, video_output2])

demo.launch(share=True)