pytorch 我们如何提取yolov5检测图像的xyxy坐标并对其进行裁剪?

pcww981p  于 2022-12-29  发布在  其他
关注(0)|答案(1)|浏览(275)
import cv2
import numpy as np
import torch
import pandas
import os

# Load the trained model
model = torch.hub.load('ultralytics/yolov5', 'custom', path='C:/Users/ASUS/Desktop/yolov5-fastapi-main/model/best.pt',force_reload=True) 

def detect_logo(image):
    result = model(image)
    labels, cord_thres = result.xyxyn[0][:, -1].numpy(), result.xyxyn[0][:, :-1].numpy()
    cropped_images = []
    for i,box in enumerate(cord_thres):
        x_min = int(box[0])
        y_min = int(box[1])
        x_max = int(box[2])
        y_max = int(box[3])
        image2 = image.astype(np.uint8)
        cropped_image = image2[y_min:y_max,x_min:x_max]
        cropped_images.append(cropped_image)
        
        cv2.imwrite("image.jpg", cropped_images) 
        
        

    return cropped_images

# Read the image and convert it to a NumPy array
image = cv2.imread("C:/Users/ASUS/Desktop/yolov5-fastapi-main/aetna_003.png")

# Detect and crop the coordinates of the detected logos
cropped_images = detect_logo(image)

我想通过提取坐标来裁剪我在yolov5中检测到的图像,它正在工作,但它没有保存图像,错误在哪里,你能帮忙吗?

qnzebej0

qnzebej01#

我解决了这个问题,下面的代码是工作.

import cv2
import numpy as np
import torch

model = torch.hub.load('ultralytics/yolov5', 'custom', path='C:/Users/ASUS/Desktop/yolov5-fastapi-main/model/best.pt',force_reload=True)
def get_text(image_path, threshold):
    Y, X,d = image_path.shape
    result = model(image_path)
    df = result.pandas().xyxy[0]
    df = df[df['confidence']>=threshold]
    df = df.drop(['class', 'name','confidence'], axis=1)
    print(df.head())
    output = []
    values = df.values.astype(int)
    for xmin,ymin,xmax,ymax in df.values.astype(int):
        xmin = max(0, xmin)
        ymin = max(0, ymin)
        xmax = min(X, xmax)
        ymax = min(Y, ymax)
        crop_img = image_path[ymin:ymax, xmin:xmax]
        
        cv2.imwrite("C:/Users/ASUS/Desktop/yolov5-fastapi-main/image.jpg",crop_img)
        return output,image_path
image =cv2.imread('C:/Users/ASUS/Desktop/yolov5-fastapi-main/papara_003.png')
detection= get_text(image, 0.5)

相关问题