3 posts tagged with "Python"

- orig_img: 原始图像数据
- orig_shape: 原始图像尺寸(高, 宽)
- path: 输入图像路径
- save_dir: 结果保存路径
- speed: 预测耗时信息

import os
from ultralytics import YOLO
import cv2
import os
import glob
import shutil

OBJECT_DETECTION_MODEL_PATH = './models/object_detection.onnx'
TASK_NAME = 'detect'

def generate_colors(names):
    colors = {}
    for name in names:
        hash_object = hashlib.md5(name.encode())
        hash_int = int(hash_object.hexdigest(), 16)
        b = (hash_int & 0xFF0000) >> 16
        g = (hash_int & 0x00FF00) >> 8
        r = hash_int & 0x0000FF
        colors[name] = (b, g, r)  # OpenCV 使用 BGR 顺序
    return colors

# 单张图像目标检测预测
def predict_single_image_by_detect(image_path, out_image_file):
    # 获取输出文件`out_image_path`文件所在的目录
    out_dir = os.path.dirname(out_image_path)
    os.makedirs(out_dir, exist_ok=True)

    image_list = [image_path]
    results = model(image_list)

    for result in results:
        boxes = result.boxes
        if boxes is None:
            cv2.imwrite(out_image_file, result.orig_img)
            continue
        boxes_data = boxes.data.cpu().numpy()
        names = result.names
        class_names = list(names.values())

        color_map = generate_colors(class_names)

        img = result.orig_img

        for box in boxes_data:
            x1, y1, x2, y2, score, class_id = box
            x1, y1, x2, y2 = map(int, [x1, y1, x2, y2])
            class_name = names[int(class_id)]
            color = color_map[class_name]
            cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
            label = f'{class_name} {score:.2f}'
            cv2.putText(img, label, (x1, max(y1 - 10, 0)),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)
        print(f"图像写入路径: {out_image_file}")
        cv2.imwrite(out_image_file, img)

if __name__ == '__main__':
    # 预测单张图像
    image_path = 'bus.jpg'
    out_image_path = image_path + '_predicted.jpg'
    predict_single_image_by_detect(image_path, out_image_path)

[x1, y1, x2, y2, score, class_id]
# x1, y1: 左上角坐标
# x2, y2: 右下角坐标
# score: 检测置信度
# class_id: 类别ID

import os
import hashlib
import cv2
import numpy as np
from ultralytics import YOLO
import glob
import shutil

SEGMENT_MODEL_PATH = "./models/segmentation.onnx"
TASK_NAME = 'segment'
model = YOLO(SEGMENT_MODEL_PATH, task=TASK_NAME)

# 单张图像的分割模型预测函数
def predict_single_image_by_segment(image_path, out_image_path):
    out_dir = os.path.dirname(out_image_path)
    os.makedirs(out_dir, exist_ok=True)

    results = model.predict(source=image_path)

    for result in results:
        if result.masks is None:
            cv2.imwrite(out_image_path, result.orig_img)
            continue
        masks = result.masks.data.cpu().numpy()
        boxes = result.boxes.data.cpu().numpy()
        label_map = result.names
        color_map = generate_colors(label_map.values())

        img_with_masks = result.orig_img.copy()

        for i, mask in enumerate(masks):
            mask = mask.astype(np.uint8)
            mask = cv2.resize(mask, (result.orig_shape[1], result.orig_shape[0]))

            color = np.random.randint(0, 255, (3,), dtype=np.uint8)
            colored_mask = np.zeros_like(result.orig_img, dtype=np.uint8)
            colored_mask[mask > 0] = color

            img_with_masks = cv2.addWeighted(img_with_masks, 1, colored_mask, 0.5, 0)

            box_data = boxes[i]
            x1, y1, x2, y2 = map(int, box_data[:4])
            class_name = label_map[int(box_data[5])]
            score = box_data[4]
            cv2.rectangle(img_with_masks, (x1, y1), (x2, y2), color_map[class_name], 2)
            label = f"{class_name}: {score:.4f}"
            cv2.putText(img_with_masks, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)

        cv2.imwrite(out_image_path, img_with_masks)
        print(f"Prediction saved to {out_image_path}")

if __name__ == '__main__':
    image_path = 'bus.jpg'
    out_image_path = image_path + '_segmented.jpg'
    predict_single_image_by_segment(image_path, out_image_path)

import cv2
from ultralytics import YOLO
import os

POSE_MODEL_PATH = './models/pose.onnx'
TASK_NAME = 'pose'
model = YOLO(POSE_MODEL_PATH, task=TASK_NAME)

def predict_single_image_by_pose(image_path, out_image_path):
    out_dir = os.path.dirname(out_image_path)
    os.makedirs(out_dir, exist_ok=True)

    results = model.predict(source=image_path)

    for result in results:
        if result.keypoints is None:
            continue
        if result.boxes is None:
            continue

        orig_img = result.orig_img
        keypoints = result.keypoints.data.cpu().numpy()
        boxes = result.boxes.data.cpu().numpy()

        for box_data, kpts in zip(boxes, keypoints):
            for keypoint in kpts:
                x, y, score = keypoint
                cv2.circle(orig_img, (int(x), int(y)), 3, (255, 0, 0), -1)

            for connection in skeleton:
                part_a, part_b = connection
                if kpts[part_a][2] > 0.5 and kpts[part_b][2] > 0.5:
                    x1, y1 = int(kpts[part_a][0]), int(kpts[part_a][1])
                    x2, y2 = int(kpts[part_b][0]), int(kpts[part_b][1])
                    cv2.line(orig_img, (x1, y1), (x2, y2), (0, 255, 255), 1)

        cv2.imwrite(out_image_path, orig_img)

if __name__ == '__main__':
    image_path = 'bus.jpg'
    out_image_path = image_path + '_posed.jpg'
    predict_single_image_by_pose(image_path, out_image_path)

[x, y, confidence]  # 每个关键点包含坐标和置信度

def main():
	print("Run Rye main")

def main():
	print("Hello Rye")