Jianfeng777/Computer_Vision_AI_Platform
1
1import cv22from mmdeploy_runtime import Detector, Segmentor, Classifier3import numpy as np4import gradio as gr5import math6import os7 8 9# Load models globally to avoid redundancy10 11helmet_detector = Detector(model_path='/mnt/e/AI/mmdeploy/output/helmet', device_name='cuda', device_id=0)12red_tree_segmentor = Segmentor(model_path='/mnt/e/AI/mmdeploy/output/red_tree', device_name='cuda', device_id=0)13vest_detector = Detector(model_path='/mnt/e/AI/mmdeploy/output/vest_detection', device_name='cuda', device_id=0)14car_detector = Detector(model_path='/mnt/e/AI/mmdeploy/output/car_calculation', device_name='cuda', device_id=0)15crack_classifier = Classifier(model_path='/mnt/e/AI/mmdeploy/output/crack_classification', device_name='cuda', device_id=0)16disease_object_detector = Detector(model_path='/mnt/e/AI/mmdeploy/output/disease_object_detection', device_name='cuda', device_id=0)17crack_segmentor = Segmentor(model_path='/mnt/e/AI/mmdeploy/output/crack_detection2', device_name='cuda', device_id=0)18leaf_disease_segmentor = Segmentor(model_path='/mnt/e/AI/mmdeploy/output/disease_leaf', device_name='cuda', device_id=0)19single_label_disease_segmentor = Segmentor(model_path='/mnt/e/AI/mmdeploy/output/disease_detection', device_name='cuda', device_id=0)20fall_detector = Detector(model_path='/mnt/e/AI/mmdeploy/output/fall_detection_fastercnn', device_name='cuda', device_id=0)21mask_detector = Detector(model_path='/mnt/e/AI/mmdeploy/output/mask_detection', device_name='cuda', device_id=0)22smoker_detector_object = Detector(model_path='/mnt/e/AI/mmdeploy/output/smoker_nonsmoker', device_name='cuda', device_id=0)23 24def smoker_detector(frame, confidence_threshold=0.3):25 SMOKE_LABELS = ['smoker', 'nonsmoker'] # 新的标签列表26 bboxes, labels, masks = smoker_detector_object(frame) # 修改检测器名字27 28 # 获取有效的bbox索引29 valid_indices = [(i, SMOKE_LABELS[label]) for i, label in enumerate(labels) if SMOKE_LABELS[label] == 'smoker' and bboxes[i][4] >= confidence_threshold]30 31 smoker_count = 032 33 for i, label_name in valid_indices:34 bbox = bboxes[i]35 [left, top, right, bottom], score = bbox[0:4].astype(int), bbox[4]36 37 if label_name == 'smoker':38 color = (255, 0, 0) # 绿色用于'smoker'39 smoker_count += 140 41 line_thickness = 242 font_scale = 0.843 cv2.rectangle(frame, (left, top), (right, bottom), color, thickness=line_thickness)44 label_text = f"{label_name} ({score:.2f})"45 cv2.putText(frame, label_text, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, font_scale, color, line_thickness)46 47 if masks and masks[i].size:48 mask = masks[i]49 blue, green, red = cv2.split(frame)50 if mask.shape == frame.shape[:2]:51 mask_img = blue52 else:53 x0 = int(max(math.floor(bbox[0]) - 1, 0))54 y0 = int(max(math.floor(bbox[1]) - 1, 0))55 mask_img = blue[y0:y0 + mask.shape[0], x0:x0 + mask.shape[1]]56 cv2.bitwise_or(mask, mask_img, mask_img)57 frame = cv2.merge([blue, green, red])58 59 # 显示smoker的数量60 frame_height, frame_width = frame.shape[:2]61 summary_text = f"Smokers: {smoker_count}"62 cv2.putText(frame, summary_text, (frame_width - 200, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)63 64 return frame, smoker_count65 66 67 68def crack_classification(frame, confidence_threshold=0.5):69 # 定义标签70 labels_dict = {0: 'Negative', 1: 'Positive'}71 72 # 使用裂缝分类器进行预测73 result = crack_classifier(frame)74 75 # 获取最大置信度的标签ID76 label_id, score = max(result, key=lambda x: x[1])77 78 if label_id == 1 and score > confidence_threshold: # 如果检测到有裂缝，并且置信度超过阈值79 seg = crack_segmentor(frame)80 crack_pixel_count = np.sum(seg == 1)81 current_palette = [(255, 255, 255), (255, 0, 0)] # 背景为白色，裂缝为红色82 color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)83 for label, color in enumerate(current_palette):84 color_seg[seg == label, :] = color85 frame = frame * 0.5 + color_seg * 0.586 frame = frame.astype(np.uint8)87 elif label_id == 0 and score <= confidence_threshold:88 crack_pixel_count = None89 else:90 crack_pixel_count = None91 label_id = 0 # 这里我默认设置为0，即"Negative"，但你可以根据实际情况进行调整92 93 # 在图像右上角显示预测结果和置信度94 label_text = labels_dict[label_id] + f" ({score:.2f})"95 color = (255, 0, 0) if label_id == 1 else (0, 255, 0) # 裂缝为红色，否则为绿色96 font_scale = 0.897 line_thickness = 298 text_size = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, line_thickness)[0]99 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)100 101 return frame, labels_dict[label_id], crack_pixel_count102 103 104def crack_detection(frame):105 # 使用裂缝检测器进行检测106 seg = crack_segmentor(frame)107 crack_pixel_count = np.sum(seg == 1)108 109 # 如果检测到裂缝，进行可视化处理110 if crack_pixel_count > 0:111 current_palette = [(255, 255, 255), (255, 0, 0)] # 背景为白色，裂缝为红色112 color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)113 for label, color in enumerate(current_palette):114 color_seg[seg == label, :] = color115 frame = frame * 0.5 + color_seg * 0.5116 frame = frame.astype(np.uint8)117 118 # 在图像右上角显示检测到的裂缝像素数量119 label_text = f"Crack Pixels: {crack_pixel_count}"120 color = (255, 0, 0) if crack_pixel_count > 0 else (0, 255, 0) # 如果有裂缝则为红色，否则为绿色121 font_scale = 0.8122 line_thickness = 2123 text_size = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, line_thickness)[0]124 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)125 126 return frame, crack_pixel_count127 128 129def car_calculation(frame, confidence_threshold=0.7):130 CAR_LABEL = 'car' # 这里只有一个车辆标签131 bboxes, labels, masks = car_detector(frame)132 valid_indices = [i for i, label in enumerate(labels) if bboxes[i][4] >= confidence_threshold]133 134 car_count = 0135 136 for i in valid_indices:137 bbox = bboxes[i]138 [left, top, right, bottom], score = bbox[0:4].astype(int), bbox[4]139 140 color = (0, 255, 0) # 使用绿色标记车辆141 line_thickness = 2142 font_scale = 0.8143 144 cv2.rectangle(frame, (left, top), (right, bottom), color, thickness=line_thickness)145 label_text = CAR_LABEL + f" ({score:.2f})"146 cv2.putText(frame, label_text, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, font_scale, color, line_thickness)147 148 if masks and masks[i].size:149 mask = masks[i]150 blue, green, red = cv2.split(frame)151 if mask.shape == frame.shape[:2]:152 mask_img = blue153 else:154 x0 = int(max(math.floor(bbox[0]) - 1, 0))155 y0 = int(max(math.floor(bbox[1]) - 1, 0))156 mask_img = blue[y0:y0 + mask.shape[0], x0:x0 + mask.shape[1]]157 cv2.bitwise_or(mask, mask_img, mask_img)158 frame = cv2.merge([blue, green, red])159 160 car_count += 1161 162 return frame, car_count163 164 165 166def vest_detection(frame, confidence_threshold=0.3):167 VEST_LABELS = ['other_clothes', 'vest'] # 新的标签列表168 bboxes, labels, masks = vest_detector(frame)169 170 # 获取有效的bbox索引171 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]172 173 vest_count = 0174 other_clothes_count = 0175 176 for i, label_name in valid_indices:177 bbox = bboxes[i]178 [left, top, right, bottom], score = bbox[0:4].astype(int), bbox[4]179 180 if label_name == 'vest':181 color = (0, 255, 255) # 黄色用于'vest'182 vest_count += 1183 else:184 color = (255, 0, 0) # 蓝色用于'other_clothes'185 other_clothes_count += 1186 187 line_thickness = 2188 font_scale = 0.8189 cv2.rectangle(frame, (left, top), (right, bottom), color, thickness=line_thickness)190 label_text = f"{label_name} ({score:.2f})"191 cv2.putText(frame, label_text, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, font_scale, color, line_thickness)192 193 if masks and masks[i].size:194 mask = masks[i]195 blue, green, red = cv2.split(frame)196 if mask.shape == frame.shape[:2]:197 mask_img = blue198 else:199 x0 = int(max(math.floor(bbox[0]) - 1, 0))200 y0 = int(max(math.floor(bbox[1]) - 1, 0))201 mask_img = blue[y0:y0 + mask.shape[0], x0:x0 + mask.shape[1]]202 cv2.bitwise_or(mask, mask_img, mask_img)203 frame = cv2.merge([blue, green, red])204 205 # 显示vest和other_clothes的数量和置信度206 frame_height, frame_width = frame.shape[:2]207 summary_text = f"Vests: {vest_count}, Other Clothes: {other_clothes_count}"208 cv2.putText(frame, summary_text, (frame_width - 300, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)209 210 return frame, vest_count, other_clothes_count211 212def detect_falls(frame, confidence_threshold=0.5):213 # 假设输入图像是RGB格式，转换为BGR214 frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)215 216 LABELS = ['fall', 'person']217 # 初始化摔倒计数器218 fall_count = 0219 220 # 使用模型进行检测221 bboxes, labels, masks = fall_detector(frame)222 223 for bbox, label_id in zip(bboxes, labels):224 [left, top, right, bottom], score = bbox[0:4].astype(int), bbox[4]225 if score < confidence_threshold:226 continue227 if LABELS[label_id] == 'fall': # 仅显示摔倒的标注框228 cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255), 2)229 label_text = f"{LABELS[label_id]}: {int(score*100)}%"230 cv2.putText(frame, label_text, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)231 # 递增摔倒计数器232 fall_count += 1233 234 # 转换图像回RGB格式235 frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)236 237 # 返回处理后的图像和摔倒的数量238 return frame, fall_count239 240def leaf_disease_detection(frame, confidence_threshold=0.3):241 # 假设输入图像是RGB格式，转换为BGR242 frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)243 244 LABELS = ['disease']245 # 初始化病害计数器246 disease_count = 0247 bboxes, labels, masks = disease_object_detector(frame)248 indices = [i for i in range(len(bboxes))]249 for index, bbox, label_id in zip(indices, bboxes, labels):250 [left, top, right, bottom], score = bbox[0:4].astype(int), bbox[4]251 if score < confidence_threshold:252 continue253 cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255), 1)254 label_text = f"{LABELS[label_id]}: {int(score*100)}%"255 cv2.putText(frame, label_text, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)256 257 if masks[index].size:258 mask = masks[index]259 blue, green, red = cv2.split(frame)260 if mask.shape == frame.shape[:2]:261 mask_img = blue262 else:263 x0 = int(max(math.floor(bbox[0]) - 1, 0))264 y0 = int(max(math.floor(bbox[1]) - 1, 0))265 mask_img = blue[y0:y0 + mask.shape[0], x0:x0 + mask.shape[1]]266 cv2.bitwise_or(mask, mask_img, mask_img)267 frame = cv2.merge([blue, green, red])268 # 递增病害计数器269 disease_count += 1270 271 # 转换图像回RGB格式272 frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)273 274 # 返回处理后的图像、病害计数和保存的图像路径275 return frame, disease_count276 277def detect_masks(frame, confidence_threshold=0.5):278 # 假设输入图像是RGB格式，转换为BGR279 frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)280 281 LABELS = ['unfit', 'mask', 'nomask']282 # 初始化三个标签的计数器283 mask_count, nomask_count, unfit_count = 0, 0, 0284 285 # 使用模型进行检测286 bboxes, labels, masks = mask_detector(frame)287 288 for bbox, label_id in zip(bboxes, labels):289 [left, top, right, bottom], score = bbox[0:4].astype(int), bbox[4]290 if score < confidence_threshold:291 continue292 293 # 根据标签ID判断类别，并进行相应的计数294 if LABELS[label_id] == 'mask':295 mask_count += 1296 color = (0, 255, 0)297 elif LABELS[label_id] == 'nomask':298 nomask_count += 1299 color = (0, 0, 255)300 elif LABELS[label_id] == 'unfit':301 unfit_count += 1302 color = (255, 0, 0)303 304 cv2.rectangle(frame, (left, top), (right, bottom), color, 2)305 label_text = f"{LABELS[label_id]}: {int(score*100)}%"306 cv2.putText(frame, label_text, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)307 308 # 转换图像回RGB格式309 frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)310 311 # 返回处理后的图像和每个标签的数量312 return frame, mask_count, nomask_count, unfit_count313 314def helmet_detection(frame, confidence_threshold=0.3):315 316 HEL_LABELS = ['head', 'helmet']317 bboxes, labels, masks = helmet_detector(frame)318 valid_indices = [i for i, bbox in enumerate(bboxes) if bbox[4] >= confidence_threshold]319 320 helmet_count = 0321 head_count = 0322 323 for i in valid_indices:324 bbox = bboxes[i]325 label_id = labels[i]326 [left, top, right, bottom], score = bbox[0:4].astype(int), bbox[4]327 328 if HEL_LABELS[label_id] == 'helmet':329 color = (0, 255, 0) # Green color for 'helmet'330 line_thickness = 1331 font_scale = 0.5332 elif HEL_LABELS[label_id] == 'head':333 color = (255, 0, 0) # Red color for 'head'334 line_thickness = 1 # Increased line thickness for 'head' boxes335 font_scale = 0.5 # Decreased font size for 'head' labels336 337 cv2.rectangle(frame, (left, top), (right, bottom), color, thickness=line_thickness)338 label_text = HEL_LABELS[label_id] + f" ({score:.2f})"339 cv2.putText(frame, label_text, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, font_scale, color, line_thickness)340 341 if HEL_LABELS[label_id] == 'helmet':342 helmet_count += 1343 elif HEL_LABELS[label_id] == 'head':344 head_count += 1345 346 return frame, helmet_count, head_count347 348 349 350def human_calculation(frame, confidence_threshold=0.3):351 """352 Process the given image to count the number of humans.353 """354 HEL_LABELS = ['head', 'helmet']355 bboxes, labels, masks = helmet_detector(frame)356 357 human_count = 0 # Initialize human count358 359 for i in range(len(bboxes)):360 bbox = bboxes[i]361 label_id = labels[i]362 score = bbox[4]363 364 # Check if the label is 'head' or 'helmet' and the score is greater than confidence_threshold365 if HEL_LABELS[label_id] in ['head', 'helmet'] and score > confidence_threshold:366 human_count += 1367 [left, top, right, bottom] = bbox[0:4].astype(int)368 cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255), thickness=1) # Red color for boxes369 label_text = f"human ({score:.2f})" # Include confidence score in label_text370 cv2.putText(frame, label_text, (left, top - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)371 372 373 return frame, human_count374 375 376def red_tree(img):377 img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)378 def get_palette(num_classes=2):379 return [(255, 255, 255), (255, 0, 0)]380 seg = red_tree_segmentor(img_bgr)381 red_tree_pixel_count = np.sum(seg == 1)382 current_palette = get_palette()383 color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)384 for label, color in enumerate(current_palette):385 color_seg[seg == label, :] = color386 color_seg_bgr = color_seg[..., ::-1]387 388 img_bgr = img_bgr * 0.5 + color_seg_bgr * 0.5389 img_bgr = img_bgr.astype(np.uint8)390 img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)391 392 return img_rgb, red_tree_pixel_count393 394 395def leaf_disease(img):396 img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)397 398 def get_palette(num_classes=3):399 return [(255, 255, 255), (0, 255, 0), (255, 0, 0)]400 401 seg = leaf_disease_segmentor(img_bgr)402 403 leaf_pixel_count = np.sum(seg == 1)404 disease_pixel_count = np.sum(seg == 2)405 406 current_palette = get_palette()407 color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)408 409 for label, color in enumerate(current_palette):410 color_seg[seg == label, :] = color411 412 color_seg_bgr = color_seg[..., ::-1]413 img_bgr = img_bgr * 0.5 + color_seg_bgr * 0.5414 img_bgr = img_bgr.astype(np.uint8)415 img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)416 417 return img_rgb, leaf_pixel_count, disease_pixel_count418 419def single_label_disease(img):420 img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)421 422 def get_palette(num_classes=2):423 return [(255, 255, 255), (255, 0, 0)]424 425 seg = single_label_disease_segmentor(img_bgr)426 427 disease_pixel_count = np.sum(seg == 1)428 429 current_palette = get_palette()430 color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)431 432 for label, color in enumerate(current_palette):433 color_seg[seg == label, :] = color434 435 color_seg_bgr = color_seg[..., ::-1]436 img_bgr = img_bgr * 0.5 + color_seg_bgr * 0.5437 img_bgr = img_bgr.astype(np.uint8)438 img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)439 440 441 return img_rgb, disease_pixel_count442 443 444 445def get_image_examples():446 image_dir = "/mnt/e/AI/mmdeploy/gradio/photo"447 image_files = [f for f in os.listdir(image_dir) if f.endswith(('.png', '.jpg', '.jpeg'))]448 image_files.sort(key=lambda f: int(''.join(filter(str.isdigit, f)))) # 按数字排序449 example_choices = [450 '红树林识别', '红树林识别', '红树林识别', 451 '安全帽检测', '安全帽检测', '安全帽检测',452 '人数统计', '人数统计', '人数统计',453 '反光衣检测','反光衣检测','反光衣检测',454 '道路车辆统计', '道路车辆统计', '道路车辆统计',455 '裂缝识别', '裂缝识别', '裂缝识别',456 '吸烟检测','吸烟检测','吸烟检测',457 '树叶病害识别1','树叶病害识别1','树叶病害识别1',458 '树叶病害识别2','树叶病害识别2','树叶病害识别2',459 '树叶病害检测3','树叶病害检测3','树叶病害检测3',460 '摔倒检测','摔倒检测','摔倒检测',461 '口罩佩戴检测','口罩佩戴检测','口罩佩戴检测',462 ]463 464 confidence_thresholds = [465 0, 0, 0, 466 0.7, 0.8, 0.6, 467 0.3, 0.8, 0.5, 468 0.8, 0.7, 0.8, 469 0.5, 0.2, 0.7, 470 0, 0, 0, 471 0.6, 0.9, 0.5, 472 0, 0, 0, 473 0, 0, 0, 474 0.4, 0.4, 0.5, 475 0.9, 0.9, 0.5, 476 0.8, 0.6, 0.5477]478 examples = [[example_choices[i], f"{image_dir}/{image_file}", confidence_thresholds[i]] for i, image_file in enumerate(image_files)]479 return examples480 481 482 483 484model_choices = ['红树林识别','裂缝识别','树叶病害识别1','树叶病害识别2','树叶病害检测3', '安全帽检测','反光衣检测', '吸烟检测','摔倒检测', '口罩佩戴检测','人数统计','道路车辆统计']485 486 487def create_blank_image(width=640, height=480, color=(255, 255, 255)):488 blank_image = np.zeros((height, width, 3), np.uint8)489 blank_image[:, :] = color490 return blank_image491 492def process_image(model_choice, image_array=None, confidence_threshold=0.3):493 output_text = '当前未有图片输入，请上传图片后再次点击运行。'494 495 if image_array is None:496 img = create_blank_image()497 else:498 if model_choice not in model_choices:499 model_choice = "安全帽检测"500 # 以下是模型选择和执行逻辑501 if model_choice == "红树林识别":502 img, red_tree_pixel_count = red_tree(image_array) # 语义分割模型503 output_text = f"红树林的像素点有 {red_tree_pixel_count} 个。"504 elif model_choice == "安全帽检测":505 img, helmet_count, head_count = helmet_detection(image_array, confidence_threshold)506 output_text = f"佩戴安全帽的人数为：{helmet_count}，未佩戴安全帽的人数为：{head_count}。"507 elif model_choice == "人数统计":508 img, human_count = human_calculation(image_array, confidence_threshold)509 output_text = f"该图片人员总人数为: {human_count}。"510 elif model_choice == "反光衣检测":511 img, vest_count, other_clothes_count= vest_detection(image_array, confidence_threshold)512 output_text = f"该图片中总有 {vest_count} 人配备了反光衣，{other_clothes_count} 人没有配备反光衣。"513 elif model_choice == "道路车辆统计":514 img, car_count = car_calculation(image_array, confidence_threshold)515 output_text = f"该道路上目前共有 {car_count} 台车辆。"516 elif model_choice == "裂缝识别":517 img, crack_result, crack_pixel_count = crack_classification(image_array, confidence_threshold)518 if crack_result == "Positive":519 output_text = f"该图片内存在裂缝，裂缝的像素点有 {crack_pixel_count} 个。"520 else:521 output_text = "该图片不存在裂缝。"522 elif model_choice == "树叶病害检测3": 523 img, disease_count = leaf_disease_detection(image_array, confidence_threshold)524 if disease_count > 0:525 output_text = f"共检测到 {disease_count} 处病害。"526 else:527 output_text = "并未检测到病害。"528 elif model_choice == "吸烟检测":529 img, smoker_count = smoker_detector(image_array, confidence_threshold)530 output_text = f"当前图片有 {smoker_count} 人在吸烟。"531 elif model_choice == "树叶病害识别1":532 img, leaf_pixel_count, disease_pixel_count = leaf_disease(image_array) # 语义分割模型533 if disease_pixel_count == 0:534 output_text = "该树叶并未出现病害。"535 else:536 output_text = f"病害的像素点有 {disease_pixel_count} 个。"537 elif model_choice == "树叶病害识别2": 538 img, disease_pixel_count = single_label_disease(image_array) # 语义分割模型539 output_text = f"病害的像素点有 {disease_pixel_count} 个。"540 elif model_choice == "摔倒检测": # 您可以根据实际情况调整模型选择的名称541 img, fall_count = detect_falls(image_array,confidence_threshold)542 output_text = f"图像中摔倒的人数为 {fall_count} 人。"543 elif model_choice == "口罩佩戴检测": # 您可以根据实际情况调整模型选择的名称544 img, mask_count, nomask_count, unfit_count = detect_masks(image_array,confidence_threshold)545 output_text = f"当前佩戴口罩的人数为 {mask_count}，未正确佩戴口罩的人数为 {unfit_count}，没有佩戴口罩的人数为 {nomask_count}。"546 547 return img, output_text548 549def process_video(model_choice, video=None, confidence_threshold=0.3):550 551 # 内部函数：创建空白视频552 def create_blank_video(filename, duration=5, fps=30, width=640, height=480, color=(255, 255, 255)):553 fourcc = cv2.VideoWriter_fourcc(*'mp4v') # 使用mp4v编解码器554 out = cv2.VideoWriter(filename, fourcc, fps, (width, height))555 blank_image = np.zeros((height, width, 3), np.uint8)556 blank_image[:, :] = color557 for _ in range(int(fps * duration)):558 out.write(blank_image)559 out.release()560 561 # 检查视频是否存在562 if video is None:563 video_output_path = '/mnt/e/AI/mmdeploy/gradio/video/none.mp4'564 create_blank_video(video_output_path)565 output_text2 = '当前未有视频输入，请上传视频后再次点击运行。'566 return video_output_path, output_text2567 else:568 video_output_path = '/mnt/e/AI/mmdeploy/gradio/video/output_video.mp4'569 cap = cv2.VideoCapture(video)570 if not cap.isOpened():571 raise ValueError("无法打开视频文件")572 fps = int(cap.get(cv2.CAP_PROP_FPS))573 num_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))574 # 获取输入视频的分辨率575 width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))576 height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))577 # 使用帧采样的逻辑，但考虑到所有帧都需要处理，我们使用间隔为1的采样。578 clip_len, frame_interval, num_clips = 1, 1, num_frames579 avg_interval = (num_frames - clip_len * frame_interval + 1) / float(num_clips)580 frame_inds = []581 for i in range(num_clips):582 clip_offset = int(i * avg_interval + avg_interval / 2.0)583 for j in range(clip_len):584 ind = (j * frame_interval + clip_offset) % num_frames585 if num_frames <= clip_len * frame_interval - 1:586 ind = j % num_frames587 frame_inds.append(ind)588 589 fourcc = cv2.VideoWriter_fourcc(*'mp4v')590 processed_frames = []591 for frame_id in sorted(frame_inds): 592 cap.set(cv2.CAP_PROP_POS_FRAMES, frame_id) # 设置读取特定的帧593 ret, frame = cap.read()594 if not ret:595 break596 # 将帧率添加到视频的左上角597 cv2.putText(frame, "FPS: {}".format(fps), (10, 30), 598 cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2, cv2.LINE_AA)599 600 if model_choice == "红树林识别":601 # 在此处调用红树林模型处理帧602 processed_frame, red_tree_pixel_count = red_tree(frame)603 # 在处理后的帧的右上角添加文字604 cv2.putText(processed_frame, "Number of pixels in this frame: {}".format(red_tree_pixel_count), 605 (processed_frame.shape[1] - 300, 30), cv2.FONT_HERSHEY_SIMPLEX, 606 0.7, (255, 255, 255), 2)607 608 609 elif model_choice == "安全帽检测":610 # 在此处调用安全帽检测模型处理帧611 processed_frame, helmet_count, head_count = helmet_detection(frame, confidence_threshold)612 613 cv2.putText(processed_frame, "Number of people wearing helmets: {}".format(helmet_count), 614 (processed_frame.shape[1] - 400, 30), cv2.FONT_HERSHEY_SIMPLEX, 615 0.7, (255, 255, 255), 2)616 617 # 在上一行文字下方添加表示未佩戴安全帽的人数的文字618 cv2.putText(processed_frame, "Number of people without helmets: {}".format(head_count - helmet_count), 619 (processed_frame.shape[1] - 450, 60), cv2.FONT_HERSHEY_SIMPLEX, 620 0.7, (255, 255, 255), 2)621 622 623 elif model_choice == "人数统计":624 # 在此处调用人数统计模型处理帧625 processed_frame, human_count = human_calculation(frame, confidence_threshold)626 cv2.putText(processed_frame, "Current number of people: {}".format(human_count), 627 (processed_frame.shape[1] - 300, 30), cv2.FONT_HERSHEY_SIMPLEX, 628 0.7, (255, 255, 255), 2)629 630 elif model_choice == "反光衣检测":631 # 在此处调用反光衣检测模型处理帧632 processed_frame, vest_count, other_clothes_count= vest_detection(image_array, confidence_threshold)633 cv2.putText(processed_frame, "Number of reflective vests: {}".format(vest_count), 634 (processed_frame.shape[1] - 350, 30), cv2.FONT_HERSHEY_SIMPLEX, 635 0.7, (255, 255, 255), 2)636 cv2.putText(processed_frame, "Number without reflective vests: {}".format(other_clothes_count), 637 (processed_frame.shape[1] - 450, 60), cv2.FONT_HERSHEY_SIMPLEX, 638 0.7, (255, 255, 255), 2)639 640 elif model_choice == "道路车辆统计":641 # 在此处调用道路车辆统计模型处理帧642 processed_frame, car_count = car_calculation(frame, confidence_threshold)643 cv2.putText(processed_frame, "Number of vehicles: {}".format(car_count), 644 (processed_frame.shape[1] - 250, 30), cv2.FONT_HERSHEY_SIMPLEX, 645 0.7, (255, 255, 255), 2) 646 647 elif model_choice == "裂缝识别":648 # 在此处调用裂缝识别模型处理帧649 processed_frame, crack_pixel_count= crack_detection(frame)650 651 elif model_choice == "树叶病害检测3":652 # 在此处调用树叶病害检测模型处理帧653 processed_frame, disease_count= leaf_disease_detection(frame, confidence_threshold)654 # 在图像右上角显示叶片的病害数量655 label_text = f"Leaf Disease Count: {disease_count}"656 color = (0, 0, 255) # 红色657 font_scale = 0.8658 line_thickness = 2659 text_size = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, line_thickness)[0]660 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)661 662 663 664 elif model_choice == "吸烟检测":665 # 在此处调用吸烟检测模型处理帧666 processed_frame, smoker_count = smoker_detector(frame, confidence_threshold)667 668 # 准备要显示的文本669 text = f"吸烟者数量: {smoker_count}"670 671 # 获取文本大小672 text_size = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]673 674 # 计算文本的位置，以便它出现在帧的右上角675 text_position = (processed_frame.shape[1] - text_size[0] - 10, text_size[1] + 10)676 677 # 将文本绘制到帧上678 cv2.putText(processed_frame, text, text_position, cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)679 680 elif model_choice == "树叶病害识别1":681 # 在此处调用树叶病害识别模型处理帧682 processed_frame, _, disease_pixel_count = leaf_disease(frame)683 text = f"Current disease pixel count on the leaf: {disease_pixel_count}"684 text_size = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]685 cv2.putText(processed_frame, text, 686 (processed_frame.shape[1] - text_size[0] - 10, text_size[1] + 10), 687 cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)688 689 690 elif model_choice == "树叶病害识别2": 691 # 在此处调用树叶病害识别模型处理帧692 processed_frame, disease_pixel_count= single_label_disease(frame)693 text = f"Current disease pixel count on the leaf: {disease_pixel_count}"694 text_size = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]695 cv2.putText(processed_frame, text, 696 (processed_frame.shape[1] - text_size[0] - 10, text_size[1] + 10), 697 cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)698 699 elif model_choice == "口罩佩戴检测": # 您可以根据实际情况调整模型选择的名称700 processed_frame, mask_count, nomask_count, unfit_count = detect_masks(frame,confidence_threshold)701 cv2.putText(processed_frame, "Number wearing masks: {}".format(mask_count), 702 (processed_frame.shape[1] - 350, 30), cv2.FONT_HERSHEY_SIMPLEX, 703 0.7, (255, 255, 255), 2)704 705 cv2.putText(processed_frame, "Number not wearing masks: {}".format(nomask_count), 706 (processed_frame.shape[1] - 400, 60), cv2.FONT_HERSHEY_SIMPLEX, 707 0.7, (255, 255, 255), 2)708 709 cv2.putText(processed_frame, "Number wearing masks incorrectly: {}".format(unfit_count), 710 (processed_frame.shape[1] - 500, 90), cv2.FONT_HERSHEY_SIMPLEX, 711 0.7, (255, 255, 255), 2)712 713 elif model_choice == "摔倒检测": 714 715 processed_frame, fall_count= detect_falls(frame,confidence_threshold)716 cv2.putText(processed_frame, "Number of people who fell: {}".format(fall_count), 717 (processed_frame.shape[1] - 350, 30), cv2.FONT_HERSHEY_SIMPLEX, 718 0.7, (255, 255, 255), 2)719 720 processed_frames.append(processed_frame)721 out = cv2.VideoWriter(video_output_path, fourcc, fps, (width,height))722 for frame in processed_frames:723 out.write(frame)724 out.release()725 cap.release()726 output_text2 = '请点击蓝色按钮下载视频。'727 return video_output_path, output_text2728 729with gr.Blocks() as demo:730 gr.Markdown("# <center>启云科技AI识别示例样板V1.12</center>")731 gr.Markdown("请上传图像或视频进行预测")732 with gr.Tab("AI图像处理"):733 with gr.Row():734 image_input2 = gr.Image(label="上传图像", type="numpy")735 with gr.Column():736 image_input1 = gr.Dropdown(choices=model_choices, label="模型选择")737 image_input3 = gr.Slider(minimum=0, maximum=1, step=0.1, label="置信度阈值")738 with gr.Row():739 image_output1 = gr.Image(label="处理后的图像", type="numpy")740 with gr.Column():741 image_output2 = gr.Textbox(label="图像输出信息")742 image_button = gr.Button('请点击按钮进行图像预测')743 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)744 with gr.Tab("AI视频处理"):745 746 with gr.Row():747 video_input2 = gr.Video(label = '上传视频', format='mp4',interactive = True)748 with gr.Column():749 video_input1 = gr.Dropdown(choices=model_choices, label="模型选择")750 video_input3 = gr.Slider(minimum=0, maximum=1, 751 step=0.1, label="置信度阈值")752 with gr.Row():753 video_output1 = gr.File(label='处理后的视频', type='file')754 with gr.Column():755 video_output2 = gr.Textbox(label = '视频输出信息') 756 video_button = gr.Button('请点击按钮进行视频预测')757 with gr.Accordion("平台简介"):758 gr.Markdown("红树林识别模型、裂缝识别模型、树叶病害识别模型、安全帽检测模型、反光衣检测模型、吸烟检测模型、口罩佩戴检测、摔倒检测、人数统计模型及道路车辆统计模型展示平台。")759 image_button.click(process_image, inputs = [image_input1, image_input2, image_input3], outputs=[image_output1, image_output2])760 video_button.click(process_video, inputs=[video_input1,video_input2, video_input3], outputs=[video_output1, video_output2])761 762demo.launch(share=True)763 764