from fastapi import FastAPI, UploadFile, HTTPException
from pydantic import BaseModel
from PIL import Image
import io
import base64
import cv2
import numpy as np

def opencv_image_to_bytes(opencv_image):
    """将 OpenCV 图像对象转换为 bytes 对象"""
    # 使用 cv2.imencode() 将图像编码为 numpy 数组
    ret, buffer = cv2.imencode('.jpg', opencv_image)
    
    # 将 numpy 数组转换为 bytes
    img_bytes = buffer.tobytes()
    
    return img_bytes

def handleImageByCV2(bs: bytes) -> {int, bytes}:
    # 读取图像
    img = cv2.imdecode(np.frombuffer(bs, np.uint8), cv2.IMREAD_COLOR)
    
    # 将图像转换为灰度图
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    
    # 使用高斯模糊处理图像
    blurred = cv2.GaussianBlur(gray, (9, 9), 2)
    
    # 使用霍夫圆检测法检测圆圈
    circles = cv2.HoughCircles(blurred, cv2.HOUGH_GRADIENT, dp=1.2, minDist=30,
                               param1=80, param2=38, minRadius=15, maxRadius=60)
    
    # 确认至少检测到一个圆圈
    if circles is not None:
        circles = np.round(circles[0, :]).astype("int")
        
        print(f"检测到 {len(circles)} 个圆圈")
        for i, (x, y, r) in enumerate(circles, start=1): # 从1开始计数
            # 在图像中绘制圆圈
            cv2.circle(img, (x, y), r, (0, 255, 0), 4)
            # 在图像中绘制圆心
            cv2.circle(img, (x, y), 2, (0, 128, 255), 3)
            # 在图像中绘制圆圈编号
            cv2.putText(img, f'{i}', (x - 10, y + 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)

        return len(circles), opencv_image_to_bytes(img) 
    else:
        return 0, opencv_image_to_bytes(img)

app = FastAPI()

class ImageResponse(BaseModel):
    image_base64: str
    circle_count: int

@app.post("/upload_image", response_model=ImageResponse)
async def upload_image(image_file: UploadFile):
    if not(image_file.filename.endswith('.jpg') or image_file.filename.endswith('.jpeg')):
        raise HTTPException(status_code=400, detail="Only jpg or jpeg files are supported")

    image_bytes = await image_file.read()
    count, new_bytes = handleImageByCV2(image_bytes)

    img_base64 = base64.b64encode(new_bytes).decode('utf-8')

    return ImageResponse(image_base64=img_base64, circle_count=count)

if __name__ == '__main__':
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)