thebears/deepstream_yolo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add RT-DETR

Browse Source
This commit is contained in:
Marcos Luciano
2023-11-01 18:42:56 -03:00
parent 1177624dd2
commit 5af9da189d
1 changed files with 0 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
utils/export_rtdetr_pytorch.py
View File

@@ -22,17 +22,6 @@ class DeepStreamOutput(nn.Module):
return boxes, scores, classes return boxes, scores, classes
class DeepStreamInput(nn.Module):
def __init__(self, img_size, device):
self.img_size = img_size
self.device = device
super().__init__()
def forward(self, x):
size = torch.tensor([[*self.img_size]]).to(self.device)
return x, size
def suppress_warnings(): def suppress_warnings():
warnings.filterwarnings('ignore', category=torch.jit.TracerWarning) warnings.filterwarnings('ignore', category=torch.jit.TracerWarning)
warnings.filterwarnings('ignore', category=UserWarning) warnings.filterwarnings('ignore', category=UserWarning)