forked from wffjwbbf/ComDesignProject
76 lines
2.7 KiB
Python
76 lines
2.7 KiB
Python
import cv2
|
|
import torch
|
|
import numpy as np
|
|
|
|
from torchvision import transforms
|
|
|
|
from models.experimental import attempt_load
|
|
from utils.general import check_img_size, check_requirements, check_imshow, non_max_suppression, apply_classifier, \
|
|
scale_coords, xyxy2xywh,increment_path
|
|
from utils.plots import plot_one_box
|
|
|
|
import param
|
|
|
|
#Load the model as you don't want load weights every time
|
|
def creat_model(opt):
|
|
# load FP32 model
|
|
return attempt_load(weights=opt.weights, device=opt.device)
|
|
|
|
# img.size should be (Batch,3,H,W),every value should be range of [0,1],
|
|
# (Please normlize it by mean=(0,),std=(255,),)
|
|
# opt should be the class from module "param"
|
|
def detect(model,img:torch.Tensor,opt=param.Parameters()):
|
|
# Initialize
|
|
imgsz = opt.img_size
|
|
|
|
# Load model
|
|
stride = int(model.stride.max()) # model stride
|
|
imgsz = check_img_size(imgsz, s=stride) # check img_size
|
|
|
|
# Inference
|
|
with torch.no_grad(): # Calculating gradients would cause a GPU memory leak
|
|
pred = model(img, augment=opt.augment)[0]
|
|
|
|
# Apply NMS
|
|
pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
|
|
|
|
return pred
|
|
|
|
|
|
def get_img_test(path_src:str):
|
|
raw_img = cv2.imread(path_src)
|
|
raw_transform = transforms.Compose([transforms.ToPILImage(),
|
|
transforms.Resize((360,640)),
|
|
transforms.Pad((0,(640-360)//2)),
|
|
transforms.ToTensor()])
|
|
return raw_transform(raw_img).unsqueeze(dim=0)
|
|
|
|
def opencv_box_plot(img:cv2.Mat,pred_img:np.array):
|
|
pred_img = pred_img.astype(np.uint)
|
|
print(type(img))
|
|
for info in pred_img:
|
|
x1,y1,x2,y2 = info[0],info[1],info[2],info[3]
|
|
cv2.rectangle(img,(x1,y1-140),(x2,y2-140),(255,0,0),2)
|
|
print(type(img))
|
|
cv2.imshow("test",img)
|
|
cv2.waitKey(0)
|
|
cv2.imwrite("result.png",img)
|
|
|
|
|
|
# example
|
|
if __name__ == '__main__':
|
|
my_option = param.Parameters()
|
|
model = creat_model(my_option) # load model in head
|
|
with torch.no_grad():
|
|
img = get_img_test("test.png") # get image as torch.Tensor with size of [1,1,640,640]
|
|
print(f"img size is {img.size()}")
|
|
result = detect(model,img,my_option) # get the sequence of result
|
|
print(f"result is :{result}")
|
|
result = result[0].detach().cpu().numpy() # the single img is index 0
|
|
img_src = cv2.imread("test.png") # re-read for imshow
|
|
img_src = cv2.resize(img_src,(640,360))
|
|
opencv_box_plot(img_src,result) # show the result in picture
|
|
|
|
|
|
|