mirror of
https://github.com/koush/scrypted.git
synced 2026-09-19 18:20:40 +01:00
615 lines
22 KiB
Python
615 lines
22 KiB
Python
from __future__ import annotations
|
|
from scrypted_sdk.types import ObjectDetectionResult, ObjectsDetected, Setting
|
|
import io
|
|
from PIL import Image
|
|
import re
|
|
import scrypted_sdk
|
|
from typing import Any, List, Tuple, Mapping
|
|
import asyncio
|
|
import time
|
|
from .rectangle import Rectangle, intersect_area, intersect_rect, to_bounding_box, from_bounding_box, combine_rect
|
|
|
|
from detect import DetectionSession, DetectPlugin
|
|
|
|
from .sort_oh import tracker
|
|
import numpy as np
|
|
import traceback
|
|
|
|
try:
|
|
from gi.repository import Gst
|
|
except:
|
|
pass
|
|
|
|
class PredictSession(DetectionSession):
|
|
image: Image.Image
|
|
tracker: sort_oh.tracker.Sort_OH
|
|
|
|
def __init__(self, start_time: float) -> None:
|
|
super().__init__()
|
|
self.image = None
|
|
self.processed = 0
|
|
self.start_time = start_time
|
|
self.tracker = None
|
|
|
|
def parse_label_contents(contents: str):
|
|
lines = contents.splitlines()
|
|
ret = {}
|
|
for row_number, content in enumerate(lines):
|
|
pair = re.split(r'[:\s]+', content.strip(), maxsplit=1)
|
|
if len(pair) == 2 and pair[0].strip().isdigit():
|
|
ret[int(pair[0])] = pair[1].strip()
|
|
else:
|
|
ret[row_number] = content.strip()
|
|
return ret
|
|
|
|
|
|
class RawImage:
|
|
jpegMediaObject: scrypted_sdk.MediaObject
|
|
|
|
def __init__(self, image: Image.Image):
|
|
self.image = image
|
|
self.jpegMediaObject = None
|
|
|
|
def is_same_box(bb1, bb2, threshold = .7):
|
|
r1 = from_bounding_box(bb1)
|
|
r2 = from_bounding_box(bb2)
|
|
ia = intersect_area(r1, r2)
|
|
|
|
if not ia:
|
|
return False, None
|
|
|
|
a1 = bb1[2] * bb1[3]
|
|
a2 = bb2[2] * bb2[3]
|
|
|
|
# if area intersect area is too small, these are different boxes
|
|
if ia / a1 < threshold or ia / a2 < threshold:
|
|
return False, None
|
|
|
|
l = min(bb1[0], bb2[0])
|
|
t = min(bb1[1], bb2[1])
|
|
r = max(bb1[0] + bb1[2], bb2[0] + bb2[2])
|
|
b = max(bb1[1] + bb1[3], bb2[1] + bb2[3])
|
|
|
|
w = r - l
|
|
h = b - t
|
|
|
|
return True, (l, t, w, h)
|
|
|
|
def is_same_detection(d1: ObjectDetectionResult, d2: ObjectDetectionResult):
|
|
if d1['className'] != d2['className']:
|
|
return False, None
|
|
|
|
return is_same_box(d1['boundingBox'], d2['boundingBox'])
|
|
|
|
def dedupe_detections(input: List[ObjectDetectionResult], is_same_detection = is_same_detection):
|
|
input = input.copy()
|
|
detections = []
|
|
while len(input):
|
|
d = input.pop()
|
|
found = False
|
|
for c in detections:
|
|
same, box = is_same_detection(d, c)
|
|
if same:
|
|
# encompass this box and score
|
|
d['boundingBox'] = box
|
|
d['score'] = max(d['score'], c['score'])
|
|
# remove from current detections list
|
|
detections = list(filter(lambda r: r != c, detections))
|
|
# run dedupe again with this new larger item
|
|
input.append(d)
|
|
found = True
|
|
break
|
|
|
|
if not found:
|
|
detections.append(d)
|
|
return detections
|
|
|
|
class Prediction:
|
|
def __init__(self, id: int, score: float, bbox: Tuple[float, float, float, float]):
|
|
self.id = id
|
|
self.score = score
|
|
self.bbox = bbox
|
|
|
|
class PredictPlugin(DetectPlugin, scrypted_sdk.BufferConverter, scrypted_sdk.Settings):
|
|
labels: dict
|
|
|
|
def __init__(self, PLUGIN_MIME_TYPE: str, nativeId: str | None = None):
|
|
super().__init__(nativeId=nativeId)
|
|
|
|
self.fromMimeType = PLUGIN_MIME_TYPE
|
|
self.toMimeType = scrypted_sdk.ScryptedMimeTypes.MediaObject.value
|
|
|
|
self.crop = False
|
|
self.trackers: Mapping[str, tracker.Sort_OH] = {}
|
|
|
|
# periodic restart because there seems to be leaks in tflite or coral API.
|
|
loop = asyncio.get_event_loop()
|
|
loop.call_later(4 * 60 * 60, lambda: self.requestRestart())
|
|
|
|
def getClasses(self) -> list[str]:
|
|
return list(self.labels.values())
|
|
|
|
def getTriggerClasses(self) -> list[str]:
|
|
return ['motion']
|
|
|
|
async def createMedia(self, data: RawImage) -> scrypted_sdk.MediaObject:
|
|
mo = await scrypted_sdk.mediaManager.createMediaObject(data, self.fromMimeType)
|
|
return mo
|
|
|
|
def end_session(self, detection_session: PredictSession):
|
|
image = detection_session.image
|
|
if image:
|
|
detection_session.image = None
|
|
image.close()
|
|
|
|
dps = detection_session.processed / (time.time() - detection_session.start_time)
|
|
print("Detections per second %s" % dps)
|
|
return super().end_session(detection_session)
|
|
|
|
def invalidateMedia(self, detection_session: PredictSession, data: RawImage):
|
|
if not data:
|
|
return
|
|
image = data.image
|
|
data.image = None
|
|
if image:
|
|
if not detection_session.image:
|
|
detection_session.image = image
|
|
else:
|
|
image.close()
|
|
data.jpegMediaObject = None
|
|
|
|
async def convert(self, data: RawImage, fromMimeType: str, toMimeType: str, options: scrypted_sdk.BufferConvertorOptions = None) -> Any:
|
|
mo = data.jpegMediaObject
|
|
if not mo:
|
|
image = data.image
|
|
if not image:
|
|
raise Exception('data is no longer valid')
|
|
|
|
bio = io.BytesIO()
|
|
image.save(bio, format='JPEG')
|
|
jpegBytes = bio.getvalue()
|
|
mo = await scrypted_sdk.mediaManager.createMediaObject(jpegBytes, 'image/jpeg')
|
|
data.jpegMediaObject = mo
|
|
return mo
|
|
|
|
def requestRestart(self):
|
|
asyncio.ensure_future(scrypted_sdk.deviceManager.requestRestart())
|
|
|
|
# width, height, channels
|
|
def get_input_details(self) -> Tuple[int, int, int]:
|
|
pass
|
|
|
|
def getModelSettings(self, settings: Any = None) -> list[Setting]:
|
|
allowList: Setting = {
|
|
'title': 'Detections Types',
|
|
# 'subgroup': 'Advanced',
|
|
'description': 'The detections that will be reported. If none are specified, all detections will be reported. Select only detection types of interest for optimal performance.',
|
|
'choices': self.getClasses(),
|
|
'multiple': True,
|
|
'key': 'allowList',
|
|
'value': [
|
|
'person',
|
|
'dog',
|
|
'cat',
|
|
'car',
|
|
'truck',
|
|
'bus',
|
|
'motorcycle',
|
|
],
|
|
}
|
|
|
|
trackerWindow: Setting = {
|
|
'title': 'Tracker Window',
|
|
'subgroup': 'Advanced',
|
|
'description': 'Internal Setting. Do not change.',
|
|
'key': 'trackerWindow',
|
|
'value': 3,
|
|
'type': 'number',
|
|
}
|
|
trackerCertainty: Setting = {
|
|
'title': 'Tracker Certainty',
|
|
'subgroup': 'Advanced',
|
|
'description': 'Internal Setting. Do not change.',
|
|
'key': 'trackerCertainty',
|
|
'value': .2,
|
|
'type': 'number',
|
|
}
|
|
return [allowList, trackerWindow, trackerCertainty]
|
|
|
|
def create_detection_result(self, objs: List[Prediction], size, allowList, convert_to_src_size=None) -> ObjectsDetected:
|
|
detections: List[ObjectDetectionResult] = []
|
|
detection_result: ObjectsDetected = {}
|
|
detection_result['detections'] = detections
|
|
detection_result['inputDimensions'] = size
|
|
|
|
for obj in objs:
|
|
className = self.labels.get(obj.id, obj.id)
|
|
if allowList and len(allowList) and className not in allowList:
|
|
continue
|
|
detection: ObjectDetectionResult = {}
|
|
detection['boundingBox'] = (
|
|
obj.bbox.xmin, obj.bbox.ymin, obj.bbox.xmax - obj.bbox.xmin, obj.bbox.ymax - obj.bbox.ymin)
|
|
detection['className'] = className
|
|
detection['score'] = obj.score
|
|
detections.append(detection)
|
|
|
|
if convert_to_src_size:
|
|
detections = detection_result['detections']
|
|
detection_result['detections'] = []
|
|
for detection in detections:
|
|
bb = detection['boundingBox']
|
|
x, y, valid = convert_to_src_size((bb[0], bb[1]), True)
|
|
x2, y2, valid2 = convert_to_src_size(
|
|
(bb[0] + bb[2], bb[1] + bb[3]), True)
|
|
if not valid or not valid2:
|
|
# print("filtering out", detection['className'])
|
|
continue
|
|
detection['boundingBox'] = (x, y, x2 - x + 1, y2 - y + 1)
|
|
detection_result['detections'].append(detection)
|
|
|
|
# print(detection_result)
|
|
return detection_result
|
|
|
|
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, _ = await self.run_detection_image(detection_session, image, settings, image.size)
|
|
return detections
|
|
|
|
def get_detection_input_size(self, src_size):
|
|
# signals to pipeline that any input size is fine
|
|
# previous code used to resize to correct size and run detection that way.
|
|
# new code will resize the frame and potentially do multiple passes.
|
|
# this is useful for high quality thumbnails.
|
|
return (None, None)
|
|
|
|
def get_input_size(self) -> Tuple[int, int]:
|
|
pass
|
|
|
|
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:
|
|
src_size = videoFrame.width, videoFrame.height
|
|
w, h = self.get_input_size()
|
|
iw, ih = src_size
|
|
ws = w / iw
|
|
hs = h / ih
|
|
s = max(ws, hs)
|
|
if ws == 1 and hs == 1:
|
|
def cvss(point, normalize=False):
|
|
return point[0], point[1], True
|
|
|
|
data = await videoFrame.toBuffer({
|
|
'format': 'rgb',
|
|
})
|
|
image = Image.frombuffer('RGB', (w, h), data)
|
|
try:
|
|
ret = await self.detect_once(image, settings, src_size, cvss)
|
|
return ret
|
|
finally:
|
|
image.close()
|
|
|
|
sw = int(w / s)
|
|
sh = int(h / s)
|
|
first_crop = (0, 0, sw, sh)
|
|
|
|
|
|
ow = iw - sw
|
|
oh = ih - sh
|
|
second_crop = (ow, oh, ow + sw, oh + sh)
|
|
|
|
firstData, secondData = await asyncio.gather(
|
|
videoFrame.toBuffer({
|
|
'resize': {
|
|
'width': w,
|
|
'height': h,
|
|
},
|
|
'crop': {
|
|
'left': 0,
|
|
'top': 0,
|
|
'width': sw,
|
|
'height': sh,
|
|
},
|
|
'format': 'rgb',
|
|
}),
|
|
videoFrame.toBuffer({
|
|
'resize': {
|
|
'width': w,
|
|
'height': h,
|
|
},
|
|
'crop': {
|
|
'left': ow,
|
|
'top': oh,
|
|
'width': sw,
|
|
'height': sh,
|
|
},
|
|
'format': 'rgb',
|
|
})
|
|
)
|
|
|
|
first = Image.frombuffer('RGB', (w, h), firstData)
|
|
second = Image.frombuffer('RGB', (w, h), secondData)
|
|
|
|
def cvss1(point, normalize=False):
|
|
return point[0] / s, point[1] / s, True
|
|
def cvss2(point, normalize=False):
|
|
return point[0] / s + ow, point[1] / s + oh, True
|
|
|
|
ret1 = await self.detect_once(first, settings, src_size, cvss1)
|
|
first.close()
|
|
ret2 = await self.detect_once(second, settings, src_size, cvss2)
|
|
second.close()
|
|
|
|
two_intersect = intersect_rect(Rectangle(*first_crop), Rectangle(*second_crop))
|
|
|
|
def is_same_detection_middle(d1: ObjectDetectionResult, d2: ObjectDetectionResult):
|
|
same, ret = is_same_detection(d1, d2)
|
|
if same:
|
|
return same, ret
|
|
|
|
if d1['className'] != d2['className']:
|
|
return False, None
|
|
|
|
r1 = from_bounding_box(d1['boundingBox'])
|
|
m1 = intersect_rect(two_intersect, r1)
|
|
if not m1:
|
|
return False, None
|
|
|
|
r2 = from_bounding_box(d2['boundingBox'])
|
|
m2 = intersect_rect(two_intersect, r2)
|
|
if not m2:
|
|
return False, None
|
|
|
|
same, ret = is_same_box(to_bounding_box(m1), to_bounding_box(m2))
|
|
if not same:
|
|
return False, None
|
|
c = to_bounding_box(combine_rect(r1, r2))
|
|
return True, c
|
|
|
|
ret = ret1
|
|
ret['detections'] = dedupe_detections(ret1['detections'] + ret2['detections'], is_same_detection=is_same_detection_middle)
|
|
return ret
|
|
|
|
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
|
|
|
|
if detection_session and not detection_session.tracker:
|
|
t = self.trackers.get(detection_session.id)
|
|
if not t:
|
|
t = tracker.Sort_OH(scene=np.array([iw, ih]))
|
|
trackerCertainty = settings.get('trackerCertainty')
|
|
if not isinstance(trackerCertainty, int):
|
|
trackerCertainty = .2
|
|
t.conf_three_frame_certainty = trackerCertainty * 3
|
|
trackerWindow = settings.get('trackerWindow')
|
|
if not isinstance(trackerWindow, int):
|
|
trackerWindow = 3
|
|
t.conf_unmatched_history_size = trackerWindow
|
|
self.trackers[detection_session.id] = t
|
|
detection_session.tracker = t
|
|
# conf_trgt = 0.35
|
|
# conf_objt = 0.75
|
|
# detection_session.tracker.conf_trgt = conf_trgt
|
|
# detection_session.tracker.conf_objt = conf_objt
|
|
|
|
# this a single pass or the second pass. detect once and return results.
|
|
if multipass_crop:
|
|
(l, t, dx, dy) = multipass_crop
|
|
|
|
# find center
|
|
cx = l + dx / 2
|
|
cy = t + dy / 2
|
|
|
|
# fix aspect ratio on box
|
|
if dx / w > dy / h:
|
|
dy = dx / w * h
|
|
else:
|
|
dx = dy / h * w
|
|
|
|
if dx > image.width:
|
|
s = image.width / dx
|
|
dx = image.width
|
|
dy *= s
|
|
|
|
if dy > image.height:
|
|
s = image.height / dy
|
|
dy = image.height
|
|
dx *= s
|
|
|
|
# crop size to fit input size
|
|
if dx < w:
|
|
dx = w
|
|
if dy < h:
|
|
dy = h
|
|
|
|
l = cx - dx / 2
|
|
t = cy - dy / 2
|
|
if l < 0:
|
|
l = 0
|
|
if t < 0:
|
|
t = 0
|
|
if l + dx > iw:
|
|
l = iw - dx
|
|
if t + dy > ih:
|
|
t = ih - dy
|
|
crop_box = (l, t, l + dx, t + dy)
|
|
if dx == w and dy == h:
|
|
input = image.crop(crop_box)
|
|
else:
|
|
input = image.resize((w, h), Image.ANTIALIAS, crop_box)
|
|
|
|
def cvss(point, normalize=False):
|
|
unscaled = ((point[0] / w) * dx + l, (point[1] / h) * dy + t)
|
|
converted = convert_to_src_size(unscaled, normalize) if convert_to_src_size else (unscaled[0], unscaled[1], True)
|
|
return converted
|
|
|
|
ret = await self.detect_once(input, settings, src_size, cvss)
|
|
input.close()
|
|
detection_session.processed = detection_session.processed + 1
|
|
return ret, RawImage(image)
|
|
|
|
ws = w / iw
|
|
hs = h / ih
|
|
s = max(ws, hs)
|
|
if ws == 1 and hs == 1:
|
|
def cvss(point, normalize=False):
|
|
converted = convert_to_src_size(point, normalize) if convert_to_src_size else (point[0], point[1], True)
|
|
return converted
|
|
|
|
ret = await self.detect_once(image, settings, src_size, cvss)
|
|
if detection_session:
|
|
detection_session.processed = detection_session.processed + 1
|
|
else:
|
|
sw = int(w / s)
|
|
sh = int(h / s)
|
|
first_crop = (0, 0, sw, sh)
|
|
first = image.resize((w, h), Image.ANTIALIAS, first_crop)
|
|
ow = iw - sw
|
|
oh = ih - sh
|
|
second_crop = (ow, oh, ow + sw, oh + sh)
|
|
second = image.resize((w, h), Image.ANTIALIAS, second_crop)
|
|
|
|
def cvss1(point, normalize=False):
|
|
unscaled = (point[0] / s, point[1] / s)
|
|
converted = convert_to_src_size(unscaled, normalize) if convert_to_src_size else (unscaled[0], unscaled[1], True)
|
|
return converted
|
|
def cvss2(point, normalize=False):
|
|
unscaled = (point[0] / s + ow, point[1] / s + oh)
|
|
converted = convert_to_src_size(unscaled, normalize) if convert_to_src_size else (unscaled[0], unscaled[1], True)
|
|
return converted
|
|
|
|
ret1 = await self.detect_once(first, settings, src_size, cvss1)
|
|
first.close()
|
|
if detection_session:
|
|
detection_session.processed = detection_session.processed + 1
|
|
ret2 = await self.detect_once(second, settings, src_size, cvss2)
|
|
if detection_session:
|
|
detection_session.processed = detection_session.processed + 1
|
|
second.close()
|
|
|
|
two_intersect = intersect_rect(Rectangle(*first_crop), Rectangle(*second_crop))
|
|
|
|
def is_same_detection_middle(d1: ObjectDetectionResult, d2: ObjectDetectionResult):
|
|
same, ret = is_same_detection(d1, d2)
|
|
if same:
|
|
return same, ret
|
|
|
|
if d1['className'] != d2['className']:
|
|
return False, None
|
|
|
|
r1 = from_bounding_box(d1['boundingBox'])
|
|
m1 = intersect_rect(two_intersect, r1)
|
|
if not m1:
|
|
return False, None
|
|
|
|
r2 = from_bounding_box(d2['boundingBox'])
|
|
m2 = intersect_rect(two_intersect, r2)
|
|
if not m2:
|
|
return False, None
|
|
|
|
same, ret = is_same_box(to_bounding_box(m1), to_bounding_box(m2))
|
|
if not same:
|
|
return False, None
|
|
c = to_bounding_box(combine_rect(r1, r2))
|
|
return True, c
|
|
|
|
ret = ret1
|
|
ret['detections'] = dedupe_detections(ret1['detections'] + ret2['detections'], is_same_detection=is_same_detection_middle)
|
|
|
|
if detection_session:
|
|
self.track(detection_session, ret)
|
|
|
|
if not len(ret['detections']):
|
|
return ret, RawImage(image)
|
|
|
|
return ret, RawImage(image)
|
|
|
|
def track(self, detection_session: PredictSession, ret: ObjectsDetected):
|
|
detections = ret['detections']
|
|
sort_input = []
|
|
for d in ret['detections']:
|
|
r: ObjectDetectionResult = d
|
|
l, t, w, h = r['boundingBox']
|
|
sort_input.append([l, t, l + w, t + h, r['score']])
|
|
trackers, unmatched_trckr, unmatched_gts = detection_session.tracker.update(np.array(sort_input), [])
|
|
for td in trackers:
|
|
x0, y0, x1, y1, trackID = td[0].item(), td[1].item(
|
|
), td[2].item(), td[3].item(), td[4].item()
|
|
slop = 0
|
|
obj: ObjectDetectionResult = None
|
|
ta = (x1 - x0) * (y1 - y0)
|
|
box = Rectangle(x0, y0, x1, y1)
|
|
for d in detections:
|
|
if d.get('id'):
|
|
continue
|
|
ob: ObjectDetectionResult = d
|
|
dx0, dy0, dw, dh = ob['boundingBox']
|
|
dx1 = dx0 + dw
|
|
dy1 = dy0 + dh
|
|
da = dw * dh
|
|
area = intersect_area(Rectangle(dx0, dy0, dx1, dy1), box)
|
|
if not area:
|
|
continue
|
|
# intersect area always gonna be smaller than
|
|
# the detection or tracker area.
|
|
# greater numbers, ie approaching 2, is better.
|
|
dslop = area / ta + area / da
|
|
if (dslop > slop):
|
|
slop = dslop
|
|
obj = ob
|
|
if obj:
|
|
obj['id'] = str(trackID)
|
|
# this may happen if tracker predicts something is still in the scene
|
|
# but was not detected
|
|
# else:
|
|
# print('unresolved tracker')
|
|
# for d in detections:
|
|
# if not d.get('id'):
|
|
# # this happens if the tracker is not confident in a new detection yet due
|
|
# # to low score or has not been found in enough frames
|
|
# if d['className'] == 'person':
|
|
# print('untracked %s: %s' % (d['className'], d['score']))
|
|
|
|
|
|
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
|
|
|
|
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')
|
|
width = caps.get_structure(0).get_value('width')
|
|
gst_buffer = gstsample.get_buffer()
|
|
result, info = gst_buffer.map(Gst.MapFlags.READ)
|
|
if not result:
|
|
return
|
|
try:
|
|
image = detection_session.image
|
|
detection_session.image = None
|
|
|
|
if image and (image.width != width or image.height != height):
|
|
image.close()
|
|
image = None
|
|
if image:
|
|
image.frombytes(bytes(info.data))
|
|
else:
|
|
image = Image.frombuffer('RGB', (width, height), bytes(info.data))
|
|
finally:
|
|
gst_buffer.unmap(info)
|
|
|
|
try:
|
|
return await self.run_detection_image(detection_session, image, settings, src_size, convert_to_src_size)
|
|
except:
|
|
image.close()
|
|
traceback.print_exc()
|
|
raise
|
|
|
|
def create_detection_session(self):
|
|
return PredictSession(start_time=time.time())
|