|
@@ -0,0 +1,202 @@
|
|
|
|
|
+import io
|
|
|
|
|
+from PIL import Image
|
|
|
|
|
+from fastapi import FastAPI, File, UploadFile, HTTPException
|
|
|
|
|
+from fastapi.responses import JSONResponse
|
|
|
|
|
+from ultralytics import YOLO
|
|
|
|
|
+from contextlib import asynccontextmanager
|
|
|
|
|
+import base64
|
|
|
|
|
+import json
|
|
|
|
|
+import os
|
|
|
|
|
+import uvicorn
|
|
|
|
|
+import numpy as np
|
|
|
|
|
+
|
|
|
|
|
+COLOR_LIST = {
|
|
|
|
|
+ "Single_Door": (255, 0, 0), "Double_Door": (0, 255, 0), "Sliding_Door": (0, 0, 255),
|
|
|
|
|
+ "Doorway_Opening": (255, 255, 0), "Mother_Son_Door": (0, 255, 255), "Revolving_Door": (255, 0, 255),
|
|
|
|
|
+ "Rolling_Gate": (192, 192, 192), "Rolling_Shutter_Door": (128, 0, 0), "Door_Window_Combo": (128, 128, 0),
|
|
|
|
|
+ "Casement_Window": (0, 128, 0), "Sliding_Window": (128, 0, 128), "French_Window": (0, 128, 128),
|
|
|
|
|
+ "Bay_Window": (0, 0, 128), "Single_Bay_Window": (255, 165, 0), "Three_Seat_Sofa": (255, 192, 203),
|
|
|
|
|
+ "Single_Seat_Sofa": (100, 149, 237), "Corner_Sofa": (255, 215, 0), "Coffee_Table": (218, 165, 32),
|
|
|
|
|
+ "Side_Table": (255, 170, 170), "TV_Cabinet": (170, 255, 170), "Wardrobe": (170, 170, 255),
|
|
|
|
|
+ "Dressing_Table": (255, 255, 170), "Nightstand": (170, 255, 255), "Shoe_Cabinet": (255, 170, 255),
|
|
|
|
|
+ "Sideboard": (255, 220, 190), "Kitchen_Cabinet": (190, 255, 220), "Bathroom_Cabinet": (220, 190, 255),
|
|
|
|
|
+ "Bookshelf": (255, 190, 220), "Dining_Table": (190, 220, 255), "Desk": (220, 255, 190),
|
|
|
|
|
+ "Chair": (130, 220, 130), "Stool": (220, 130, 130), "Bench": (130, 130, 220),
|
|
|
|
|
+ "Double_Bed": (220, 220, 130), "Single_Bed": (130, 220, 220), "Gas_Stove": (220, 130, 220),
|
|
|
|
|
+ "Sink": (150, 50, 50), "Bathtub": (50, 150, 50), "Toilet": (50, 50, 150),
|
|
|
|
|
+ "Squat_Toilet": (150, 150, 50), "Vanity_Unit": (50, 150, 150), "Shower_Head": (150, 50, 150),
|
|
|
|
|
+ "Mop_Sink": (100, 0, 0), "Rug": (0, 100, 0), "Washing_Machine": (0, 0, 100),
|
|
|
|
|
+ "TV": (100, 100, 0), "Refrigerator": (0, 100, 100), "Wall_Mounted_AC": (100, 0, 100),
|
|
|
|
|
+ "Cabinet_AC": (75, 75, 75), "Water_Dispenser": (175, 175, 175), "Hair_Dryer": (25, 25, 25),
|
|
|
|
|
+ "Microwave": (225, 225, 225), "Grill": (10, 140, 10), "Robot_Vacuum": (140, 10, 10),
|
|
|
|
|
+ "Table_Lamp": (10, 10, 140), "Living_Room": (140, 140, 10), "Dining_Room": (10, 140, 140),
|
|
|
|
|
+ "Bedroom": (140, 10, 140), "Bathroom": (240, 100, 0), "Balcony": (0, 240, 100),
|
|
|
|
|
+ "Entrance": (100, 0, 240), "Storage_Room": (240, 0, 100), "Community_Name": (0, 100, 240),
|
|
|
|
|
+ "Community_Gate": (100, 240, 0), "Building_Number": (100, 100, 100), "Floor_Number": (200, 200, 200),
|
|
|
|
|
+ "House_Number": (255, 10, 10), "Room_Number": (10, 255, 10)
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+english_lables = [
|
|
|
|
|
+ "Single_Door", "Double_Door", "Sliding_Door", "Doorway_Opening", "Mother_Son_Door", "Rolling_Gate",
|
|
|
|
|
+ "Rolling_Shutter_Door", "Revolving_Door", "Door_Window_Combo", "Casement_Window", "Sliding_Window",
|
|
|
|
|
+ "French_Window", "Bay_Window", "Single_Bay_Window", "Three_Seat_Sofa", "Single_Seat_Sofa",
|
|
|
|
|
+ "Corner_Sofa", "Coffee_Table", "Side_Table", "TV_Cabinet", "Wardrobe", "Dressing_Table",
|
|
|
|
|
+ "Nightstand", "Shoe_Cabinet", "Sideboard", "Kitchen_Cabinet", "Bathroom_Cabinet", "Bookshelf",
|
|
|
|
|
+ "Dining_Table", "Desk", "Chair", "Stool", "Bench", "Double_Bed", "Single_Bed", "Gas_Stove",
|
|
|
|
|
+ "Sink", "Bathtub", "Toilet", "Squat_Toilet", "Vanity_Unit", "Shower_Head", "Mop_Sink", "Rug",
|
|
|
|
|
+ "Washing_Machine", "TV", "Refrigerator", "Wall_Mounted_AC", "Cabinet_AC", "Water_Dispenser",
|
|
|
|
|
+ "Hair_Dryer", "Microwave", "Grill", "Robot_Vacuum", "Table_Lamp", "Living_Room", "Dining_Room",
|
|
|
|
|
+ "Bedroom", "Bathroom", "Balcony", "Entrance", "Storage_Room", "Community_Name", "Community_Gate",
|
|
|
|
|
+ "Building_Number", "Floor_Number", "House_Number", "Room_Number"
|
|
|
|
|
+]
|
|
|
|
|
+
|
|
|
|
|
+label_dic = {
|
|
|
|
|
+ "单开门": "Single_Door", "双开门": "Double_Door", "移门": "Sliding_Door", "垭口": "Doorway_Opening",
|
|
|
|
|
+ "子母门": "Mother_Son_Door", "拉闸门": "Rolling_Gate", "卷闸门": "Rolling_Shutter_Door",
|
|
|
|
|
+ "旋转门": "Revolving_Door", "门联窗": "Door_Window_Combo", "平开窗": "Casement_Window",
|
|
|
|
|
+ "推拉窗": "Sliding_Window", "落地窗": "French_Window", "飘窗": "Bay_Window",
|
|
|
|
|
+ "单面飘窗": "Single_Bay_Window", "三人沙发": "Three_Seat_Sofa", "单人沙发": "Single_Seat_Sofa",
|
|
|
|
|
+ "转角沙发": "Corner_Sofa", "茶几": "Coffee_Table", "边几": "Side_Table",
|
|
|
|
|
+ "电视柜": "TV_Cabinet", "衣柜": "Wardrobe", "梳妆台": "Dressing_Table",
|
|
|
|
|
+ "床头柜": "Nightstand", "鞋柜": "Shoe_Cabinet", "边柜": "Sideboard",
|
|
|
|
|
+ "橱柜": "Kitchen_Cabinet", "浴室柜": "Bathroom_Cabinet", "书柜": "Bookshelf",
|
|
|
|
|
+ "餐桌": "Dining_Table", "书桌": "Desk", "椅子": "Chair", "凳子": "Stool",
|
|
|
|
|
+ "条凳": "Bench", "双人床": "Double_Bed", "单人床": "Single_Bed", "燃气灶": "Gas_Stove",
|
|
|
|
|
+ "水槽": "Sink", "浴缸": "Bathtub", "马桶": "Toilet", "蹲便器": "Squat_Toilet",
|
|
|
|
|
+ "洗漱台": "Vanity_Unit", "花洒": "Shower_Head", "拖把池": "Mop_Sink", "地毯": "Rug",
|
|
|
|
|
+ "洗衣机": "Washing_Machine", "电视": "TV", "冰箱": "Refrigerator",
|
|
|
|
|
+ "挂式空调": "Wall_Mounted_AC", "柜式空调": "Cabinet_AC", "饮水机": "Water_Dispenser",
|
|
|
|
|
+ "吹风机": "Hair_Dryer", "微波炉": "Microwave", "烧烤架": "Grill",
|
|
|
|
|
+ "扫地机器人": "Robot_Vacuum", "台灯": "Table_Lamp", "客厅": "Living_Room",
|
|
|
|
|
+ "餐厅": "Dining_Room", "卧室": "Bedroom", "卫生间": "Bathroom", "阳台": "Balcony",
|
|
|
|
|
+ "玄关": "Entrance", "储物间": "Storage_Room", "小区名": "Community_Name",
|
|
|
|
|
+ "小区大门": "Community_Gate", "楼号": "Building_Number", "楼层": "Floor_Number",
|
|
|
|
|
+ "门牌号": "House_Number", "房号": "Room_Number"
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def parse_seg_results(results, filename):
|
|
|
|
|
+ """解析旋转框检测结果"""
|
|
|
|
|
+ detect_results = []
|
|
|
|
|
+ imageHeight, imageWidth = results[0].orig_shape
|
|
|
|
|
+ imagePath = os.path.basename(filename)
|
|
|
|
|
+ img_inf = (imagePath, imageWidth, imageHeight)
|
|
|
|
|
+
|
|
|
|
|
+ segs = []
|
|
|
|
|
+ for result in results:
|
|
|
|
|
+ xywh = result.boxes.xywh.cpu().numpy()
|
|
|
|
|
+ box = result.boxes.cpu().numpy()
|
|
|
|
|
+ x_c, y_c, w, h = int(xywh[0, 0]), int(xywh[0, 1]), int(xywh[0, 2]), int(xywh[0, 3])
|
|
|
|
|
+ x = int(x_c - w / 2)
|
|
|
|
|
+ y = int(y_c - h / 2)
|
|
|
|
|
+ x1, y1, x2, y2 = x, y, x + w, y + h # 左上角顶点和右下角顶点
|
|
|
|
|
+
|
|
|
|
|
+ ##############
|
|
|
|
|
+ points_list = []
|
|
|
|
|
+ tlp = (x1, y1)
|
|
|
|
|
+ trp = (x2, y1)
|
|
|
|
|
+ blp = (x1, y2)
|
|
|
|
|
+ brp = (x2, y2)
|
|
|
|
|
+ points_list.append(tlp)
|
|
|
|
|
+ points_list.append(trp)
|
|
|
|
|
+ points_list.append(blp)
|
|
|
|
|
+ points_list.append(brp)
|
|
|
|
|
+ all_x = [point[0] for point in points_list]
|
|
|
|
|
+ all_y = [point[1] for point in points_list]
|
|
|
|
|
+ ##############
|
|
|
|
|
+
|
|
|
|
|
+ segments_xy = result.masks.xy[0] # 归一化坐标,形状通常是 (N, 2) 的一维数组
|
|
|
|
|
+ imageHeight, imageWidth = result.orig_shape
|
|
|
|
|
+ seg_points = []
|
|
|
|
|
+ length = segments_xy.shape[0]
|
|
|
|
|
+
|
|
|
|
|
+ # for i in range(length):
|
|
|
|
|
+ # seg_point = (int(segments_xy[i, :][0]), int(segments_xy[i, :][1]))
|
|
|
|
|
+ # seg_points.append(seg_point)
|
|
|
|
|
+ # points_array = np.array(seg_points).reshape((-1, 2))
|
|
|
|
|
+ # points_array = points_array.astype(np.int32)
|
|
|
|
|
+ # pts_for_polylines = [points_array]
|
|
|
|
|
+
|
|
|
|
|
+ for i in range(length):
|
|
|
|
|
+ x, y = (int(segments_xy[i, :][0]), int(segments_xy[i, :][1]))
|
|
|
|
|
+ seg_points.append(x)
|
|
|
|
|
+ seg_points.append(y)
|
|
|
|
|
+
|
|
|
|
|
+ yolov_x = [round(x / imageWidth, 6) for x in all_x]
|
|
|
|
|
+ yolov_y = [round(y / imageHeight, 6) for y in all_y]
|
|
|
|
|
+
|
|
|
|
|
+ segs.append({
|
|
|
|
|
+ # "bbox": points_list,
|
|
|
|
|
+ "bbox": [x1, y1, x2, y1, x2, y2, x1, y2],
|
|
|
|
|
+ "seg_points": seg_points,
|
|
|
|
|
+ "category": label_dic[result.names[int(box.cls[0])]],
|
|
|
|
|
+ "color": COLOR_LIST[label_dic[result.names[int(box.cls[0])]]],
|
|
|
|
|
+ "label": f"{int(box.cls[0])} {yolov_x[0]} {yolov_y[0]} {yolov_x[1]}{yolov_y[1]}{yolov_x[2]}{yolov_y[2]} {yolov_x[3]} {yolov_y[3]}",
|
|
|
|
|
+ "name": result.names[int(box.cls[0])],
|
|
|
|
|
+ "score": float(box.conf[0])
|
|
|
|
|
+ })
|
|
|
|
|
+ return segs, img_inf
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+models = {}
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+@asynccontextmanager
|
|
|
|
|
+async def lifespan(app: FastAPI):
|
|
|
|
|
+ # 在应用启动时执行
|
|
|
|
|
+ print("应用启动...开始加载模型...")
|
|
|
|
|
+
|
|
|
|
|
+ models["seg"] = YOLO('../../checkpoints/seg_detect_Panoramic_Home.pt')
|
|
|
|
|
+
|
|
|
|
|
+ print("模型加载完成!")
|
|
|
|
|
+ yield
|
|
|
|
|
+ # 在应用关闭时执行 (如果需要)
|
|
|
|
|
+ print("应用关闭...")
|
|
|
|
|
+ models.clear()
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+app = FastAPI(lifespan=lifespan)
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+@app.get("/")
|
|
|
|
|
+def read_root():
|
|
|
|
|
+ return {"message": "欢迎使用 YOLO 系列模型部署 API"}
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+@app.post("/seg") # 假设您将此端点修改为接收 JSON 数据
|
|
|
|
|
+async def predict_seg(data: dict):
|
|
|
|
|
+ """旋转框检测任务,接收 Base64 编码的图片"""
|
|
|
|
|
+ try:
|
|
|
|
|
+ # 从 JSON 中获取 Base64 字符串
|
|
|
|
|
+ base64_string = data.get("image_data")
|
|
|
|
|
+ if not base64_string:
|
|
|
|
|
+ raise HTTPException(status_code=400, detail="未找到 'image_data' 字段")
|
|
|
|
|
+
|
|
|
|
|
+ # 解码 Base64 字符串
|
|
|
|
|
+ image_bytes = base64.b64decode(base64_string)
|
|
|
|
|
+
|
|
|
|
|
+ # 将字节流转换为 PIL.Image 对象
|
|
|
|
|
+ image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
|
|
|
|
|
+
|
|
|
|
|
+ # 使用加载好的模型进行推理
|
|
|
|
|
+ # results = models["seg"](image, imgsz=2560)
|
|
|
|
|
+ results = models["seg"](image, imgsz=640)
|
|
|
|
|
+ segs, imginf = parse_seg_results(results, data.get("filename"))
|
|
|
|
|
+
|
|
|
|
|
+ content_data = JSONResponse(content={"shapes": segs}).body.decode("utf-8") # 获取原始JSON字符串
|
|
|
|
|
+ data_dict = json.loads(content_data) # 解析为字典
|
|
|
|
|
+ # ✅ 添加新数据到shapes列表
|
|
|
|
|
+ data_dict.update({
|
|
|
|
|
+ "imageHeight": imginf[2],
|
|
|
|
|
+ "imagePath": imginf[0],
|
|
|
|
|
+ "imageWidth": imginf[1],
|
|
|
|
|
+ "version": "4dage_cul_rel_dis_det_0.0.1"
|
|
|
|
|
+ })
|
|
|
|
|
+
|
|
|
|
|
+ return JSONResponse(content=data_dict)
|
|
|
|
|
+ except Exception as e:
|
|
|
|
|
+ raise HTTPException(status_code=500, detail=f"处理图片时出错: {str(e)}")
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+if __name__ == "__main__":
|
|
|
|
|
+ uvicorn.run("seg_server:app", host="0.0.0.0", port=7000, reload=True)
|