# from ultralytics import YOLO
#
# model = YOLO(r"D:\ultralytics\runs\segment\train2\weights\best.pt")
# results = model.predict(source=r"D:\dataset\tooth\images", save=True,save_txt=False) # Display preds. Accepts all YOLO predict arguments
import os

from ultralytics import YOLO
import cv2
import numpy as np
import math

model = YOLO(r"D:\ultralytics\runs\segment\train2\weights\best.pt")
for file in os.listdir(r'D:\dataset\tooth\images'):
    # results = model.predict(source=r"D:\dataset\tooth\images\PX20240605_120714_0000_001642A7.png", save=False,save_txt=False)
    results = model.predict(source=r"D:\dataset\tooth\images\{}".format(file), save=False,save_txt=False)
    masks = results[0].masks.data
    boxes = results[0].boxes.data
    # Contours = []
    centers = []
    H = 0
    W = 0
    src = cv2.imread(r"D:\dataset\tooth\images\{}".format(file))
    # src = cv2.imread(r"D:\dataset\tooth\images\PX20240605_120714_0000_001642A7.png")
    occlusal = 0
    occlusalH = 0

    for index,mask in enumerate(masks):
        box = boxes[index]
        cls = box[-1].cpu().detach().numpy()
        img = mask.cpu().detach().numpy()
        H,W = img.shape
        # background = np.zeros_like(img)
        # index_obj = img > 0
        # background[index_obj] = 255
        # cv2.imwrite('mask.jpg',background)
        # background = cv2.imread('mask.jpg',0)
        # contours, hierarchy = cv2.findContours(background, cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)
        # contourObj = contours[0]
        # for contour in contours:
        #     if len(contour) > len(contourObj):
        #         contourObj = contour
        x1 = int(box[0])
        y1 = int(box[1])
        x2 = int(box[2])
        y2 = int(box[3])
        if cls == 1:
            w = x2 - x1
            h = y2 - y1
            occlusalH += h
            occlusal += 1
        # Contours.append(contourObj)
        if cls == 5:
            src = cv2.rectangle(src, (x1, y1), (x2, y2), (0, 0, 255), 2)
            src = cv2.putText(src, "badTooth", (x1, y1 - 2), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255,0), 2)
        else:
            src = cv2.rectangle(src, (x1, y1), (x2, y2), (0, 255, 0), 2)
        centers.append((x1,y1,index,x2,y2,cls))
    if occlusal != 0:
        occlusalMeanH = occlusalH / occlusal
    else:
        occlusalMeanH = 0

    sortedCenters = sorted(centers, key=lambda x: x[1])
    centersTop = []
    centersBottom = []
    Max = 0
    boundary = False
    centersTop.append(sortedCenters[0])
    for index in range(len(sortedCenters) - 1):
        div = sortedCenters[index + 1][1] - sortedCenters[index][1]
        if div > Max:
            Max = div
    for index in range(len(sortedCenters) - 1):
        div = sortedCenters[index + 1][1] - sortedCenters[index][1]
        if div == Max:
            boundary = True
        if boundary:
            centersBottom.append(sortedCenters[index + 1])
        else:
            centersTop.append(sortedCenters[index + 1])
    centersTop = sorted(centersTop, key=lambda x: x[0])
    centersBottom = sorted(centersBottom, key=lambda x: x[0])

    indexs = []
    for center in centersTop:
        indexs.append(center[2])
    for center in centersBottom:
        indexs.append(center[2])

    # src = cv2.resize(src,(W,H))
    # mask = np.zeros_like(src)

    tlOpsigenes = False
    blOpsigenes = False
    trOpsigenes = False
    brOpsigenes = False

    for i in range(len(centersTop) - 1):
        index1 = indexs[i]
        index2 = indexs[i + 1]
        X2 = int(boxes[index2][0])
        X1 = int(boxes[index1][2])
        if X2 - X1 > 20:
            x1 = int(boxes[index1][2])
            y1 = int(boxes[index1][1])
            x2 = int(boxes[index2][0])
            y2 = int(boxes[index2][3])
            src = cv2.rectangle(src,(x1,y1),(x2,y2), (0, 0, 255), 2)
            src = cv2.putText(src,"agomphosis",(x1,y1 - 2),cv2.FONT_HERSHEY_SIMPLEX,1,(0, 255, 255),2)
        if i == 0:
            w = centersTop[i][3] - centersTop[i][0]
            h = centersTop[i][4] - centersTop[i][1]
            if w / h > 1.05 and occlusalMeanH > h and centersTop[i][5] == 1:
                src = cv2.rectangle(src, (centersTop[i][0],centersTop[i][1]), (centersTop[i][3],centersTop[i][4]), (0, 0, 255), 2)
                src = cv2.putText(src, "opsigenes", (centersTop[i][0],centersTop[i][1] - 2), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2)
                tlOpsigenes = True
        if i == len(centersTop) - 2:
            w = centersTop[i + 1][3] - centersTop[i + 1][0]
            h = centersTop[i + 1][4] - centersTop[i + 1][1]
            if w / h > 1.05 and occlusalMeanH > h and centersTop[i][5] == 1:
                src = cv2.rectangle(src, (centersTop[i + 1][0],centersTop[i + 1][1]), (centersTop[i + 1][3],centersTop[i + 1][4]), (0, 0, 255), 2)
                src = cv2.putText(src, "opsigenes", (centersTop[i + 1][0], centersTop[i + 1][1] - 2), cv2.FONT_HERSHEY_SIMPLEX,1, (0, 255, 255), 2)
                trOpsigenes = True

    for i in range(len(centersBottom) - 1):
        index1 = indexs[i + len(centersTop)]
        index2 = indexs[i + 1 + len(centersTop)]
        X2 = int(boxes[index2][0])
        X1 = int(boxes[index1][2])
        if X2 - X1 > 20:
            x1 = int(boxes[index1][2])
            y1 = int(boxes[index1][1])
            x2 = int(boxes[index2][0])
            y2 = int(boxes[index2][3])
            src = cv2.rectangle(src, (x1, y1), (x2, y2), (0, 0, 255), 2)
            src = cv2.putText(src, "agomphosis", (x1, y1 - 2), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2)
        if i == 0:
            w = centersBottom[i][3] - centersBottom[i][0]
            h = centersBottom[i][4] - centersBottom[i][1]
            if w / h > 1.05 and occlusalMeanH > h and centersBottom[i][5] == 1:
                src = cv2.rectangle(src, (centersBottom[i][0],centersBottom[i][1]), (centersBottom[i][3],centersBottom[i][4]), (0, 0, 255), 2)
                src = cv2.putText(src, "opsigenes", (centersBottom[i][0],centersBottom[i][1] - 2), cv2.FONT_HERSHEY_SIMPLEX,1, (0, 255, 255), 2)
                blOpsigenes  =True

        if i == len(centersBottom) - 2:
            w = centersBottom[i + 1][3] - centersBottom[i + 1][0]
            h = centersBottom[i + 1][4] - centersBottom[i + 1][1]
            if w / h > 1.05 and occlusalMeanH > h and centersBottom[i][5] == 1:
                src = cv2.rectangle(src, (centersBottom[i + 1][0], centersBottom[i + 1][1]),(centersBottom[i + 1][3], centersBottom[i + 1][4]), (0, 0, 255), 2)
                src = cv2.putText(src, "opsigenes", (centersBottom[i + 1][0], centersBottom[i + 1][1] - 2),cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2)
                brOpsigenes  =True

    topLeftX1 = centersTop[0][0]
    topLeftY1 = centersTop[0][1]

    topRightX1 = centersTop[-1][0]
    topRightY1 = centersTop[-1][1]

    bottomLeftX1 = centersBottom[0][0]
    bottomLeftY1 = centersBottom[0][1]

    bottomRightX1 = centersBottom[-1][0]
    bottomRightY1 = centersBottom[-1][1]

    leftMinX = min(topLeftX1,bottomLeftX1)
    leftMaxX = max(topLeftX1,bottomLeftX1)

    rightMinX = min(topRightX1,bottomRightX1)
    rightMaxX = max(topRightX1,bottomRightX1)

    leftMinY = min(topLeftY1, bottomLeftY1)
    leftMaxY = max(topLeftY1, bottomLeftY1)

    rightMinY = min(topRightY1, bottomRightY1)
    rightMaxY = max(topRightY1, bottomRightY1)

    leftDiv = leftMaxX - leftMinX
    rightDiv = rightMaxX - rightMinX

    solo = False
    if len(centersTop) == 1:
        if centersBottom[0][1] - centersTop[0][1] < 100:
            solo = True
    if solo:
        cv2.rectangle(src, (centersBottom[0][0] - 100, centersBottom[0][1] - 300), (centersBottom[len(centersBottom) - 1][3] + 100,centersBottom[len(centersBottom) - 1][4] - 300), (0, 0, 255),2)
        src = cv2.putText(src, "agomphosis", (centersBottom[0][0] - 100, centersBottom[0][1] - 302), cv2.FONT_HERSHEY_SIMPLEX, 1,(0, 255, 255), 2)

    if leftDiv > 90:
        if bottomLeftX1 < topLeftX1 and blOpsigenes == False and solo == False:#左上
            cv2.rectangle(src,(centersBottom[0][0],centersTop[0][1]),(centersTop[0][0],centersTop[0][4]),(0,0,255),2)
            src = cv2.putText(src, "agomphosis", (centersBottom[0][0],centersTop[0][1] - 2), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2)
        if bottomLeftX1 > topLeftX1 and tlOpsigenes == False and solo == False:#左下
            cv2.rectangle(src, (centersTop[0][0], centersBottom[0][1]), (centersBottom[0][0], centersBottom[0][4]),(0, 0, 255), 2)
            src = cv2.putText(src, "agomphosis", (centersTop[0][0], centersBottom[0][1] - 2), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2)
    if rightDiv > 90:
        if bottomRightX1 > topRightX1 and brOpsigenes == False and solo == False:#右上
            cv2.rectangle(src, (centersTop[-1][3],centersTop[-1][1]), (centersBottom[-1][3],centersTop[-1][4]), (0, 0, 255), 2)
            src = cv2.putText(src, "agomphosis", (centersTop[-1][3],centersTop[-1][1] - 2), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2)
        if bottomRightX1 < topRightX1 and trOpsigenes == False and solo == False:#右下
            cv2.rectangle(src, (centersBottom[-1][3],centersBottom[-1][1]), (centersTop[-1][3], centersBottom[-1][4]),(0, 0, 255), 2)
            src = cv2.putText(src, "agomphosis", (centersBottom[-1][3],centersBottom[-1][1] - 2), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2)

    cv2.namedWindow('0', 0)
    cv2.imshow('0',src)
    cv2.waitKey(0)

# for index in indexs:
#     print()
    # cv2.drawContours(src,Contours,index, (0, 0, 255), 1)