394 lines
14 KiB
Python
394 lines
14 KiB
Python
from typing import List, Optional
|
||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
import base64
|
||
import io
|
||
import cv2
|
||
import numpy as np
|
||
from PIL import Image
|
||
import os
|
||
import uuid
|
||
from datetime import datetime
|
||
|
||
|
||
from app.core.database import get_db
|
||
from app.services.imageServices import ImageService
|
||
from app.schemas.image import ImageBase
|
||
from app.schemas.ocr import ImageBase64Request
|
||
from app.util.responseHttp import ResponseUtil
|
||
from app.util.baiduOCR import BaiduOCR, BaiduOCRONNX
|
||
from app.util.yolov8Obj import Yolov8Obj, YOLOv8ONNX
|
||
|
||
# from app.crud.event import event
|
||
# from app.schemas.event import EventList, EventDetail, EventUpdate, EventQuery, TestEvent
|
||
|
||
baiduOCR = BaiduOCR()
|
||
baiduOcrOnnx = BaiduOCRONNX()
|
||
yolov8Obj = Yolov8Obj()
|
||
yolov8ONNX = YOLOv8ONNX()
|
||
|
||
router = APIRouter(prefix="/api/v1", tags=["ocr"])
|
||
|
||
@router.get("/hello")
|
||
async def get_hello():
|
||
|
||
# return {"data":"hello"}
|
||
return ResponseUtil.error(msg=f"OCR识别失败", data=None)
|
||
|
||
@router.get("/test_select")
|
||
async def test_select(
|
||
db: AsyncSession = Depends(get_db),
|
||
image_query: ImageBase = None
|
||
):
|
||
"""
|
||
测试查询
|
||
"""
|
||
result = await ImageService.get_image_list(db,image_query)
|
||
return ResponseUtil.success(msg="success", data=result)
|
||
|
||
|
||
@router.get("/test_ocr")
|
||
async def test_ocr(
|
||
# image_path: str = None
|
||
):
|
||
"""
|
||
测试OCR
|
||
"""
|
||
|
||
image_path = '/home/admin-root/haotian/康达瑞贝斯机器狗/data_image/001读表图片/2c7cc83019e7388a7041101da92c9829_frame_000000.jpg'
|
||
|
||
|
||
result = baiduOCR.ocr(image_path)
|
||
# print(result)
|
||
return ResponseUtil.success(msg="success", data=result)
|
||
# return ResponseUtil.success(msg="success")
|
||
|
||
|
||
@router.post("/ocr_from_base64")
|
||
async def ocr_from_base64(request: ImageBase64Request):
|
||
"""
|
||
从base64图片数据进行OCR识别
|
||
"""
|
||
try:
|
||
# 移除base64数据的前缀（如果有）
|
||
image_base64 = request.image_base64
|
||
if image_base64.startswith('data:image'):
|
||
# 格式如: data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQ...
|
||
image_base64 = image_base64.split(',')[1]
|
||
|
||
# 解码base64数据
|
||
image_data = base64.b64decode(image_base64)
|
||
|
||
# 将字节数据转换为PIL Image对象
|
||
image = Image.open(io.BytesIO(image_data))
|
||
|
||
# 创建临时文件路径
|
||
temp_dir = "tmp/ocr_images"
|
||
os.makedirs(temp_dir, exist_ok=True)
|
||
temp_filename = f"{uuid.uuid4()}.{request.image_type or 'jpg'}"
|
||
temp_path = os.path.join(temp_dir, temp_filename)
|
||
|
||
# 保存图片到临时文件
|
||
image.save(temp_path)
|
||
|
||
# 使用PaddleOCR进行识别
|
||
result = baiduOCR.ocr(temp_path)
|
||
|
||
# 删除临时文件
|
||
# os.remove(temp_path)
|
||
|
||
return ResponseUtil.success(msg="OCR识别成功", data=result)
|
||
except Exception as e:
|
||
return ResponseUtil.error(msg=f"OCR识别失败: {str(e)}", data=None)
|
||
|
||
|
||
@router.post("/ocr_onnx_from_base64")
|
||
async def ocr_from_base64(request: ImageBase64Request):
|
||
"""
|
||
从base64图片数据进行OCR识别
|
||
"""
|
||
try:
|
||
# 移除base64数据的前缀（如果有）
|
||
image_base64 = request.image_base64
|
||
if image_base64.startswith('data:image'):
|
||
# 格式如: data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQ...
|
||
image_base64 = image_base64.split(',')[1]
|
||
|
||
# 解码base64数据
|
||
image_data = base64.b64decode(image_base64)
|
||
|
||
# 将字节数据转换为PIL Image对象
|
||
image = Image.open(io.BytesIO(image_data))
|
||
|
||
# 创建临时文件路径
|
||
temp_dir = "./tmp/ocr_images"
|
||
os.makedirs(temp_dir, exist_ok=True)
|
||
temp_filename = f"{uuid.uuid4()}.{request.image_type or 'jpg'}"
|
||
temp_path = os.path.join(temp_dir, temp_filename)
|
||
|
||
# 保存图片到临时文件
|
||
image.save(temp_path)
|
||
|
||
# 使用PaddleOCR进行识别
|
||
result = baiduOcrOnnx.ocr(temp_path)
|
||
print(result)
|
||
|
||
# 删除临时文件
|
||
# os.remove(temp_path)
|
||
|
||
# return ResponseUtil.success(msg="OCR识别成功", data=[result['text'], result['confidence']])
|
||
return ResponseUtil.success(msg="OCR识别成功", data=result)
|
||
except Exception as e:
|
||
return ResponseUtil.error(msg=f"OCR识别失败: {str(e)}", data=None)
|
||
|
||
|
||
@router.post("/detect_from_base64_0")
|
||
async def ocr_from_base64(request: ImageBase64Request):
|
||
""" 从base64图片进行目标检测, 检测是否侵占消防区域"""
|
||
|
||
try:
|
||
image_base64 = request.image_base64
|
||
if image_base64.startswith('data:image'):
|
||
image_base64 = image_base64.split(',')[1]
|
||
image_data = base64.b64decode(image_base64)
|
||
|
||
image = Image.open(io.BytesIO(image_data))
|
||
|
||
temp_dir = "./tmp/detect_images"
|
||
os.makedirs(temp_dir, exist_ok=True)
|
||
temp_filename = f"{uuid.uuid4()}.{request.image_type or 'jpg'}"
|
||
temp_path = os.path.join(temp_dir, temp_filename)
|
||
|
||
image.save(temp_path)
|
||
|
||
cls, conf, coords = yolov8Obj.detect(temp_path)
|
||
|
||
os.remove(temp_path)
|
||
|
||
return ResponseUtil.success(msg="侵占消防区域目标检测成功,是否有遮挡", data=(len(cls)==0))
|
||
|
||
except Exception as e:
|
||
return ResponseUtil.error(msg=f"检测是否侵占消防区域失败: {str(e)}", data=None)
|
||
|
||
|
||
@router.post("/detect_onnx_from_base64_0")
|
||
async def ocr_from_base64(request: ImageBase64Request):
|
||
""" 从base64图片进行目标检测, 检测是否侵占消防区域"""
|
||
|
||
# 为不同类别生成不同颜色
|
||
def get_color(class_id):
|
||
np.random.seed(class_id)
|
||
return tuple(map(int, np.random.randint(0, 255, 3)))
|
||
|
||
try:
|
||
# image_base64 = request.image_base64
|
||
# if image_base64.startswith('data:image'):
|
||
# image_base64 = image_base64.split(',')[1]
|
||
# image_data = base64.b64decode(image_base64)
|
||
|
||
# image = Image.open(io.BytesIO(image_data))
|
||
|
||
# temp_dir = "./tmp/detect_images"
|
||
# os.makedirs(temp_dir, exist_ok=True)
|
||
# temp_filename = f"{uuid.uuid4()}.{request.image_type or 'jpg'}"
|
||
# temp_path = os.path.join(temp_dir, temp_filename)
|
||
|
||
# image.save(temp_path)
|
||
|
||
# boxes, scores, class_ids = yolov8ONNX.detect(image_data)
|
||
|
||
# os.remove(temp_path)
|
||
|
||
image_base64 = request.image_base64
|
||
if image_base64.startswith('data:image'):
|
||
image_base64 = image_base64.split(',')[1]
|
||
image_data = base64.b64decode(image_base64)
|
||
|
||
# 将字节流转换为OpenCV格式的BGR图像
|
||
image_np = np.frombuffer(image_data, np.uint8)
|
||
image_cv = cv2.imdecode(image_np, cv2.IMREAD_COLOR) # 这会得到BGR格式的图像
|
||
|
||
# 不再需要保存到临时文件，直接使用内存中的图像
|
||
boxes, scores, class_ids = yolov8ONNX.detect(image_cv)
|
||
|
||
|
||
# 绘制检测结果
|
||
image_with_boxes = image_cv.copy()
|
||
|
||
for box, score, class_id in zip(boxes, scores, class_ids):
|
||
# 获取边界框坐标
|
||
x1, y1, x2, y2 = map(int, box)
|
||
|
||
# 生成随机颜色或使用固定颜色方案
|
||
color = get_color(class_id) # 绿色，也可以根据class_id设置不同颜色
|
||
|
||
# 绘制边界框
|
||
cv2.rectangle(image_with_boxes, (x1, y1), (x2, y2), color, 2)
|
||
|
||
# 准备标签文本
|
||
label = f"Class {class_id}: {score:.2f}"
|
||
# 如果你有类别名称字典，可以这样使用：
|
||
# label = f"{class_names[class_id]}: {score:.2f}"
|
||
|
||
# 计算文本大小以绘制背景
|
||
(text_width, text_height), baseline = cv2.getTextSize(
|
||
label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1
|
||
)
|
||
|
||
# 绘制文本背景
|
||
cv2.rectangle(
|
||
image_with_boxes,
|
||
(x1, y1 - text_height - baseline - 5),
|
||
(x1 + text_width, y1),
|
||
color,
|
||
-1
|
||
)
|
||
|
||
# 绘制文本
|
||
cv2.putText(
|
||
image_with_boxes,
|
||
label,
|
||
(x1, y1 - 5),
|
||
cv2.FONT_HERSHEY_SIMPLEX,
|
||
0.5,
|
||
(0, 0, 0), # 黑色文字
|
||
1,
|
||
cv2.LINE_AA
|
||
)
|
||
|
||
# 创建保存目录
|
||
save_dir = "tmp/detect_images"
|
||
os.makedirs(save_dir, exist_ok=True)
|
||
|
||
# 生成带时间戳的随机文件名
|
||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
random_id = uuid.uuid4().hex[:8] # 取UUID的前8位
|
||
filename = f"detect_{timestamp}_{random_id}.jpg"
|
||
output_path = os.path.join(save_dir, filename)
|
||
|
||
# 保存图像
|
||
cv2.imwrite(output_path, image_with_boxes)
|
||
print(f"检测结果已保存到: {output_path}")
|
||
|
||
return ResponseUtil.success(msg="侵占消防区域目标检测成功,是否有遮挡", data=(len(boxes)==0))
|
||
|
||
except Exception as e:
|
||
return ResponseUtil.error(msg=f"检测是否侵占消防区域失败: {str(e)}", data=None)
|
||
|
||
@router.post("/detect_from_base64_1")
|
||
async def ocr_from_base64(request: ImageBase64Request):
|
||
""" 从吧色图64图片进行目标检测, 检测是否存在灭火器"""
|
||
|
||
# try:
|
||
# 为不同类别生成不同颜色
|
||
def get_color(class_id):
|
||
np.random.seed(class_id)
|
||
return tuple(map(int, np.random.randint(0, 255, 3)))
|
||
|
||
try:
|
||
|
||
image_base64 = request.image_base64
|
||
if image_base64.startswith('data:image'):
|
||
image_base64 = image_base64.split(',')[1]
|
||
image_data = base64.b64decode(image_base64)
|
||
|
||
# 将字节流转换为OpenCV格式的BGR图像
|
||
image_np = np.frombuffer(image_data, np.uint8)
|
||
image_cv = cv2.imdecode(image_np, cv2.IMREAD_COLOR) # 这会得到BGR格式的图像
|
||
|
||
# 不再需要保存到临时文件，直接使用内存中的图像
|
||
boxes, scores, class_ids = yolov8ONNX.detect(image_cv)
|
||
|
||
|
||
# 绘制检测结果
|
||
image_with_boxes = image_cv.copy()
|
||
|
||
for box, score, class_id in zip(boxes, scores, class_ids):
|
||
# 获取边界框坐标
|
||
x1, y1, x2, y2 = map(int, box)
|
||
|
||
# 生成随机颜色或使用固定颜色方案
|
||
color = get_color(class_id) # 绿色，也可以根据class_id设置不同颜色
|
||
|
||
# 绘制边界框
|
||
cv2.rectangle(image_with_boxes, (x1, y1), (x2, y2), color, 2)
|
||
|
||
# 准备标签文本
|
||
label = f"Class {class_id}: {score:.2f}"
|
||
# 如果你有类别名称字典，可以这样使用：
|
||
# label = f"{class_names[class_id]}: {score:.2f}"
|
||
|
||
# 计算文本大小以绘制背景
|
||
(text_width, text_height), baseline = cv2.getTextSize(
|
||
label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1
|
||
)
|
||
|
||
# 绘制文本背景
|
||
cv2.rectangle(
|
||
image_with_boxes,
|
||
(x1, y1 - text_height - baseline - 5),
|
||
(x1 + text_width, y1),
|
||
color,
|
||
-1
|
||
)
|
||
|
||
# 绘制文本
|
||
cv2.putText(
|
||
image_with_boxes,
|
||
label,
|
||
(x1, y1 - 5),
|
||
cv2.FONT_HERSHEY_SIMPLEX,
|
||
0.5,
|
||
(0, 0, 0), # 黑色文字
|
||
1,
|
||
cv2.LINE_AA
|
||
)
|
||
|
||
# 创建保存目录
|
||
save_dir = "tmp/detect_images"
|
||
os.makedirs(save_dir, exist_ok=True)
|
||
|
||
# 生成带时间戳的随机文件名
|
||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
random_id = uuid.uuid4().hex[:8] # 取UUID的前8位
|
||
filename = f"detect_{timestamp}_{random_id}.jpg"
|
||
output_path = os.path.join(save_dir, filename)
|
||
|
||
# 保存图像
|
||
cv2.imwrite(output_path, image_with_boxes)
|
||
print(f"检测结果已保存到: {output_path}")
|
||
|
||
return ResponseUtil.success(msg="灭火器目标检测成功,是否存在灭火器", data=(len(class_ids)!=0 and 0 in class_ids))
|
||
# return ResponseUtil.success(msg="目标检测成功", data=cls)
|
||
|
||
except Exception as e:
|
||
return ResponseUtil.error(msg=f"检测是否存在灭火器失败: {str(e)}", data=None)
|
||
|
||
# @router.post("/detect_from_base64_1")
|
||
# async def ocr_from_base64(request: ImageBase64Request):
|
||
# """ 从吧色图64图片进行目标检测, 检测是否存在灭火器"""
|
||
|
||
# try:
|
||
# image_base64 = request.image_base64
|
||
# if image_base64.startswith('data:image'):
|
||
# image_base64 = image_base64.split(',')[1]
|
||
# image_data = base64.b64decode(image_base64)
|
||
|
||
# image = Image.open(io.BytesIO(image_data))
|
||
|
||
# temp_dir = "./tmp/detect_images"
|
||
# os.makedirs(temp_dir, exist_ok=True)
|
||
# temp_filename = f"{uuid.uuid4()}.{request.image_type or 'jpg'}"
|
||
# temp_path = os.path.join(temp_dir, temp_filename)
|
||
|
||
# image.save(temp_path)
|
||
|
||
# cls, conf, coords = yolov8Obj.detect(temp_path)
|
||
|
||
# os.remove(temp_path)
|
||
|
||
# return ResponseUtil.success(msg="灭火器目标检测成功,是否存在灭火器", data=(len(cls)!=0 and 0 in cls))
|
||
# # return ResponseUtil.success(msg="目标检测成功", data=cls)
|
||
|
||
# except Exception as e:
|
||
# return ResponseUtil.error(msg=f"检测是否存在灭火器失败: {str(e)}", data=None) |