administrator/frigate
1
0
Fork 0
mirror of https://github.com/blakeblackshear/frigate.git synced 2026-02-04 18:25:22 +03:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

Black format

Browse Source
This commit is contained in:
Sergey Krashevich 2023-04-23 18:44:07 +03:00
parent 1a2f8c568e
commit aa6d99a67e
No known key found for this signature in database
GPG Key ID: 625171324E7D3856
1 changed files with 12 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
frigate/detectors/plugins/deepstack.py
View File

@ -17,10 +17,13 @@ DETECTOR_KEY = "deepstack"
class DeepstackDetectorConfig(BaseDetectorConfig): class DeepstackDetectorConfig(BaseDetectorConfig):
type: Literal[DETECTOR_KEY] type: Literal[DETECTOR_KEY]
api_url: str = Field(default="http://localhost:80/v1/vision/detection", title="DeepStack API URL") api_url: str = Field(
default="http://localhost:80/v1/vision/detection", title="DeepStack API URL"
)
api_timeout: float = Field(default=0.1, title="DeepStack API timeout (in seconds)") api_timeout: float = Field(default=0.1, title="DeepStack API timeout (in seconds)")
api_key: str = Field(default="", title="DeepStack API key (if required)") api_key: str = Field(default="", title="DeepStack API key (if required)")
class DeepStack(DetectionApi): class DeepStack(DetectionApi):
type_key = DETECTOR_KEY type_key = DETECTOR_KEY
@ -32,15 +35,15 @@ class DeepStack(DetectionApi):
self.h = detector_config.model.height self.h = detector_config.model.height
self.w = detector_config.model.width self.w = detector_config.model.width
def get_label_index(self, label_value): def get_label_index(self, label_value):
if label_value.lower() == 'truck': if label_value.lower() == "truck":
label_value = 'car' label_value = "car"
for index, value in self.labels.items(): for index, value in self.labels.items():
if value == label_value.lower(): if value == label_value.lower():
return index return index
return -1 return -1
def detect_raw(self, tensor_input): def detect_raw(self, tensor_input):
image_data = np.squeeze(tensor_input).astype(np.uint8) image_data = np.squeeze(tensor_input).astype(np.uint8)
image = Image.fromarray(image_data) image = Image.fromarray(image_data)
@ -48,10 +51,12 @@ class DeepStack(DetectionApi):
image.save(output, format="JPEG") image.save(output, format="JPEG")
image_bytes = output.getvalue() image_bytes = output.getvalue()
data = {"api_key": self.api_key} data = {"api_key": self.api_key}
response = requests.post(self.api_url, files={"image": image_bytes}, timeout=self.api_timeout) response = requests.post(
self.api_url, files={"image": image_bytes}, timeout=self.api_timeout
)
response_json = response.json() response_json = response.json()
detections = np.zeros((20, 6), np.float32) detections = np.zeros((20, 6), np.float32)
for i, detection in enumerate(response_json["predictions"]): for i, detection in enumerate(response_json["predictions"]):
logger.debug(f"Response: {detection}") logger.debug(f"Response: {detection}")
if detection["confidence"] < 0.4: if detection["confidence"] < 0.4: