michael/scrypted
1
0
Fork 0
mirror of https://github.com/koush/scrypted.git synced 2026-09-19 10:10:39 +01:00
Code Issues Actions Packages Projects Releases Wiki Activity

tensorflow-lite: use multiple tpu

Browse Source
This commit is contained in:
Koushik Dutta
2023-03-15 23:08:48 -07:00
parent 1c8ff2493b
commit 38ba31ca7d
5 changed files with 72 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
plugins/tensorflow-lite/package-lock.json generated
View File

@@ -1,12 +1,12 @@
{
"name": "@scrypted/tensorflow-lite",
"version": "0.0.112",
"version": "0.0.113",
"lockfileVersion": 2,
"requires": true,
"packages": {
"": {
"name": "@scrypted/tensorflow-lite",
"version": "0.0.112",
"version": "0.0.113",
"devDependencies": {
"@scrypted/sdk": "file:../../sdk"
}

2
plugins/tensorflow-lite/package.json
View File

@@ -44,5 +44,5 @@
"devDependencies": {
"@scrypted/sdk": "file:../../sdk"
},
"version": "0.0.112"
"version": "0.0.113"
}

12
plugins/tensorflow-lite/src/detect/__init__.py
View File

@@ -209,11 +209,11 @@ class DetectPlugin(scrypted_sdk.ScryptedDeviceBase, ObjectDetection):
async def run_detection_videoframe(self, videoFrame: scrypted_sdk.VideoFrame) -> ObjectsDetected:
pass
def run_detection_avframe(self, detection_session: DetectionSession, avframe, settings: Any, src_size, convert_to_src_size) -> Tuple[ObjectsDetected, Any]:
async def run_detection_avframe(self, detection_session: DetectionSession, avframe, settings: Any, src_size, convert_to_src_size) -> Tuple[ObjectsDetected, Any]:
pil: Image.Image = avframe.to_image()
return self.run_detection_image(detection_session, pil, settings, src_size, convert_to_src_size)
return await self.run_detection_image(detection_session, pil, settings, src_size, convert_to_src_size)
def run_detection_image(self, detection_session: DetectionSession, image: Image.Image, settings: Any, src_size, convert_to_src_size) -> Tuple[ObjectsDetected, Any]:
async def run_detection_image(self, detection_session: DetectionSession, image: Image.Image, settings: Any, src_size, convert_to_src_size) -> Tuple[ObjectsDetected, Any]:
pass
def run_detection_crop(self, detection_session: DetectionSession, sample: Any, settings: Any, src_size, convert_to_src_size, bounding_box: Tuple[float, float, float, float]) -> ObjectsDetected:
@@ -335,7 +335,7 @@ class DetectPlugin(scrypted_sdk.ScryptedDeviceBase, ObjectDetection):
finally:
detection_session.running = False
else:
return self.run_detection_jpeg(detection_session, bytes(await scrypted_sdk.mediaManager.convertMediaObjectToBuffer(mediaObject, 'image/jpeg')), settings)
return await self.run_detection_jpeg(detection_session, bytes(await scrypted_sdk.mediaManager.convertMediaObjectToBuffer(mediaObject, 'image/jpeg')), settings)
if not create:
# a detection session may have been created, but not started
@@ -479,7 +479,7 @@ class DetectPlugin(scrypted_sdk.ScryptedDeviceBase, ObjectDetection):
if not current_data:
raise Exception('no sample')
detection_result = self.run_detection_crop(
detection_result = await self.run_detection_crop(
detection_session, current_data, detection_session.settings, current_src_size, current_convert_to_src_size, boundingBox)
return detection_result['detections']
@@ -493,7 +493,7 @@ class DetectPlugin(scrypted_sdk.ScryptedDeviceBase, ObjectDetection):
first_frame = False
print("first frame received", detection_session.id)
detection_result, data = run_detection(
detection_result, data = await run_detection(
detection_session, sample, detection_session.settings, src_size, convert_to_src_size)
if detection_result:
detection_result['running'] = True

30
plugins/tensorflow-lite/src/predict/__init__.py
View File

@@ -250,13 +250,13 @@ class PredictPlugin(DetectPlugin, scrypted_sdk.BufferConverter, scrypted_sdk.Set
# print(detection_result)
return detection_result
def run_detection_jpeg(self, detection_session: PredictSession, image_bytes: bytes, settings: Any) -> ObjectsDetected:
async def run_detection_jpeg(self, detection_session: PredictSession, image_bytes: bytes, settings: Any) -> ObjectsDetected:
stream = io.BytesIO(image_bytes)
image = Image.open(stream)
if image.mode == 'RGBA':
image = image.convert('RGB')
detections, _ = self.run_detection_image(detection_session, image, settings, image.size)
detections, _ = await self.run_detection_image(detection_session, image, settings, image.size)
return detections
def get_detection_input_size(self, src_size):
@@ -269,7 +269,7 @@ class PredictPlugin(DetectPlugin, scrypted_sdk.BufferConverter, scrypted_sdk.Set
def get_input_size(self) -> Tuple[int, int]:
pass
def detect_once(self, input: Image.Image, settings: Any, src_size, cvss) -> ObjectsDetected:
async def detect_once(self, input: Image.Image, settings: Any, src_size, cvss) -> ObjectsDetected:
pass
async def run_detection_videoframe(self, videoFrame: scrypted_sdk.VideoFrame, settings: Any) -> ObjectsDetected:
@@ -288,7 +288,7 @@ class PredictPlugin(DetectPlugin, scrypted_sdk.BufferConverter, scrypted_sdk.Set
})
image = Image.frombuffer('RGB', (w, h), data)
try:
ret = self.detect_once(image, settings, src_size, cvss)
ret = await self.detect_once(image, settings, src_size, cvss)
return ret
finally:
image.close()
@@ -339,9 +339,9 @@ class PredictPlugin(DetectPlugin, scrypted_sdk.BufferConverter, scrypted_sdk.Set
def cvss2(point, normalize=False):
return point[0] / s + ow, point[1] / s + oh, True
ret1 = self.detect_once(first, settings, src_size, cvss1)
ret1 = await self.detect_once(first, settings, src_size, cvss1)
first.close()
ret2 = self.detect_once(second, settings, src_size, cvss2)
ret2 = await self.detect_once(second, settings, src_size, cvss2)
second.close()
two_intersect = intersect_rect(Rectangle(*first_crop), Rectangle(*second_crop))
@@ -374,7 +374,7 @@ class PredictPlugin(DetectPlugin, scrypted_sdk.BufferConverter, scrypted_sdk.Set
ret['detections'] = dedupe_detections(ret1['detections'] + ret2['detections'], is_same_detection=is_same_detection_middle)
return ret
def run_detection_image(self, detection_session: PredictSession, image: Image.Image, settings: Any, src_size, convert_to_src_size: Any = None, multipass_crop: Tuple[float, float, float, float] = None):
async def run_detection_image(self, detection_session: PredictSession, image: Image.Image, settings: Any, src_size, convert_to_src_size: Any = None, multipass_crop: Tuple[float, float, float, float] = None):
(w, h) = self.get_input_size() or image.size
(iw, ih) = image.size
@@ -448,7 +448,7 @@ class PredictPlugin(DetectPlugin, scrypted_sdk.BufferConverter, scrypted_sdk.Set
converted = convert_to_src_size(unscaled, normalize) if convert_to_src_size else (unscaled[0], unscaled[1], True)
return converted
ret = self.detect_once(input, settings, src_size, cvss)
ret = await self.detect_once(input, settings, src_size, cvss)
input.close()
detection_session.processed = detection_session.processed + 1
return ret, RawImage(image)
@@ -461,7 +461,7 @@ class PredictPlugin(DetectPlugin, scrypted_sdk.BufferConverter, scrypted_sdk.Set
converted = convert_to_src_size(point, normalize) if convert_to_src_size else (point[0], point[1], True)
return converted
ret = self.detect_once(image, settings, src_size, cvss)
ret = await self.detect_once(image, settings, src_size, cvss)
if detection_session:
detection_session.processed = detection_session.processed + 1
else:
@@ -483,11 +483,11 @@ class PredictPlugin(DetectPlugin, scrypted_sdk.BufferConverter, scrypted_sdk.Set
converted = convert_to_src_size(unscaled, normalize) if convert_to_src_size else (unscaled[0], unscaled[1], True)
return converted
ret1 = self.detect_once(first, settings, src_size, cvss1)
ret1 = await self.detect_once(first, settings, src_size, cvss1)
first.close()
if detection_session:
detection_session.processed = detection_session.processed + 1
ret2 = self.detect_once(second, settings, src_size, cvss2)
ret2 = await self.detect_once(second, settings, src_size, cvss2)
if detection_session:
detection_session.processed = detection_session.processed + 1
second.close()
@@ -576,11 +576,11 @@ class PredictPlugin(DetectPlugin, scrypted_sdk.BufferConverter, scrypted_sdk.Set
# print('untracked %s: %s' % (d['className'], d['score']))
def run_detection_crop(self, detection_session: DetectionSession, sample: RawImage, settings: Any, src_size, convert_to_src_size, bounding_box: Tuple[float, float, float, float]) -> ObjectsDetected:
(ret, _) = self.run_detection_image(detection_session, sample.image, settings, src_size, convert_to_src_size, bounding_box)
async def run_detection_crop(self, detection_session: DetectionSession, sample: RawImage, settings: Any, src_size, convert_to_src_size, bounding_box: Tuple[float, float, float, float]) -> ObjectsDetected:
(ret, _) = await self.run_detection_image(detection_session, sample.image, settings, src_size, convert_to_src_size, bounding_box)
return ret
def run_detection_gstsample(self, detection_session: PredictSession, gstsample, settings: Any, src_size, convert_to_src_size) -> Tuple[ObjectsDetected, Image.Image]:
async def run_detection_gstsample(self, detection_session: PredictSession, gstsample, settings: Any, src_size, convert_to_src_size) -> Tuple[ObjectsDetected, Image.Image]:
caps = gstsample.get_caps()
# can't trust the width value, compute the stride
height = caps.get_structure(0).get_value('height')
@@ -604,7 +604,7 @@ class PredictPlugin(DetectPlugin, scrypted_sdk.BufferConverter, scrypted_sdk.Set
gst_buffer.unmap(info)
try:
return self.run_detection_image(detection_session, image, settings, src_size, convert_to_src_size)
return await self.run_detection_image(detection_session, image, settings, src_size, convert_to_src_size)
except:
image.close()
traceback.print_exc()

69
plugins/tensorflow-lite/src/tflite/__init__.py
View File

@@ -3,7 +3,6 @@ import threading
from .common import *
from PIL import Image
from pycoral.adapters import detect
from pycoral.adapters.common import input_size
loaded_py_coral = False
try:
from pycoral.utils.edgetpu import list_edge_tpus
@@ -19,6 +18,9 @@ import scrypted_sdk
from scrypted_sdk.types import Setting
from typing import Any, Tuple
from predict import PredictPlugin
import concurrent.futures
import queue
import asyncio
def parse_label_contents(contents: str):
lines = contents.splitlines()
@@ -41,6 +43,9 @@ class TensorFlowLitePlugin(PredictPlugin, scrypted_sdk.BufferConverter, scrypted
labels_contents = scrypted_sdk.zip.open(
'fs/coco_labels.txt').read().decode('utf8')
self.labels = parse_label_contents(labels_contents)
self.interpreters = queue.Queue()
self.interpreter_count = 0
try:
edge_tpus = list_edge_tpus()
print('edge tpus', edge_tpus)
@@ -53,7 +58,21 @@ class TensorFlowLitePlugin(PredictPlugin, scrypted_sdk.BufferConverter, scrypted
'fs/mobilenet_ssd_v2_coco_quant_postprocess_edgetpu.tflite').read()
# face_model = scrypted_sdk.zip.open(
# 'fs/mobilenet_ssd_v2_face_quant_postprocess.tflite').read()
self.interpreter = make_interpreter(model)
for idx, edge_tpu in enumerate(edge_tpus):
try:
interpreter = make_interpreter(model, ":%s" % idx)
interpreter.allocate_tensors()
_, height, width, channels = interpreter.get_input_details()[
0]['shape']
self.input_details = int(width), int(height), int(channels)
self.interpreters.put(interpreter)
self.interpreter_count = self.interpreter_count + 1
print('added tpu %s' % (edge_tpu))
except Exception as e:
print('unable to use Coral Edge TPU', e)
if not self.interpreter_count:
raise Exception('all tpus failed to load')
# self.face_interpreter = make_interpreter(face_model)
except Exception as e:
print('unable to use Coral Edge TPU', e)
@@ -62,10 +81,16 @@ class TensorFlowLitePlugin(PredictPlugin, scrypted_sdk.BufferConverter, scrypted
'fs/mobilenet_ssd_v2_coco_quant_postprocess.tflite').read()
# face_model = scrypted_sdk.zip.open(
# 'fs/mobilenet_ssd_v2_face_quant_postprocess.tflite').read()
self.interpreter = tflite.Interpreter(model_content=model)
interpreter = tflite.Interpreter(model_content=model)
interpreter.allocate_tensors()
_, height, width, channels = interpreter.get_input_details()[
0]['shape']
self.input_details = int(width), int(height), int(channels)
self.interpreters.put(interpreter)
self.interpreter_count = self.interpreter_count + 1
# self.face_interpreter = make_interpreter(face_model)
self.interpreter.allocate_tensors()
self.mutex = threading.Lock()
self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=self.interpreter_count, thread_name_prefix="tflite", )
async def getSettings(self) -> list[Setting]:
ret = await super().getSettings()
@@ -83,30 +108,32 @@ class TensorFlowLitePlugin(PredictPlugin, scrypted_sdk.BufferConverter, scrypted
# width, height, channels
def get_input_details(self) -> Tuple[int, int, int]:
with self.mutex:
_, height, width, channels = self.interpreter.get_input_details()[
0]['shape']
return int(width), int(height), int(channels)
return self.input_details
def get_input_size(self) -> Tuple[int, int]:
w, h = input_size(self.interpreter)
return int(w), int(h)
return self.input_details[0:2]
def detect_once(self, input: Image.Image, settings: Any, src_size, cvss):
try:
with self.mutex:
async def detect_once(self, input: Image.Image, settings: Any, src_size, cvss):
def predict():
interpreter = self.interpreters.get()
try:
common.set_input(
self.interpreter, input)
interpreter, input)
scale = (1, 1)
# _, scale = common.set_resized_input(
# self.interpreter, cropped.size, lambda size: cropped.resize(size, Image.ANTIALIAS))
self.interpreter.invoke()
interpreter.invoke()
objs = detect.get_objects(
self.interpreter, score_threshold=.2, image_scale=scale)
except:
print('tensorflow-lite encountered an error while detecting. requesting plugin restart.')
self.requestRestart()
raise e
interpreter, score_threshold=.2, image_scale=scale)
return objs
except:
print('tensorflow-lite encountered an error while detecting. requesting plugin restart.')
self.requestRestart()
raise e
finally:
self.interpreters.put(interpreter)
objs = await asyncio.get_event_loop().run_in_executor(self.executor, predict)
allowList = settings.get('allowList', None) if settings else None
ret = self.create_detection_result(objs, src_size, allowList, cvss)