forked from wffjwbbf/ComDesignProject
34 lines
1004 B
Python
34 lines
1004 B
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,non_max_suppression
|
|
|
|
import param
|
|
|
|
|
|
class YOLOv7(object):
|
|
def __init__(self,option:param.Parameters) -> None:
|
|
# load FP32 model
|
|
self.opt = option
|
|
self.model = attempt_load(weights=self.opt.weights_yolo, device=self.opt.device)
|
|
|
|
def detect(self,imgs:torch.Tensor) -> list:
|
|
# img.size should be (Batch,3,H,W),every value should be range of [0,1]
|
|
# Both H and W % 64 == 0
|
|
# Initialize
|
|
model = self.model
|
|
opt = self.opt
|
|
# Inference
|
|
with torch.no_grad(): # Calculating gradients would cause a GPU memory leak
|
|
pred = model(imgs, 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
|
|
|
|
|
|
|
|
|