Spaces:
Sleeping
Sleeping
File size: 1,450 Bytes
e1832f4 |
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 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
# Mikel Broström 🔥 Yolo Tracking 🧾 AGPL-3.0 license
from pathlib import Path
import numpy as np
import torch
from ultralytics.engine.results import Results
from abc import ABC, abstractmethod
class YoloInterface(ABC):
@abstractmethod
def __call__(self, im):
pass
@abstractmethod
def preprocess(self, ims):
pass
@abstractmethod
def postprocess(self, preds):
pass
def get_scaling_factors(self, im, im0):
# im to im0 factor for predictions
im0_w = im0.shape[1]
im0_h = im0.shape[0]
im_w = im.shape[2]
im_h = im.shape[1]
w_r = im0_w / im_w
h_r = im0_h / im_h
return im_w, im_h, w_r, h_r
def scale_and_clip(self, preds, im_w, im_h, w_r, h_r):
# scale bboxes to original image
preds[:, [0, 2]] = preds[:, [0, 2]] * self.w_r
preds[:, [1, 3]] = preds[:, [1, 3]] * self.h_r
if not isinstance(preds, (torch.Tensor)):
preds = torch.from_numpy(preds)
preds[:, [0, 2]] = torch.clip(preds[:, [0, 2]], min=0) # max=im_w
preds[:, [1, 3]] = torch.clip(preds[:, [1, 3]], min=0) # max=im_h
return preds
def get_model_from_weigths(self, l, model):
model_type = None
for key in l:
if Path(key).stem in str(model.name):
model_type = str(Path(key).with_suffix(''))
break
return model_type
|