22 lines
744 B
Python
22 lines
744 B
Python
|
import torch
|
||
|
|
from algorithm.yolov5.models.common import DetectMultiBackend
|
||
|
|
import os
|
||
|
|
from algorithm.yolov5.models.yolo import Detect, Model,RotationDetect
|
||
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
|
|
|
||
|
|
# model_state_dict = torch.load('weight/remote_sensing/oriented.pt')
|
||
|
|
model = DetectMultiBackend(weights='weight/remote_sensing/oriented.pt', device=device, dnn=True)
|
||
|
|
# model.load_state_dict(model_state_dict)
|
||
|
|
# model = torch.hub.load((os.getcwd()) + "/algorithm/yolov5", 'custom', source='local', path='./weight//remote_sensing/oriented.pt', force_reload=True)
|
||
|
|
print(RotationDetect())
|
||
|
|
for m in model.modules():
|
||
|
|
t = type(m)
|
||
|
|
print(t)
|
||
|
|
if Detect:
|
||
|
|
print(1)
|
||
|
|
m = RotationDetect()
|
||
|
|
|
||
|
|
# print(model)
|
||
|
|
|
||
|
|
|