睿抗平安城市火灾识别代码讲解见B站

# 初始化OCR引擎
ocr_engine = PaddleOCR(use_angle_cls=True, use_gpu=False, lang="ch")

# 定义路径
IMAGE_DIR = r'E:\2024RAICOM\P3_1\images\train_img'
RESULT_FILEPATH = r'E:\2024RAICOM\P3_1\result3\result.txt'

def fetch_label_data(label_filepath):
"""从标签文件中读取数据"""
try:
with open(label_filepath, mode='rt', encoding='utf-8') as file:
content = file.read()
except FileNotFoundError:
print("标签文件未找到，请检查路径。")
return {}
except Exception as e:
print(f"读取文件时发生错误: {e}")
return {}

data_dict = {}
for line in content.splitlines():
parts = line.strip().split('\t')
if len(parts) == 2:
img_name, json_str = parts
try:
data_dict[img_name] = json.loads(json_str)
except json.JSONDecodeError:
print(f"解析JSON数据失败: {json_str}")
return data_dict

def organize_floors(floors_info):
"""组织楼层数据，按Y坐标降序排列"""
return sorted(floors_info, key=lambda x: x['points'][0][1], reverse=True)

def pinpoint_fire_location(spark_instances, ordered_floors):
"""依据火花位置识别火灾楼层"""
detected_fires = {}
for spark in spark_instances:
spark_midpoint_y = (spark['points'][0][1] + spark['points'][2][1]) / 2
for index, floor in enumerate(ordered_floors, start=1):
floor_y_range = (floor['points'][0][1], floor['points'][2][1])
if floor_y_range[0] <= spark_midpoint_y <= floor_y_range[1]:
detected_fires[index] = True
break
return sorted(detected_fires, reverse=False)

def identify_building(image_path):
"""识别建筑物名称"""
ocr_results = ocr_engine.ocr(image_path, cls=True)
for result_group in ocr_results:
for line in result_group:
text_content = line[1][0]
if "大厦" in text_content:
return text_content
return "未知大厦名称"

def process_images_and_output_results(labels_file, image_directory, output_file):
"""处理图片并输出结果"""
if not os.path.exists(labels_file):
print("标签文件不存在。")
return

labels_data = fetch_label_data(labels_file)
processed_images_log = []

try:
with open(output_file, 'w', encoding='utf-8') as out_file:
for image_name, label_info in labels_data.items():
clean_img_name = image_name.replace('train\\', '') # 这里的img/要看替换一下
img_path = os.path.join(image_directory, clean_img_name)
if not os.path.exists(img_path):
print(f"图片'{image_name}'未找到。")
continue

building_title = identify_building(img_path)
floor_segments = [item for item in label_info if item.get('transcription') == '0']
spark_instances = [item for item in label_info if item.get('transcription') == '1']

if not floor_segments and not spark_instances:
print(f"'{image_name}'缺少必要标签数据，跳过。")
continue

floors_sorted = organize_floors(floor_segments)
fire_floor_numbers = pinpoint_fire_location(spark_instances, floors_sorted)

processed_images_log.append(image_name)
output_message = generate_output_message(image_name, building_title, fire_floor_numbers)
print(output_message)
out_file.write(output_message + '\n')

except Exception as e:
print(f"处理过程中出现错误: {e}")

def generate_output_message(image_name, building_name, fire_floors):
"""生成输出消息"""
clean_img_name = image_name.replace('train\\', '') # 这里的img/要看替换一下
if fire_floors:
fire_floor_str = "、".join(map(str, fire_floors))

return f"图片'{clean_img_name}'中的楼宇为{building_name}，火灾发生在第{fire_floor_str}层。"
else:
return f"图片'{clean_img_name}'中的楼宇为{building_name}，未检测到火灾。"

if __name__ == "__main__":
process_images_and_output_results(r'E:\2024RAICOM\2024RAICOM_city\Label.txt', IMAGE_DIR, RESULT_FILEPATH)