import cv2
from matplotlib import pyplot as plt
import numpy as np
from ultralytics import YOLO
import torch

# Data
test01 = "data/16_3450.png"
test_image = test01

# Load a model
model = YOLO("weight/yolov8l-seg-pre100.onnx")  # load a custom model


with torch.no_grad():
    results = model(
        test_image,
        save=True,
        show_boxes=False,
        project="./result/",
        conf=0.5,
        retina_masks=False
    )


with torch.no_grad():
    for result in results:
        masks = result.masks.data
        boxes = result.boxes.data

    #ALL
    canvas = torch.any(masks, dim=0).int() * 255

    clss = boxes[:, 5]
    obj_indices = torch.where(clss == 4)

    # Cable
    obj_masks = masks[obj_indices]
    obj_mask = torch.any(obj_masks, dim=0).int() * 255
    # cropped_image = result.orig_img[obj_mask.cpu().numpy()]


fig, axs = plt.subplots(nrows=2, ncols=2, figsize=(27,15))
axs[0][0].imshow(result.orig_img)
axs[0][0].set_title("Orijinal Görüntü")

axs[0][1].imshow(canvas.cpu().numpy())
axs[0][1].set_title("Segmentasyon Maskesi")

mask = np.array(obj_mask.cpu().numpy())*255
cv2.imwrite("cable_mask.png", mask)
axs[1][0].imshow(obj_mask.cpu().numpy())
axs[1][0].set_title("Seçilen")

axs[1][1].imshow(result.plot())
axs[1][1].set_title("Sonuç")
plt.show()