You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
81 lines
2.4 KiB
81 lines
2.4 KiB
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
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|