Spaces:
Running
on
Zero
Running
on
Zero
File size: 848 Bytes
e85fecb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 |
"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""
import torch
from ...core import register
__all__ = [
"YOLO",
]
@register()
class YOLO(torch.nn.Module):
__inject__ = [
"backbone",
"neck",
"head",
]
def __init__(self, backbone: torch.nn.Module, neck, head):
super().__init__()
self.backbone = backbone
self.neck = neck
self.head = head
def forward(self, x, **kwargs):
x = self.backbone(x)
x = self.neck(x)
x = self.head(x)
return x
def deploy(
self,
):
self.eval()
for m in self.modules():
if m is not self and hasattr(m, "deploy"):
m.deploy()
return self
|