import cv2
import torch
import numpy as np
from PIL import Image
from ultralytics import YOLO


# 3个输入参数
img_path = 'img0002.jpg'
iou = 0.1
conf = 0.25

def split_image(img_path, size=(800, 800)):
    img = cv2.imread(img_path)
    height, width = img.shape[:2]
    rows = (height + size[1] - 1) // size[1]
    cols = (width + size[0] - 1) // size[0]
    img_list = []
    indexes = []
    for r in range(rows):
        for c in range(cols):
            y1 = r * size[1]
            y2 = min((r + 1) * size[1], height)
            x1 = c * size[0]
            x2 = min((c + 1) * size[0], width)
            split = img[y1:y2, x1:x2]
            img_list.append(split)
            indexes.append((r, c))
    return img_list, indexes, (height, width)


def combine_images(pred_imgs, indexes, size=(800, 800), img_shape=(3000, 4000)):
    combined_img = np.zeros((img_shape[0], img_shape[1], 3), dtype=np.uint8)
    for idx, (r, c) in enumerate(indexes):
        y1 = r * size[1]
        y2 = min((r + 1) * size[1], img_shape[0])
        x1 = c * size[0]
        x2 = min((c + 1) * size[0], img_shape[1])
        combined_img[y1:y2, x1:x2] = pred_imgs[idx][:y2 - y1, :x2 - x1]
    return combined_img

follicle_groups_detector = YOLO('follicle_groups.pt')
follicles_detector = YOLO('follicles.pt')

results = follicle_groups_detector(img_path, iou=iou, conf=conf)
for r in results:
    num_follicle_groups = len(r.boxes)
    im_array = r.plot()
    im = Image.fromarray(im_array[..., ::-1])
    im.save('results_1.jpg') #输出结果图1

img_list, indexes, (height, width) = split_image(img_path)
print(f"Number of image blocks: {len(img_list)}")
num_small_follicles = 0
num_big_follicles = 0
pred_imgs = []
for img in img_list:
    results = follicles_detector(img, iou=iou, conf=conf)
    for r in results:
        num_small_follicles += torch.sum(r.boxes.cls == 0).item()
        num_big_follicles += torch.sum(r.boxes.cls == 1).item()
        im_array = r.plot()
        pred_imgs.append(im_array)

# 输出的3个结果文本
print('毛囊群数量：', num_follicle_groups)
print('大毛囊数量：', num_big_follicles)
print('小毛囊数量：', num_small_follicles)

combined_img = combine_images(pred_imgs, indexes, size=(800, 800), img_shape=(height, width))
combined_image_pil = Image.fromarray(combined_img[..., ::-1])
combined_image_pil.save('results_2.jpg') #输出结果图2