# yolo-prg2 新增標籤在照片上

#!pip install ultralytics
import cv2
from ultralytics import YOLO
from google.colab.patches import cv2_imshow # Import the cv2_imshow function from google.colab.patches

# 加载YOLOv8模型
model = YOLO('yolov8n.pt')  # 或使用其他预训练模型

# 读取图片
image_path = '/content/table.png'
image = cv2.imread(image_path)

# 运行YOLOv8检测
results = model(image)

# Get the class names from the model
names = model.names

# 在图片上绘制检测结果
for result in results:
    boxes = result.boxes.xyxy.cpu().numpy().astype(int) 
    # Get class names and confidence scores as numpy arrays
    classNames = result.boxes.cls.cpu().numpy()
    conf = result.boxes.conf.cpu().numpy()
    for i, box in enumerate(boxes):
        cv2.rectangle(image, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2)
        # 使用类别索引获取类别名称
        label = f'{names[int(classNames[i])]}: {conf[i]:.2f}'
        cv2.putText(image, label, (box[0], box[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)

# 显示结果
cv2_imshow(image) # Use cv2_imshow instead of cv2.imshow
cv2.waitKey(0)
cv2.destroyAllWindows()

# 保存结果（可选）
cv2.imwrite('output_image.jpg', image)

print('检测完成，结果已保存。')