【深度学习】使用 YOLOv9+SAM 进行动态物体检测和分割

from google.colab import drivedrive.mount('/content/drive')
or
%cd {HOME}/!pip install -q roboflow
from roboflow import Roboflowrf = Roboflow(api_key="YOUR API KEY")project = rf.workspace("roboflow-100").project("construction-safety-gsnvb")dataset = project.version(2).download("yolov7")

!git clone https://github.com/SkalskiP/yolov9.git%cd yolov9!pip3 install -r requirements.txt -q

import osHOME = os.getcwd()print(HOME)

!mkdir -p {HOME}/weights !wget -P {HOME}/weights -q https://github.com/WongKinYiu/yolov9/releases/download/v0 .1 /yolov9-c.pt !wget -P {HOME}/weights -q https://github.com/WongKinYiu/yolov9/releases/download/v0 .1 /yolov9-e.pt !wget -P {HOME}/weights -q https://github.com/WongKinYiu/yolov9/releases/download/v0 .1 /gelan-c.pt !wget -P {HOME}/weights -q https://github.com/WongKinYiu/yolov9/releases/download/v0 .1 /gelan-e.pt

!python detect.py --weights {HOME}/weights/gelan-c.pt --conf 0.1 --source /content/drive/MyDrive/data/image9.jpeg --device 0 --save-txt --save-conf

!mkdir -p {HOME}/data!wget -P {HOME}/data -q /content/drive/MyDrive/data/image9.jpegSOURCE_IMAGE_PATH = f"{HOME}/image9.jpeg"

import matplotlib.patches as patchesfrom matplotlib import pyplot as pltimport numpy as npimport yaml
with open('/content/yolov9/data/coco.yaml', 'r') as file:    coco_data = yaml.safe_load(file)    class_names = coco_data['names']
for class_id, bbox, conf in zip(class_ids, bboxes, conf_scores):    class_name = class_names[class_id]# print(f'Class ID: {class_id}, Class Name: {class_name}, BBox coordinates: {bbox}')
color_map = {}for class_id in class_ids:    color_map[class_id] = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
def show_mask(mask, ax, color):    h, w = mask.shape[-2:]    mask_image = mask.reshape(h, w, 1) * np.array(color).reshape(1, 1, -1)    ax.imshow(mask_image)
def show_box(box, label, conf_score, color, ax):    x0, y0 = box[0], box[1]    w, h = box[2] - box[0], box[3] - box[1]    rect = plt.Rectangle((x0, y0), w, h, edgecolor=color, facecolor='none', lw=2)    ax.add_patch(rect)
    label_offset = 10
# Construct the label with the class name and confidence score    label_text = f'{label} {conf_score:.2f}'
    ax.text(x0, y0 - label_offset, label_text, color='black', fontsize=10, va='top', ha='left',


    
            bbox=dict(facecolor=color, alpha=0.7, edgecolor='none', boxstyle='square,pad=0.4'))
plt.figure(figsize=(10, 10))ax = plt.gca()plt.imshow(image)
# Display and process each bounding box with the corresponding maskfor class_id, bbox in zip(class_ids, bboxes):    class_name = class_names[class_id]# print(f'Class ID: {class_id}, Class Name: {class_name}, BBox coordinates: {bbox}')
    color = color_map[class_id]    input_box = np.array(bbox)
# Generate the mask for the current bounding box    masks, _, _ = predictor.predict(        point_coords=None,        point_labels=None,        box=input_box,        multimask_output=False,    )
    show_mask(masks[0], ax, color=color)    show_box(bbox, class_name, conf, color, ax)
# Show the final plotplt.axis('off')plt.show()

!pip install ultralyticsfrom ultralytics import YOLO

!pip 安装'git+https://github.com/facebookresearch/segment-anything.git'
 !wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth

import cv2


    

# Specify the path to your imageimage_path = '/content/drive/MyDrive/data/image9.jpeg'
# Read the image to get its dimensionsimage = cv2.imread(image_path)image_height, image_width, _ = image.shape
detections_path = '/content/yolov9/runs/detect/exp/labels/image9.txt'
bboxes = []class_ids = []conf_scores = []
with open(detections_path, 'r') as file:for line in file:components = line.split()class_id = int(components[0])confidence = float(components[5])cx, cy, w, h = [float(x) for x in components[1:5]]
        # Convert from normalized [0, 1] to image scalecx *= image_widthcy *= image_heightw *= image_widthh *= image_height
        # Convert the center x, y, width, and height to xmin, ymin, xmax, ymaxxmin = cx - w / 2ymin = cy - h / 2xmax = cx + w / 2ymax = cy + h / 2
class_ids.append(class_id)bboxes.append((xmin, ymin, xmax, ymax))conf_scores.append(confidence)
# Display the resultsfor class_id, bbox, conf in zip(class_ids, bboxes, conf_scores):print(f'Class ID: {class_id}, Confidence: {conf:.2f}, BBox coordinates: {bbox}')

from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictorsam_checkpoint = "/content/yolov9/sam_vit_h_4b8939.pth"model_type = "vit_h"sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)predictor = SamPredictor(sam)

import cv2image = cv2.cvtColor(cv2.imread('/content/drive/MyDrive/data/image9.jpeg'), cv2.COLOR_BGR2RGB)predictor.set_image(image)

aggregate_mask = np.zeros(image.shape[:2], dtype=np.uint8)
# Generate and accumulate masks for all bounding boxesfor bbox in bboxes:input_box = np.array(bbox).reshape(1, 4)masks, _, _ = predictor.predict(point_coords=None,point_labels=None,box=input_box,multimask_output=False,)aggregate_mask = np.where(masks[0] > 0.5, 1, aggregate_mask)
# Convert the aggregate segmentation mask to a binary maskbinary_mask = np.where(aggregate_mask == 1, 1, 0)
# Create a white background with the same size as the imagewhite_background = np.ones_like(image) * 255
# Apply the binary mask to the original image# Where the binary mask is 0 (background), use white_background; otherwise, use the original imagenew_image = white_background * (1 - binary_mask[..., np.newaxis]) + image * binary_mask[..., np.newaxis]
# Display the new image with the detections and white backgroundplt.figure(figsize=(10, 10))plt.imshow(new_image.astype(np.uint8))


    
plt.axis('off')plt.show()