fix formatting

This commit is contained in:
MarcA711 2024-05-17 16:15:39 +00:00
parent 6c5b44e3f8
commit a92bd60742

View File

@ -1,8 +1,10 @@
import logging import logging
from frigate.detectors.detector_config import ModelTypeEnum
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
import numpy as np import numpy as np
from frigate.detectors.detector_config import ModelTypeEnum
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -54,20 +56,24 @@ class DetectionApi(ABC):
boxes = np.transpose( boxes = np.transpose(
np.vstack( np.vstack(
( (
boxes[:, 1]/self.height, boxes[:, 1] / self.height,
boxes[:, 0]/self.width, boxes[:, 0] / self.width,
boxes[:, 3]/self.height, boxes[:, 3] / self.height,
boxes[:, 2]/self.width, boxes[:, 2] / self.width,
) )
) )
) )
results = np.hstack((class_ids[..., np.newaxis], scores[..., np.newaxis], boxes)) results = np.hstack(
(class_ids[..., np.newaxis], scores[..., np.newaxis], boxes)
)
return np.resize(results, (20,6)) return np.resize(results, (20, 6))
def post_process(self, output): def post_process(self, output):
if self.detector_config.model.model_type == ModelTypeEnum.yolonas: if self.detector_config.model.model_type == ModelTypeEnum.yolonas:
return self.yolonas(output) return self.yolonas(output)
else: else:
raise ValueError(f'Model type "{self.detector_config.model.model_type}" is currently not supported.') raise ValueError(
f'Model type "{self.detector_config.model.model_type}" is currently not supported.'
)