the changes I made don't use these anymore
This commit is contained in:
Cowwan
parent
1e622bbc92
commit
33007ca428
7 changed files with 0 additions and 2632 deletions
341
dshowcapture.py
341
dshowcapture.py
|
|
@ -1,341 +0,0 @@
|
||||||
import os
|
|
||||||
import platform
|
|
||||||
import sys
|
|
||||||
import numpy as np
|
|
||||||
from ctypes import *
|
|
||||||
from PIL import Image
|
|
||||||
import cv2
|
|
||||||
import json
|
|
||||||
|
|
||||||
def resolve(name):
|
|
||||||
f = os.path.join(os.path.dirname(__file__), name)
|
|
||||||
return f
|
|
||||||
|
|
||||||
lib = None
|
|
||||||
bm_lib = None
|
|
||||||
bm_options = None
|
|
||||||
bm_enabled = False
|
|
||||||
|
|
||||||
def set_bm_enabled(v):
|
|
||||||
global bm_enabled
|
|
||||||
bm_enabled = v
|
|
||||||
|
|
||||||
def set_options(str):
|
|
||||||
global bm_options
|
|
||||||
bm_options = str
|
|
||||||
|
|
||||||
def create_frame_buffer(width, height, factor):
|
|
||||||
buffer = bytearray(width * height * 4 * factor)
|
|
||||||
char_array = c_char * len(buffer)
|
|
||||||
return char_array.from_buffer(buffer)
|
|
||||||
|
|
||||||
class DShowCapture():
|
|
||||||
def __init__(self):
|
|
||||||
global lib
|
|
||||||
global bm_lib
|
|
||||||
global bm_options
|
|
||||||
global bm_enabled
|
|
||||||
if lib is None or bm_lib is None:
|
|
||||||
if platform.architecture()[0] == '32bit':
|
|
||||||
dll_path = resolve(os.path.join("dshowcapture", "dshowcapture_x86.dll"))
|
|
||||||
lib = cdll.LoadLibrary(dll_path)
|
|
||||||
if bm_enabled:
|
|
||||||
dll_path = resolve(os.path.join("dshowcapture", "libminibmcapture32.dll"))
|
|
||||||
bm_lib = cdll.LoadLibrary(dll_path)
|
|
||||||
else:
|
|
||||||
dll_path = resolve(os.path.join("dshowcapture", "dshowcapture_x64.dll"))
|
|
||||||
lib = cdll.LoadLibrary(dll_path)
|
|
||||||
if bm_enabled:
|
|
||||||
dll_path = resolve(os.path.join("dshowcapture", "libminibmcapture64.dll"))
|
|
||||||
bm_lib = cdll.LoadLibrary(dll_path)
|
|
||||||
|
|
||||||
# DirectShow
|
|
||||||
lib.create_capture.restype = c_void_p
|
|
||||||
lib.get_devices.argtypes = [c_void_p]
|
|
||||||
lib.get_device.argtypes = [c_void_p, c_int, c_char_p, c_int]
|
|
||||||
lib.capture_device.argtypes = [c_void_p, c_int, c_int, c_int, c_int]
|
|
||||||
lib.capture_device_by_dcap.argtypes = [c_void_p, c_int, c_int, c_int, c_int, c_longlong]
|
|
||||||
lib.capture_device_default.argtypes = [c_void_p, c_int]
|
|
||||||
lib.get_width.argtypes = [c_void_p]
|
|
||||||
lib.get_height.argtypes = [c_void_p]
|
|
||||||
lib.get_fps.argtypes = [c_void_p]
|
|
||||||
lib.get_flipped.argtypes = [c_void_p]
|
|
||||||
lib.get_colorspace.argtypes = [c_void_p]
|
|
||||||
lib.get_colorspace_internal.argtypes = [c_void_p]
|
|
||||||
lib.capturing.argtypes = [c_void_p]
|
|
||||||
lib.get_frame.argtypes = [c_void_p, c_int, c_char_p, c_int]
|
|
||||||
lib.get_size.argtypes = [c_void_p]
|
|
||||||
lib.stop_capture.argtypes = [c_void_p]
|
|
||||||
lib.get_json_length.argtypes = [c_void_p]
|
|
||||||
lib.get_json.argtypes = [c_void_p, c_char_p, c_int]
|
|
||||||
lib.destroy_capture.argtypes = [c_void_p]
|
|
||||||
|
|
||||||
# Blackmagic
|
|
||||||
if bm_enabled:
|
|
||||||
bm_lib.start_capture_single.argtypes = [c_int, c_int, c_void_p]
|
|
||||||
bm_lib.read_frame_bgra32_blocking.argtypes = [c_char_p, c_int]
|
|
||||||
bm_lib.stop_capture_single.argtypes = []
|
|
||||||
bm_lib.get_json_length.argtypes = []
|
|
||||||
bm_lib.get_json.argtypes = [c_char_p, c_int]
|
|
||||||
if bm_options is not None:
|
|
||||||
bm_lib.set_options.argtypes = [c_char_p]
|
|
||||||
bm_lib.set_options(bm_options.encode())
|
|
||||||
self.lib = lib
|
|
||||||
if bm_enabled:
|
|
||||||
self.bm_lib = bm_lib
|
|
||||||
self.cap = lib.create_capture()
|
|
||||||
self.name_buffer = create_string_buffer(255);
|
|
||||||
self.buffer = None
|
|
||||||
self.have_devices = False
|
|
||||||
self.size = None
|
|
||||||
self.real_size = None
|
|
||||||
self.info = None
|
|
||||||
self.type = None
|
|
||||||
|
|
||||||
def __del__(self):
|
|
||||||
if self.buffer is not None:
|
|
||||||
del self.buffer
|
|
||||||
del self.name_buffer
|
|
||||||
self.destroy_capture()
|
|
||||||
|
|
||||||
def get_devices(self):
|
|
||||||
self.have_devices = True
|
|
||||||
return self.lib.get_devices(self.cap)
|
|
||||||
|
|
||||||
def get_device(self, device_number):
|
|
||||||
self.lib.get_device(self.cap, device_number, self.name_buffer, 255)
|
|
||||||
name_str = str(self.name_buffer.value.decode('utf8', 'surrogateescape'))
|
|
||||||
return name_str
|
|
||||||
|
|
||||||
def get_info(self):
|
|
||||||
global bm_enabled
|
|
||||||
# DirectShow
|
|
||||||
json_length = self.lib.get_json_length(self.cap);
|
|
||||||
json_buffer = create_string_buffer(json_length)
|
|
||||||
self.lib.get_json(self.cap, json_buffer, json_length);
|
|
||||||
json_text = str(json_buffer.value.decode('utf8', 'surrogateescape'))
|
|
||||||
self.info = json.loads(json_text)
|
|
||||||
for cam in self.info:
|
|
||||||
cam["type"] = "DirectShow"
|
|
||||||
cam["index"] = cam["id"]
|
|
||||||
|
|
||||||
# Blackmagic
|
|
||||||
if bm_enabled:
|
|
||||||
json_length = self.bm_lib.get_json_length();
|
|
||||||
json_buffer = create_string_buffer(json_length)
|
|
||||||
self.bm_lib.get_json(json_buffer, json_length);
|
|
||||||
json_text = str(json_buffer.value.decode('utf8', 'surrogateescape'))
|
|
||||||
bm_info = json.loads(json_text)
|
|
||||||
dshow_len = len(self.info)
|
|
||||||
for cam in bm_info:
|
|
||||||
cam["type"] = "Blackmagic"
|
|
||||||
cam["index"] = cam["id"] + dshow_len
|
|
||||||
self.info.extend(bm_info)
|
|
||||||
|
|
||||||
return self.info
|
|
||||||
|
|
||||||
def capture_device(self, cam, width, height, fps):
|
|
||||||
if not self.have_devices:
|
|
||||||
self.get_devices()
|
|
||||||
ret = self.lib.capture_device(self.cap, cam, width, height, fps) == 1
|
|
||||||
if ret:
|
|
||||||
self.type = "DirectShow"
|
|
||||||
self.width = self.get_width()
|
|
||||||
self.height = self.get_height()
|
|
||||||
self.fps = self.get_fps()
|
|
||||||
self.flipped = self.get_flipped()
|
|
||||||
self.colorspace = self.get_colorspace()
|
|
||||||
self.colorspace_internal = self.get_colorspace_internal()
|
|
||||||
self.size = self.width * self.height * 4
|
|
||||||
self.buffer = create_frame_buffer(self.width, self.height, 4)
|
|
||||||
else:
|
|
||||||
self.size = None
|
|
||||||
self.real_size = None
|
|
||||||
return ret;
|
|
||||||
|
|
||||||
def capture_device_by_dcap(self, cam, dcap, width, height, fps):
|
|
||||||
global bm_enabled
|
|
||||||
if not self.have_devices:
|
|
||||||
self.get_devices()
|
|
||||||
if self.info is None:
|
|
||||||
self.get_info()
|
|
||||||
fps = int(10000000 / fps)
|
|
||||||
ret = False
|
|
||||||
if self.info[cam]['type'] == "DirectShow":
|
|
||||||
ret = self.lib.capture_device_by_dcap(self.cap, cam, dcap, width, height, fps) == 1
|
|
||||||
elif bm_enabled and self.info[cam]['type'] == "Blackmagic":
|
|
||||||
ret = self.bm_lib.start_capture_single(self.info[cam]['id'], self.info[cam]['caps'][dcap]['bmModecode'], None) == 1
|
|
||||||
if ret:
|
|
||||||
self.type = self.info[cam]['type']
|
|
||||||
if self.type == "DirectShow":
|
|
||||||
self.width = self.get_width()
|
|
||||||
self.height = self.get_height()
|
|
||||||
self.fps = self.get_fps()
|
|
||||||
self.flipped = self.get_flipped()
|
|
||||||
self.colorspace = self.get_colorspace()
|
|
||||||
self.colorspace_internal = self.get_colorspace_internal()
|
|
||||||
self.size = self.width * self.height * 4 * 4
|
|
||||||
self.buffer = create_frame_buffer(self.width, self.height, 4)
|
|
||||||
elif bm_enabled and self.type == "Blackmagic":
|
|
||||||
self.width = self.info[cam]['caps'][dcap]['minCX']
|
|
||||||
self.height = self.info[cam]['caps'][dcap]['minCY']
|
|
||||||
self.fps = int(10000000 / self.info[cam]['caps'][dcap]['minInterval'])
|
|
||||||
self.flipped = True
|
|
||||||
self.colorspace = 101
|
|
||||||
self.colorspace_internal = self.colorspace
|
|
||||||
self.size = self.width * self.height * 4
|
|
||||||
self.buffer = create_frame_buffer(self.width, self.height, 1)
|
|
||||||
self.real_size = self.size
|
|
||||||
else:
|
|
||||||
self.size = None
|
|
||||||
self.real_size = None
|
|
||||||
return ret;
|
|
||||||
|
|
||||||
def capture_device_default(self, cam):
|
|
||||||
if not self.have_devices:
|
|
||||||
self.get_devices()
|
|
||||||
ret = self.lib.capture_device_default(self.cap, cam) == 1
|
|
||||||
if ret:
|
|
||||||
self.type = "DirectShow"
|
|
||||||
self.width = self.get_width()
|
|
||||||
self.height = self.get_height()
|
|
||||||
self.fps = self.get_fps()
|
|
||||||
self.flipped = self.get_flipped()
|
|
||||||
self.colorspace = self.get_colorspace()
|
|
||||||
self.colorspace_internal = self.get_colorspace_internal()
|
|
||||||
self.size = self.width * self.height * 4 * 4
|
|
||||||
self.buffer = create_frame_buffer(self.width, self.height, 4)
|
|
||||||
else:
|
|
||||||
self.size = None
|
|
||||||
self.real_size = None
|
|
||||||
return ret;
|
|
||||||
|
|
||||||
def get_width(self):
|
|
||||||
return self.lib.get_width(self.cap)
|
|
||||||
|
|
||||||
def get_height(self):
|
|
||||||
return self.lib.get_height(self.cap)
|
|
||||||
|
|
||||||
def get_fps(self):
|
|
||||||
return self.lib.get_fps(self.cap)
|
|
||||||
|
|
||||||
def get_flipped(self):
|
|
||||||
return self.lib.get_flipped(self.cap) != 0
|
|
||||||
|
|
||||||
def get_colorspace(self):
|
|
||||||
colorspace = self.lib.get_colorspace(self.cap)
|
|
||||||
if colorspace == 0:
|
|
||||||
colorspace = self.lib.get_colorspace_internal(self.cap)
|
|
||||||
return colorspace
|
|
||||||
|
|
||||||
def get_colorspace_internal(self):
|
|
||||||
return self.lib.get_colorspace_internal(self.cap)
|
|
||||||
|
|
||||||
def capturing(self):
|
|
||||||
if self.type == "Blackmagic":
|
|
||||||
return 1
|
|
||||||
return self.lib.capturing(self.cap) == 1
|
|
||||||
|
|
||||||
def get_frame(self, timeout):
|
|
||||||
global bm_enabled
|
|
||||||
if self.size is None:
|
|
||||||
return None
|
|
||||||
if self.type == "DirectShow":
|
|
||||||
self.real_size = self.lib.get_frame(self.cap, timeout, self.buffer, self.size)
|
|
||||||
elif bm_enabled and self.type == "Blackmagic":
|
|
||||||
self.bm_lib.read_frame_bgra32_blocking(self.buffer, self.size)
|
|
||||||
else:
|
|
||||||
return None
|
|
||||||
img = np.frombuffer(self.buffer, dtype=np.uint8)[0:self.real_size]
|
|
||||||
if self.colorspace in [100, 101]:
|
|
||||||
if self.real_size == self.height * self.width * 4:
|
|
||||||
img = cv2.cvtColor(img.reshape((self.height,self.width,4)), cv2.COLOR_BGRA2BGR)
|
|
||||||
elif self.real_size == self.height * self.width * 3:
|
|
||||||
img = img.reshape((self.height,self.width,3))
|
|
||||||
else:
|
|
||||||
return None
|
|
||||||
if not self.flipped:
|
|
||||||
img = cv2.flip(img, 0)
|
|
||||||
return img
|
|
||||||
elif self.colorspace == 200 and self.real_size == (3 * self.height // 2) * self.width:
|
|
||||||
img = cv2.cvtColor(img.reshape((3 * self.height // 2,self.width,1)), cv2.COLOR_YUV2BGR_I420)
|
|
||||||
elif self.colorspace == 201 and self.real_size == (3 * self.height // 2) * self.width:
|
|
||||||
img = cv2.cvtColor(img.reshape((3 * self.height // 2,self.width,1)), cv2.COLOR_YUV2BGR_NV12)
|
|
||||||
elif self.colorspace == 202 and self.real_size == (3 * self.height // 2) * self.width:
|
|
||||||
img = cv2.cvtColor(img.reshape((3 * self.height // 2,self.width,1)), cv2.COLOR_YUV2BGR_YV12)
|
|
||||||
elif self.colorspace == 203 and self.real_size == self.height * self.width:
|
|
||||||
img = cv2.cvtColor(img.reshape((self.height,self.width,1)), cv2.COLOR_GRAY2BGR)
|
|
||||||
elif self.colorspace == 300 and self.real_size == self.height * self.width * 2:
|
|
||||||
img = cv2.cvtColor(img.reshape((self.height,self.width,2)), cv2.COLOR_YUV2BGR_YVYU)
|
|
||||||
elif self.colorspace == 301 and self.real_size == self.height * self.width * 2:
|
|
||||||
img = cv2.cvtColor(img.reshape((self.height,self.width,2)), cv2.COLOR_YUV2BGR_YUY2)
|
|
||||||
elif self.colorspace == 302 and self.real_size == self.height * self.width * 2:
|
|
||||||
img = cv2.cvtColor(img.reshape((self.height,self.width,2)), cv2.COLOR_YUV2BGR_UYVY)
|
|
||||||
elif self.colorspace == 303:
|
|
||||||
return None
|
|
||||||
elif self.colorspace == 400 and self.real_size > 4:
|
|
||||||
img = cv2.imdecode(img, cv2.IMREAD_COLOR);
|
|
||||||
else:
|
|
||||||
return None
|
|
||||||
if self.flipped:
|
|
||||||
img = cv2.flip(img, 0)
|
|
||||||
return img
|
|
||||||
|
|
||||||
def stop_capture(self):
|
|
||||||
global bm_enabled
|
|
||||||
self.size = None
|
|
||||||
self.real_size = None
|
|
||||||
if self.type == "DirectShow":
|
|
||||||
return self.lib.stop_capture(self.cap)
|
|
||||||
elif bm_enabled and self.type == "Blackmagic":
|
|
||||||
return self.bm_lib.stop_capture_single() != 0
|
|
||||||
return True
|
|
||||||
|
|
||||||
def destroy_capture(self):
|
|
||||||
if self.cap is None:
|
|
||||||
return 0
|
|
||||||
self.stop_capture()
|
|
||||||
ret = self.lib.destroy_capture(self.cap)
|
|
||||||
self.cap = None
|
|
||||||
self.size = None
|
|
||||||
self.real_size = None
|
|
||||||
return ret
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
cam = 0
|
|
||||||
width = 1280
|
|
||||||
height = 720
|
|
||||||
fps = 30
|
|
||||||
if len(sys.argv) > 1:
|
|
||||||
cam = int(sys.argv[1])
|
|
||||||
if len(sys.argv) > 2:
|
|
||||||
width = int(sys.argv[2])
|
|
||||||
if len(sys.argv) > 3:
|
|
||||||
height = int(sys.argv[3])
|
|
||||||
if len(sys.argv) > 4:
|
|
||||||
fps = int(sys.argv[4])
|
|
||||||
|
|
||||||
cap = DShowCapture()
|
|
||||||
devices = cap.get_devices()
|
|
||||||
print("Devices: ", devices)
|
|
||||||
for i in range(devices):
|
|
||||||
print(f"{i} {cap.get_device(i)}")
|
|
||||||
|
|
||||||
print(f"Capturing: {cap.capture_device(cam, width, height, fps)}")
|
|
||||||
width = cap.get_width()
|
|
||||||
height = cap.get_height()
|
|
||||||
flipped = cap.get_flipped()
|
|
||||||
print(f"Width: {width} Height: {height} FPS: {cap.get_fps()} Flipped: {flipped} Colorspace: {cap.get_colorspace()} Internal: {cap.get_colorspace_internal()}")
|
|
||||||
|
|
||||||
while True:
|
|
||||||
img = cap.get_frame(1000)
|
|
||||||
if img is not None:
|
|
||||||
cv2.imshow("DShowCapture", img)
|
|
||||||
if cv2.waitKey(1) & 0xFF == ord('q'):
|
|
||||||
sys.exit(0)
|
|
||||||
else:
|
|
||||||
print("Lost frame")
|
|
||||||
|
|
||||||
cap.stop_capture()
|
|
||||||
cap.destroy_capture()
|
|
||||||
del cap
|
|
||||||
125
escapi.py
125
escapi.py
|
|
@ -1,125 +0,0 @@
|
||||||
"""
|
|
||||||
A simple python wrapper around escapi
|
|
||||||
|
|
||||||
Usage:
|
|
||||||
|
|
||||||
from camera import Device
|
|
||||||
|
|
||||||
device = Deveice.connect(0, 500, 500)
|
|
||||||
image = device.get_image()
|
|
||||||
"""
|
|
||||||
|
|
||||||
import os
|
|
||||||
import platform
|
|
||||||
import sys
|
|
||||||
from ctypes import *
|
|
||||||
from PIL import Image
|
|
||||||
import numpy as np
|
|
||||||
import cv2
|
|
||||||
|
|
||||||
def resolve(name):
|
|
||||||
f = os.path.join(os.path.dirname(__file__), name)
|
|
||||||
return f
|
|
||||||
|
|
||||||
|
|
||||||
class CAPTURE_PROPETIES:
|
|
||||||
CAPTURE_BRIGHTNESS = 1,
|
|
||||||
CAPTURE_CONTRAST = 2,
|
|
||||||
CAPTURE_HUE = 3,
|
|
||||||
CAPTURE_SATURATION = 4,
|
|
||||||
CAPTURE_SHARPNESS = 5,
|
|
||||||
CAPTURE_GAMMA = 6,
|
|
||||||
CAPTURE_COLORENABLE = 7,
|
|
||||||
CAPTURE_WHITEBALANCE = 8,
|
|
||||||
CAPTURE_BACKLIGHTCOMPENSATION = 9,
|
|
||||||
CAPTURE_GAIN = 10,
|
|
||||||
CAPTURE_PAN = 11,
|
|
||||||
CAPTURE_TILT = 12,
|
|
||||||
CAPTURE_ROLL = 13,
|
|
||||||
CAPTURE_ZOOM = 14,
|
|
||||||
CAPTURE_EXPOSURE = 15,
|
|
||||||
CAPTURE_IRIS = 16,
|
|
||||||
CAPTURE_FOCUS = 17,
|
|
||||||
CAPTURE_PROP_MAX = 18,
|
|
||||||
|
|
||||||
|
|
||||||
class SimpleCapParms(Structure):
|
|
||||||
_fields_ = [
|
|
||||||
("buffer", POINTER(c_int)),
|
|
||||||
("width", c_int),
|
|
||||||
("height", c_int),
|
|
||||||
("fps", c_int),
|
|
||||||
]
|
|
||||||
|
|
||||||
|
|
||||||
lib = None
|
|
||||||
|
|
||||||
|
|
||||||
def init():
|
|
||||||
global lib
|
|
||||||
if platform.architecture()[0] == '32bit':
|
|
||||||
dll_path = resolve(os.path.join("escapi", "escapi_x86.dll"))
|
|
||||||
lib = cdll.LoadLibrary(dll_path)
|
|
||||||
else:
|
|
||||||
dll_path = resolve(os.path.join("escapi", "escapi_x64.dll"))
|
|
||||||
lib = cdll.LoadLibrary(dll_path)
|
|
||||||
if lib is None or lib.ESCAPIVersion() != 0xfff001:
|
|
||||||
print("Invalid ESCAPI DLL found.")
|
|
||||||
sys.exit(1)
|
|
||||||
lib.initCapture.argtypes = [c_int, POINTER(SimpleCapParms)]
|
|
||||||
lib.initCapture.restype = c_int
|
|
||||||
lib.initCOM()
|
|
||||||
|
|
||||||
def count_capture_devices():
|
|
||||||
return lib.countCaptureDevices()
|
|
||||||
|
|
||||||
def device_name(device):
|
|
||||||
"""
|
|
||||||
Get the device name for the given device
|
|
||||||
:param device: The number of the device
|
|
||||||
:return: The name of the device
|
|
||||||
"""
|
|
||||||
namearry = (c_char * 256)()
|
|
||||||
lib.getCaptureDeviceName(device, namearry, 256)
|
|
||||||
camearaname = namearry.value
|
|
||||||
return camearaname
|
|
||||||
|
|
||||||
def init_camera(device, width, height, fps):
|
|
||||||
global devices
|
|
||||||
array = (width * height * c_int)()
|
|
||||||
options = SimpleCapParms()
|
|
||||||
options.width = width
|
|
||||||
options.height = height
|
|
||||||
options.fps = fps
|
|
||||||
options.buffer = array
|
|
||||||
lib.initCapture(device, byref(options))
|
|
||||||
return array
|
|
||||||
|
|
||||||
def do_capture(device):
|
|
||||||
lib.doCapture(device)
|
|
||||||
|
|
||||||
def is_capture_done(device):
|
|
||||||
return lib.isCaptureDone(device)
|
|
||||||
|
|
||||||
def read(device, width, height, array):
|
|
||||||
if is_capture_done(device):
|
|
||||||
img = Image.frombuffer('RGBA', (width, height), array, 'raw', 'BGRA', 0, 1)
|
|
||||||
img = np.array(img)
|
|
||||||
img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGR);
|
|
||||||
return img
|
|
||||||
else:
|
|
||||||
return None
|
|
||||||
|
|
||||||
def get_image(device, width, height, array):
|
|
||||||
lib.doCapture(device)
|
|
||||||
|
|
||||||
while lib.isCaptureDone(device) == 0:
|
|
||||||
pass
|
|
||||||
|
|
||||||
img = Image.frombuffer('RGBA', (width, height), array, 'raw', 'BGRA', 0, 0)
|
|
||||||
img = np.array(img)
|
|
||||||
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR);
|
|
||||||
return img
|
|
||||||
|
|
||||||
def deinit_camera(device):
|
|
||||||
lib.deinitCapture(device)
|
|
||||||
250
input_reader.py
250
input_reader.py
|
|
@ -1,250 +0,0 @@
|
||||||
import re
|
|
||||||
import sys
|
|
||||||
import os
|
|
||||||
import cv2
|
|
||||||
import numpy as np
|
|
||||||
import escapi
|
|
||||||
import dshowcapture
|
|
||||||
import time
|
|
||||||
import traceback
|
|
||||||
import gc
|
|
||||||
|
|
||||||
class VideoReader():
|
|
||||||
def __init__(self, capture, camera=False):
|
|
||||||
if os.name == 'nt' and camera:
|
|
||||||
self.cap = cv2.VideoCapture(capture, cv2.CAP_DSHOW)
|
|
||||||
else:
|
|
||||||
self.cap = cv2.VideoCapture(capture)
|
|
||||||
if self.cap is None:
|
|
||||||
print("The video source cannot be opened")
|
|
||||||
sys.exit(0)
|
|
||||||
self.name = str(capture)
|
|
||||||
def is_open(self):
|
|
||||||
return self.cap.isOpened()
|
|
||||||
def is_ready(self):
|
|
||||||
return True
|
|
||||||
def read(self):
|
|
||||||
return self.cap.read()
|
|
||||||
def close(self):
|
|
||||||
self.cap.release()
|
|
||||||
|
|
||||||
class EscapiReader(VideoReader):
|
|
||||||
def __init__(self, capture, width, height, fps):
|
|
||||||
self.device = None
|
|
||||||
self.width = width
|
|
||||||
self.height = height
|
|
||||||
self.fps = fps
|
|
||||||
self.device = capture
|
|
||||||
escapi.count_capture_devices()
|
|
||||||
self.name = str(escapi.device_name(self.device).decode('utf8', 'surrogateescape'))
|
|
||||||
self.buffer = escapi.init_camera(self.device, self.width, self.height, self.fps)
|
|
||||||
escapi.do_capture(self.device)
|
|
||||||
def is_open(self):
|
|
||||||
return True
|
|
||||||
def is_ready(self):
|
|
||||||
return escapi.is_capture_done(self.device)
|
|
||||||
def read(self):
|
|
||||||
if escapi.is_capture_done(self.device):
|
|
||||||
image = escapi.read(self.device, self.width, self.height, self.buffer)
|
|
||||||
escapi.do_capture(self.device)
|
|
||||||
return True, image
|
|
||||||
else:
|
|
||||||
return False, None
|
|
||||||
def close(self):
|
|
||||||
escapi.deinit_camera(self.device)
|
|
||||||
|
|
||||||
class DShowCaptureReader(VideoReader):
|
|
||||||
def __init__(self, capture, width, height, fps, use_dshowcapture=True, dcap=None):
|
|
||||||
self.device = None
|
|
||||||
self.width = width
|
|
||||||
self.height = height
|
|
||||||
self.fps = fps
|
|
||||||
self.dcap = dcap;
|
|
||||||
self.device = dshowcapture.DShowCapture()
|
|
||||||
self.device.get_devices()
|
|
||||||
info = self.device.get_info()
|
|
||||||
self.name = info[capture]['name']
|
|
||||||
if info[capture]['type'] == "Blackmagic":
|
|
||||||
self.name = "Blackmagic: " + self.name
|
|
||||||
if dcap is None or dcap < 0:
|
|
||||||
dcap = 0
|
|
||||||
ret = False
|
|
||||||
if dcap is None:
|
|
||||||
ret = self.device.capture_device(capture, self.width, self.height, self.fps)
|
|
||||||
else:
|
|
||||||
if dcap < 0:
|
|
||||||
ret = self.device.capture_device_default(capture)
|
|
||||||
else:
|
|
||||||
ret = self.device.capture_device_by_dcap(capture, dcap, self.width, self.height, self.fps)
|
|
||||||
if not ret:
|
|
||||||
raise Exception("Failed to start capture.")
|
|
||||||
self.width = self.device.width
|
|
||||||
self.height = self.device.height
|
|
||||||
self.fps = self.device.fps
|
|
||||||
print(f"Camera: \"{self.name}\" Capability ID: {dcap} Resolution: {self.device.width}x{self.device.height} Frame rate: {self.device.fps} Colorspace: {self.device.colorspace} Internal: {self.device.colorspace_internal} Flipped: {self.device.flipped}")
|
|
||||||
self.timeout = 1000
|
|
||||||
def is_open(self):
|
|
||||||
return self.device.capturing()
|
|
||||||
def is_ready(self):
|
|
||||||
return self.device.capturing()
|
|
||||||
def read(self):
|
|
||||||
img = None
|
|
||||||
try:
|
|
||||||
img = self.device.get_frame(self.timeout)
|
|
||||||
except:
|
|
||||||
gc.collect()
|
|
||||||
img = self.device.get_frame(self.timeout)
|
|
||||||
if img is None:
|
|
||||||
return False, None
|
|
||||||
else:
|
|
||||||
return True, img
|
|
||||||
def close(self):
|
|
||||||
self.device.destroy_capture()
|
|
||||||
|
|
||||||
class OpenCVReader(VideoReader):
|
|
||||||
def __init__(self, capture, width, height, fps):
|
|
||||||
self.device = None
|
|
||||||
self.width = width
|
|
||||||
self.height = height
|
|
||||||
self.fps = fps
|
|
||||||
self.name = str(capture)
|
|
||||||
super(OpenCVReader, self).__init__(capture, camera=True)
|
|
||||||
self.cap.set(3, width)
|
|
||||||
self.cap.set(4, height)
|
|
||||||
self.cap.set(38, 1)
|
|
||||||
def is_open(self):
|
|
||||||
return super(OpenCVReader, self).is_open()
|
|
||||||
def is_ready(self):
|
|
||||||
return super(OpenCVReader, self).is_ready()
|
|
||||||
def read(self):
|
|
||||||
return super(OpenCVReader, self).read()
|
|
||||||
def close(self):
|
|
||||||
super(OpenCVReader, self).close()
|
|
||||||
|
|
||||||
class RawReader:
|
|
||||||
def __init__(self, width, height):
|
|
||||||
self.width = int(width)
|
|
||||||
self.height = int(height)
|
|
||||||
|
|
||||||
if self.width < 1 or self.height < 1:
|
|
||||||
print("No acceptable size was given for reading raw RGB frames.")
|
|
||||||
sys.exit(0)
|
|
||||||
|
|
||||||
self.len = self.width * self.height * 3
|
|
||||||
self.open = True
|
|
||||||
def is_open(self):
|
|
||||||
return self.open
|
|
||||||
def is_ready(self):
|
|
||||||
return True
|
|
||||||
def read(self):
|
|
||||||
frame = bytearray()
|
|
||||||
read_bytes = 0
|
|
||||||
while read_bytes < self.len:
|
|
||||||
bytes = sys.stdin.buffer.read(self.len)
|
|
||||||
read_bytes += len(bytes)
|
|
||||||
frame.extend(bytes)
|
|
||||||
return True, np.frombuffer(frame, dtype=np.uint8).reshape((self.height, self.width, 3))
|
|
||||||
def close(self):
|
|
||||||
self.open = False
|
|
||||||
|
|
||||||
def try_int(s):
|
|
||||||
try:
|
|
||||||
return int(s)
|
|
||||||
except:
|
|
||||||
return None
|
|
||||||
|
|
||||||
def test_reader(reader):
|
|
||||||
got_any = 0
|
|
||||||
try:
|
|
||||||
for i in range(30):
|
|
||||||
if not reader.is_ready():
|
|
||||||
time.sleep(0.02)
|
|
||||||
ret, frame = reader.read()
|
|
||||||
if not ret:
|
|
||||||
time.sleep(0.02)
|
|
||||||
print("No frame")
|
|
||||||
else:
|
|
||||||
print("Got frame")
|
|
||||||
got_any += 1
|
|
||||||
if got_any > 10:
|
|
||||||
break
|
|
||||||
if reader.is_open():
|
|
||||||
return got_any > 0
|
|
||||||
print("Fail")
|
|
||||||
return False
|
|
||||||
except:
|
|
||||||
traceback.print_exc()
|
|
||||||
print("Except")
|
|
||||||
return False
|
|
||||||
|
|
||||||
class InputReader():
|
|
||||||
def __init__(self, capture, raw_rgb, width, height, fps, use_dshowcapture=False, dcap=None):
|
|
||||||
self.reader = None
|
|
||||||
self.name = str(capture)
|
|
||||||
try:
|
|
||||||
if raw_rgb > 0:
|
|
||||||
self.reader = RawReader(width, height)
|
|
||||||
elif os.path.exists(capture):
|
|
||||||
self.reader = VideoReader(capture)
|
|
||||||
elif capture == str(try_int(capture)):
|
|
||||||
if os.name == 'nt':
|
|
||||||
# Try with DShowCapture
|
|
||||||
good = True
|
|
||||||
name = ""
|
|
||||||
try:
|
|
||||||
if use_dshowcapture:
|
|
||||||
self.reader = DShowCaptureReader(int(capture), width, height, fps, dcap=dcap)
|
|
||||||
name = self.reader.name
|
|
||||||
good = test_reader(self.reader)
|
|
||||||
self.name = name
|
|
||||||
else:
|
|
||||||
good = False
|
|
||||||
except:
|
|
||||||
print("DShowCapture exception: ")
|
|
||||||
traceback.print_exc()
|
|
||||||
good = False
|
|
||||||
if good:
|
|
||||||
return
|
|
||||||
# Try with Escapi
|
|
||||||
good = True
|
|
||||||
try:
|
|
||||||
print(f"DShowCapture failed. Falling back to escapi for device {name}.", file=sys.stderr)
|
|
||||||
escapi.init()
|
|
||||||
devices = escapi.count_capture_devices()
|
|
||||||
found = None
|
|
||||||
for i in range(devices):
|
|
||||||
escapi_name = str(escapi.device_name(i).decode('utf8', 'surrogateescape'))
|
|
||||||
if name == escapi_name:
|
|
||||||
found = i
|
|
||||||
if found is None:
|
|
||||||
good = False
|
|
||||||
else:
|
|
||||||
print(f"Found device {name} as {i}.", file=sys.stderr)
|
|
||||||
self.reader = EscapiReader(found, width, height, fps)
|
|
||||||
good = test_reader(self.reader)
|
|
||||||
except:
|
|
||||||
print("Escapi exception: ")
|
|
||||||
traceback.print_exc()
|
|
||||||
good = False
|
|
||||||
if good:
|
|
||||||
return
|
|
||||||
# Try with OpenCV
|
|
||||||
print(f"Escapi failed. Falling back to OpenCV. If this fails, please change your camera settings.", file=sys.stderr)
|
|
||||||
self.reader = OpenCVReader(int(capture), width, height, fps)
|
|
||||||
self.name = self.reader.name
|
|
||||||
else:
|
|
||||||
self.reader = OpenCVReader(int(capture), width, height, fps)
|
|
||||||
except Exception as e:
|
|
||||||
print("Error: " + str(e))
|
|
||||||
|
|
||||||
if self.reader is None or not self.reader.is_open():
|
|
||||||
print("There was no valid input.")
|
|
||||||
sys.exit(0)
|
|
||||||
def is_open(self):
|
|
||||||
return self.reader.is_open()
|
|
||||||
def is_ready(self):
|
|
||||||
return self.reader.is_ready()
|
|
||||||
def read(self):
|
|
||||||
return self.reader.read()
|
|
||||||
def close(self):
|
|
||||||
self.reader.close()
|
|
||||||
314
model.py
314
model.py
|
|
@ -1,314 +0,0 @@
|
||||||
# This file is not used by the tracking application and currently outdated
|
|
||||||
import torch
|
|
||||||
import torch.nn as nn
|
|
||||||
import geffnet.mobilenetv3 # geffnet.mobilenetv3._gen_mobilenet_v3 needs to be patched to return the parameters instead of instantiating the network, based on commit c450c12ae6ffb1757f62dde3c2765da3c10f6def
|
|
||||||
from geffnet.efficientnet_builder import round_channels
|
|
||||||
|
|
||||||
class DSConv2d(nn.Module):
|
|
||||||
def __init__(self, in_planes, out_planes, kernels_per_layer=4, groups=1, old=0):
|
|
||||||
super(DSConv2d, self).__init__()
|
|
||||||
if old == 2:
|
|
||||||
self.conv = nn.Sequential(
|
|
||||||
nn.Conv2d(in_planes, in_planes * kernels_per_layer, kernel_size=3, padding=1, groups=in_planes),
|
|
||||||
nn.Conv2d(in_planes * kernels_per_layer, out_planes, kernel_size=1, groups=groups)
|
|
||||||
)
|
|
||||||
elif old == 1:
|
|
||||||
self.conv = nn.Sequential(
|
|
||||||
nn.Conv2d(in_planes, in_planes * kernels_per_layer, kernel_size=3, padding=1, groups=in_planes, bias=False),
|
|
||||||
nn.BatchNorm2d(in_planes * kernels_per_layer),
|
|
||||||
nn.Conv2d(in_planes * kernels_per_layer, out_planes, kernel_size=1, groups=groups, bias=False),
|
|
||||||
nn.BatchNorm2d(out_planes),
|
|
||||||
nn.ReLU6(inplace=True)
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
self.conv = nn.Sequential(
|
|
||||||
nn.Conv2d(in_planes, in_planes * kernels_per_layer, kernel_size=3, padding=1, groups=in_planes, bias=False),
|
|
||||||
nn.BatchNorm2d(in_planes * kernels_per_layer),
|
|
||||||
nn.ReLU6(inplace=True),
|
|
||||||
nn.Conv2d(in_planes * kernels_per_layer, out_planes, kernel_size=1, groups=groups, bias=False),
|
|
||||||
nn.BatchNorm2d(out_planes),
|
|
||||||
nn.ReLU6(inplace=True)
|
|
||||||
)
|
|
||||||
def forward(self, x):
|
|
||||||
x = self.conv(x)
|
|
||||||
return x
|
|
||||||
|
|
||||||
class UNetUp(nn.Module):
|
|
||||||
def __init__(self, in_channels, residual_in_channels, out_channels, size, old=0):
|
|
||||||
super(UNetUp, self).__init__()
|
|
||||||
self.up = nn.Upsample(size=size, mode='bilinear', align_corners=True)
|
|
||||||
self.conv = DSConv2d(in_channels + residual_in_channels, out_channels, 1, 1, old=old)
|
|
||||||
def forward(self, x1, x2):
|
|
||||||
x1 = self.up(x1)
|
|
||||||
x = torch.cat([x2, x1], dim=1)
|
|
||||||
x = self.conv(x)
|
|
||||||
return x
|
|
||||||
|
|
||||||
# This is the gaze tracking model
|
|
||||||
class OpenSeeFaceGaze(geffnet.mobilenetv3.MobileNetV3):
|
|
||||||
def __init__(self):
|
|
||||||
kwargs = geffnet.mobilenetv3._gen_mobilenet_v3(['small'])
|
|
||||||
super(OpenSeeFaceGaze, self).__init__(**kwargs)
|
|
||||||
self.up1 = UNetUp(576, 48, 64, (2,2), old=2)
|
|
||||||
self.up2 = UNetUp(64, 24, 32, (4,4), old=2)
|
|
||||||
self.up3 = UNetUp(32, 16, 15, (8,8), old=2)
|
|
||||||
self.group = DSConv2d(15, 3, kernels_per_layer=4, groups=3, old=2)
|
|
||||||
def _forward_impl(self, x):
|
|
||||||
x = self.conv_stem(x)
|
|
||||||
x = self.bn1(x)
|
|
||||||
x = self.act1(x)
|
|
||||||
r1 = None
|
|
||||||
r2 = None
|
|
||||||
r3 = None
|
|
||||||
for i, feature in enumerate(self.blocks):
|
|
||||||
x = feature(x)
|
|
||||||
if i == 3:
|
|
||||||
r3 = x
|
|
||||||
if i == 1:
|
|
||||||
r2 = x
|
|
||||||
if i == 0:
|
|
||||||
r1 = x
|
|
||||||
x = self.up1(x, r3)
|
|
||||||
x = self.up2(x, r2)
|
|
||||||
x = self.up3(x, r1)
|
|
||||||
x = self.group(x)
|
|
||||||
return x
|
|
||||||
def forward(self, x):
|
|
||||||
return self._forward_impl(x)
|
|
||||||
|
|
||||||
# This is the face detection model. Because the landmark model is very robust, it gets away with predicting very rough bounding boxes. It is fully convolutional and can be made to run on different resolutions. It was trained on 224x224 crops and the most reasonable results can be found in the range of 224x224 to 640x640.
|
|
||||||
class OpenSeeFaceDetect(geffnet.mobilenetv3.MobileNetV3):
|
|
||||||
def __init__(self, size="large", channel_multiplier=0.1):
|
|
||||||
kwargs = geffnet.mobilenetv3._gen_mobilenet_v3([size], channel_multiplier=channel_multiplier)
|
|
||||||
super(OpenSeeFaceDetect, self).__init__(**kwargs)
|
|
||||||
if size == "large":
|
|
||||||
self.up1 = UNetUp(round_channels(960, channel_multiplier), round_channels(112, channel_multiplier), 256, (14,14), old=1)
|
|
||||||
self.up2 = UNetUp(256, round_channels(40, channel_multiplier), 128, (28,28), old=1)
|
|
||||||
self.up3 = UNetUp(128, round_channels(24, channel_multiplier), 64, (56,56), old=1)
|
|
||||||
self.group = DSConv2d(64, 2, kernels_per_layer=4, groups=2, old=1)
|
|
||||||
self.r1_i = 1
|
|
||||||
self.r2_i = 2
|
|
||||||
self.r3_i = 4
|
|
||||||
elif size == "small":
|
|
||||||
self.up1 = UNetUp(round_channels(576, channel_multiplier), round_channels(40, channel_multiplier), 256, (14,14), old=1)
|
|
||||||
self.up2 = UNetUp(256, round_channels(24, channel_multiplier), 128, (28,28), old=1)
|
|
||||||
self.up3 = UNetUp(128, round_channels(16, channel_multiplier), 64, (56,56), old=1)
|
|
||||||
self.group = DSConv2d(64, 2, kernels_per_layer=4, groups=2, old=1)
|
|
||||||
self.r1_i = 0
|
|
||||||
self.r2_i = 1
|
|
||||||
self.r3_i = 2
|
|
||||||
self.maxpool = nn.MaxPool2d(kernel_size=3, dilation=1, stride=1, padding=1)
|
|
||||||
def _forward_impl(self, x):
|
|
||||||
x = self.conv_stem(x)
|
|
||||||
x = self.bn1(x)
|
|
||||||
x = self.act1(x)
|
|
||||||
r2 = None
|
|
||||||
r3 = None
|
|
||||||
for i, feature in enumerate(self.blocks):
|
|
||||||
x = feature(x)
|
|
||||||
if i == self.r3_i:
|
|
||||||
r3 = x
|
|
||||||
if i == self.r2_i:
|
|
||||||
r2 = x
|
|
||||||
if i == self.r1_i:
|
|
||||||
r1 = x
|
|
||||||
x = self.up1(x, r3)
|
|
||||||
x = self.up2(x, r2)
|
|
||||||
x = self.up3(x, r1)
|
|
||||||
x = self.group(x)
|
|
||||||
x2 = self.maxpool(x)
|
|
||||||
return x, x2
|
|
||||||
def forward(self, x):
|
|
||||||
return self._forward_impl(x)
|
|
||||||
|
|
||||||
def logit_arr(p, factor=16.0):
|
|
||||||
p = p.clamp(0.0000001, 0.9999999)
|
|
||||||
return torch.log(p / (1 - p)) / factor
|
|
||||||
|
|
||||||
# Landmark detection model
|
|
||||||
# Models:
|
|
||||||
# 0: "small", 0.5
|
|
||||||
# 1: "small", 1.0
|
|
||||||
# 2: "large", 0.75
|
|
||||||
# 3: "large", 1.0
|
|
||||||
class OpenSeeFaceLandmarks(geffnet.mobilenetv3.MobileNetV3):
|
|
||||||
def __init__(self, size="large", channel_multiplier=1.0, inference=False):
|
|
||||||
kwargs = geffnet.mobilenetv3._gen_mobilenet_v3([size], channel_multiplier=channel_multiplier)
|
|
||||||
super(OpenSeeFaceLandmarks, self).__init__(**kwargs)
|
|
||||||
if size == "large":
|
|
||||||
self.up1 = UNetUp(round_channels(960, channel_multiplier), round_channels(112, channel_multiplier), 256, (14,14))
|
|
||||||
self.up2 = UNetUp(256, round_channels(40, channel_multiplier), 198 * 1, (28,28))
|
|
||||||
self.group = DSConv2d(198 * 1, 198, kernels_per_layer=4, groups=3)
|
|
||||||
self.r2_i = 2
|
|
||||||
self.r3_i = 4
|
|
||||||
elif size == "small":
|
|
||||||
self.up1 = UNetUp(round_channels(576, channel_multiplier), round_channels(40, channel_multiplier), 256, (14,14))
|
|
||||||
self.up2 = UNetUp(256, round_channels(24, channel_multiplier), 198 * 1, (28,28))
|
|
||||||
self.group = DSConv2d(198 * 1, 198, kernels_per_layer=4, groups=3)
|
|
||||||
self.r2_i = 1
|
|
||||||
self.r3_i = 2
|
|
||||||
self.inference = inference
|
|
||||||
def _forward_impl(self, x):
|
|
||||||
x = self.conv_stem(x)
|
|
||||||
x = self.bn1(x)
|
|
||||||
x = self.act1(x)
|
|
||||||
r2 = None
|
|
||||||
r3 = None
|
|
||||||
for i, feature in enumerate(self.blocks):
|
|
||||||
x = feature(x)
|
|
||||||
if i == self.r3_i:
|
|
||||||
r3 = x
|
|
||||||
if i == self.r2_i:
|
|
||||||
r2 = x
|
|
||||||
x = self.up1(x, r3)
|
|
||||||
x = self.up2(x, r2)
|
|
||||||
x = self.group(x)
|
|
||||||
|
|
||||||
if self.inference:
|
|
||||||
t_main = x[:, 0:66].reshape((-1, 66, 28*28))
|
|
||||||
t_m = t_main.argmax(dim=2)
|
|
||||||
indices = t_m.unsqueeze(2)
|
|
||||||
t_conf = t_main.gather(2, indices).squeeze(2)
|
|
||||||
t_off_x = x[:, 66:132].reshape((-1, 66, 28*28)).gather(2, indices).squeeze(2)
|
|
||||||
t_off_y = x[:, 132:198].reshape((-1, 66, 28*28)).gather(2, indices).squeeze(2)
|
|
||||||
t_off_x = (223. * logit_arr(t_off_x) + 0.5).floor()
|
|
||||||
t_off_y = (223. * logit_arr(t_off_y) + 0.5).floor()
|
|
||||||
t_x = 223. * (t_m / 28.).floor() / 27. + t_off_x
|
|
||||||
t_y = 223. * t_m.remainder(28.).float() / 27. + t_off_y
|
|
||||||
x = (t_conf.mean(1), torch.stack([t_x, t_y, t_conf], 2))
|
|
||||||
|
|
||||||
return x
|
|
||||||
def forward(self, x):
|
|
||||||
return self._forward_impl(x)
|
|
||||||
|
|
||||||
# lm_modelT for 56x56 30 point inference
|
|
||||||
class OpenSeeFaceLandmarks30Pt(geffnet.mobilenetv3.MobileNetV3):
|
|
||||||
def __init__(self, size="large", channel_multiplier=1.0, inference=False):
|
|
||||||
kwargs = geffnet.mobilenetv3._gen_mobilenet_v3([size], channel_multiplier=channel_multiplier)
|
|
||||||
super(OpenSeeFaceLandmarks30Pt, self).__init__(**kwargs)
|
|
||||||
self.up1 = UNetUp(960, 112, 256, (4,4))
|
|
||||||
self.up2 = UNetUp(256, 40, 180, (7,7))
|
|
||||||
self.group = DSConv2d(180, 90, kernels_per_layer=4, groups=3)
|
|
||||||
self.inference = inference
|
|
||||||
def _forward_impl(self, x):
|
|
||||||
x = self.conv_stem(x)
|
|
||||||
x = self.bn1(x)
|
|
||||||
x = self.act1(x)
|
|
||||||
r2 = None
|
|
||||||
r3 = None
|
|
||||||
for i, feature in enumerate(self.blocks):
|
|
||||||
x = feature(x)
|
|
||||||
if i == 4:
|
|
||||||
r3 = x
|
|
||||||
if i == 2:
|
|
||||||
r2 = x
|
|
||||||
x = self.up1(x, r3)
|
|
||||||
x = self.up2(x, r2)
|
|
||||||
x = self.group(x)
|
|
||||||
|
|
||||||
if self.inference:
|
|
||||||
t_main = x[:, 0:30].reshape((-1, 30, 7*7))
|
|
||||||
t_m = t_main.argmax(dim=2)
|
|
||||||
indices = t_m.unsqueeze(2)
|
|
||||||
t_conf = t_main.gather(2, indices).squeeze(2)
|
|
||||||
t_off_x = x[:, 30:60].reshape((-1, 30, 7*7)).gather(2, indices).squeeze(2)
|
|
||||||
t_off_y = x[:, 60:90].reshape((-1, 30, 7*7)).gather(2, indices).squeeze(2)
|
|
||||||
t_off_x = 55. * logit_arr(t_off_x, factor=8.0)
|
|
||||||
t_off_y = 55. * logit_arr(t_off_y, factor=8.0)
|
|
||||||
t_x = 55. * (t_m / 7.).floor() / 6. + t_off_x
|
|
||||||
t_y = 55. * t_m.remainder(7.).float() / 6. + t_off_y
|
|
||||||
x = (t_conf.mean(1), torch.stack([t_x, t_y, t_conf], 2))
|
|
||||||
|
|
||||||
return x
|
|
||||||
def forward(self, x):
|
|
||||||
return self._forward_impl(x)
|
|
||||||
|
|
||||||
# Adaptive Wing Loss with offset layers and emphasis for eyes and eyebrows with 66 landmark points
|
|
||||||
def AdapWingLoss(pre_hm, gt_hm):
|
|
||||||
# pre_hm = pre_hm.to('cpu')
|
|
||||||
# gt_hm = gt_hm.to('cpu')
|
|
||||||
theta = 0.5
|
|
||||||
alpha = 2.1
|
|
||||||
w = 14
|
|
||||||
e = 1
|
|
||||||
A = w * (1 / (1 + torch.pow(theta / e, alpha - gt_hm))) * (alpha - gt_hm) * torch.pow(theta / e, alpha - gt_hm - 1) * (1 / e)
|
|
||||||
C = (theta * A - w * torch.log(1 + torch.pow(theta / e, alpha - gt_hm)))
|
|
||||||
|
|
||||||
batch_size = gt_hm.size()[0]
|
|
||||||
hm_num = gt_hm.size()[1]
|
|
||||||
|
|
||||||
mask = torch.zeros_like(gt_hm)
|
|
||||||
# W = 10
|
|
||||||
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
|
|
||||||
first = True
|
|
||||||
first_mask = None
|
|
||||||
for i in range(batch_size):
|
|
||||||
img_list = []
|
|
||||||
for j in range(hm_num // 3):
|
|
||||||
hm = np.round(gt_hm[i][j].cpu().numpy() * 255)
|
|
||||||
img_list.append(hm)
|
|
||||||
img_merge = cv2.merge(img_list)
|
|
||||||
img_dilate = cv2.morphologyEx(img_merge, cv2.MORPH_DILATE, kernel)
|
|
||||||
img_dilate[img_dilate < 51] = 1 # 0*W+1
|
|
||||||
img_dilate[img_dilate >= 51] = 11 # 1*W+1
|
|
||||||
img_dilate = np.array(img_dilate, dtype=int)
|
|
||||||
img_dilate = img_dilate.transpose(2, 0, 1)
|
|
||||||
dilated = torch.from_numpy(img_dilate).float()
|
|
||||||
if first:
|
|
||||||
first_mask = dilated
|
|
||||||
first = False
|
|
||||||
# These weights varied between training runs and model sizes
|
|
||||||
dilated[17:27][dilated[17:27] > 1] *= 1.4
|
|
||||||
dilated[17, dilated[17] > 1] *= 1.6
|
|
||||||
dilated[18, dilated[18] > 1] *= 1.8
|
|
||||||
dilated[25, dilated[25] > 1] *= 1.8
|
|
||||||
dilated[26, dilated[26] > 1] *= 1.6
|
|
||||||
dilated[36:48][dilated[36:48] > 1] *= 2.8
|
|
||||||
# Used for a very small model
|
|
||||||
#dilated[[37,38,40,41,43,44,46,47]][dilated[[37,38,40,41,43,44,46,47]] > 1] *= 20.8
|
|
||||||
mask[i] = torch.cat([dilated, dilated, dilated], 0)
|
|
||||||
|
|
||||||
diff_hm = torch.abs(gt_hm - pre_hm)
|
|
||||||
AWingLoss = A * diff_hm - C
|
|
||||||
idx = diff_hm < theta
|
|
||||||
AWingLoss[idx] = w * torch.log(1 + torch.pow(diff_hm / e, alpha - gt_hm))[idx]
|
|
||||||
|
|
||||||
AWingLoss *= mask
|
|
||||||
sum_loss = torch.sum(AWingLoss)
|
|
||||||
all_pixel = torch.sum(mask)
|
|
||||||
mean_loss = sum_loss / all_pixel
|
|
||||||
|
|
||||||
return first_mask.detach(), mean_loss
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
# Checkpoint test
|
|
||||||
if __name__== "__main__":
|
|
||||||
print("Checking gaze model")
|
|
||||||
m=OpenSeeFaceGaze()
|
|
||||||
ckpt = torch.load("gaze.pth")
|
|
||||||
m.load_state_dict(ckpt)
|
|
||||||
print("Checking detection model")
|
|
||||||
m=OpenSeeFaceDetect()
|
|
||||||
ckpt = torch.load("detection.pth")
|
|
||||||
m.load_state_dict(ckpt)
|
|
||||||
print("Checking lm_model0 model")
|
|
||||||
m=OpenSeeFaceLandmarks("small", 0.5)
|
|
||||||
ckpt = torch.load("lm_model0.pth")
|
|
||||||
m.load_state_dict(ckpt)
|
|
||||||
print("Checking lm_model1 model")
|
|
||||||
m=OpenSeeFaceLandmarks("small", 1.0)
|
|
||||||
ckpt = torch.load("lm_model1.pth")
|
|
||||||
m.load_state_dict(ckpt)
|
|
||||||
print("Checking lm_model2 model")
|
|
||||||
m=OpenSeeFaceLandmarks("large", 0.75)
|
|
||||||
ckpt = torch.load("lm_model2.pth")
|
|
||||||
m.load_state_dict(ckpt)
|
|
||||||
print("Checking lm_model3 model")
|
|
||||||
m=OpenSeeFaceLandmarks("large", 1.0)
|
|
||||||
ckpt = torch.load("lm_model3.pth")
|
|
||||||
m.load_state_dict(ckpt)
|
|
||||||
print("Checking lm_modelT model")
|
|
||||||
m=OpenSeeFaceLandmarks("large", 1.0)
|
|
||||||
ckpt = torch.load("lm_modelT.pth")
|
|
||||||
m.load_state_dict(ckpt)
|
|
||||||
82
remedian.py
82
remedian.py
|
|
@ -1,82 +0,0 @@
|
||||||
"""
|
|
||||||
Watch over a stream of numbers, incrementally learning their median.
|
|
||||||
|
|
||||||
Implemented via nested lists. New numbers are added to `lst[i]` and
|
|
||||||
when it fills up, it posts its median to `lst[i+1]`. Wen `lst[i+1]`
|
|
||||||
fills up, it posts the medians of its medians to `lst[i+2]`. Etc.
|
|
||||||
When a remedian is queried for the current median, it returns the
|
|
||||||
median of the last list with any numbers.
|
|
||||||
|
|
||||||
This approach is quite space efficient . E.g. four nested lists,
|
|
||||||
each with 64 items, require memory for 4*64 items yet can hold the
|
|
||||||
median of the median of the median of the median of over 17 million
|
|
||||||
numbers.
|
|
||||||
|
|
||||||
Example usage:
|
|
||||||
|
|
||||||
z=remedian()
|
|
||||||
for i in range(1000):
|
|
||||||
z + i
|
|
||||||
if not i % 100:
|
|
||||||
print(i, z.median())
|
|
||||||
|
|
||||||
Based on [The Remedian:A Robust Averaging Method for Large Data
|
|
||||||
Sets](http://web.ipac.caltech.edu/staff/fmasci/home/astro_refs/Remedian.pdf).
|
|
||||||
by Peter J. Rousseeuw and Gilbert W. Bassett Jr. Journal of the
|
|
||||||
American Statistical Association March 1990, Vol. 85, No. 409,
|
|
||||||
Theory and Methods
|
|
||||||
|
|
||||||
The code [remedianeg.py](remedianeg.py) compares this rig to just
|
|
||||||
using Python's built-in sort then reporing the middle number.
|
|
||||||
Assuming lists of length 64 and use of pypy3:
|
|
||||||
|
|
||||||
- Remedian is getting nearly as fast (within 20%) as raw sort after 500 items;
|
|
||||||
- While at the same time, avoids having to store all the numbers in RAM;
|
|
||||||
- Further, remedian's computed median is within 1% (or less) of the medians found via Python's sort.
|
|
||||||
|
|
||||||
_____
|
|
||||||
|
|
||||||
## Programmer's Guide
|
|
||||||
|
|
||||||
"""
|
|
||||||
|
|
||||||
# If `ordered` is `False`, do not sort `lst`
|
|
||||||
def median(lst,ordered=False):
|
|
||||||
assert lst,"median needs a non-empty list"
|
|
||||||
n = len(lst)
|
|
||||||
p = q = n//2
|
|
||||||
if n < 3:
|
|
||||||
p,q = 0, n-1
|
|
||||||
else:
|
|
||||||
lst = lst if ordered else sorted(lst)
|
|
||||||
if not n % 2: # for even-length lists, use mean of mid 2 nums
|
|
||||||
q = p -1
|
|
||||||
return lst[p] if p==q else (lst[p]+lst[q])/2
|
|
||||||
|
|
||||||
class remedian:
|
|
||||||
|
|
||||||
# Initialization
|
|
||||||
def __init__(i,inits=[], k=64, # after some experimentation, 64 works ok
|
|
||||||
about = None):
|
|
||||||
i.all,i.k = [],k
|
|
||||||
i.more,i._median=None,None
|
|
||||||
[i + x for x in inits]
|
|
||||||
|
|
||||||
# When full, push the median of current values to next list, then reset.
|
|
||||||
def __add__(i,x):
|
|
||||||
i._median = None
|
|
||||||
i.all.append(x)
|
|
||||||
if len(i.all) == i.k:
|
|
||||||
i.more = i.more or remedian(k=i.k)
|
|
||||||
i.more + i._medianPrim(i.all)
|
|
||||||
i.all = [] # reset
|
|
||||||
|
|
||||||
# If there is a next list, ask its median. Else, work it out locally.
|
|
||||||
def median(i):
|
|
||||||
return i.more.median() if i.more else i._medianPrim(i.all)
|
|
||||||
|
|
||||||
# Only recompute median if we do not know it already.
|
|
||||||
def _medianPrim(i,all):
|
|
||||||
if i._median == None:
|
|
||||||
i._median = median(all,ordered=False)
|
|
||||||
return i._median
|
|
||||||
150
retinaface.py
150
retinaface.py
|
|
@ -1,150 +0,0 @@
|
||||||
import sys
|
|
||||||
import numpy as np
|
|
||||||
import cv2
|
|
||||||
import onnxruntime
|
|
||||||
import time
|
|
||||||
import queue
|
|
||||||
import threading
|
|
||||||
import json
|
|
||||||
import copy
|
|
||||||
|
|
||||||
def py_cpu_nms(dets, thresh):
|
|
||||||
""" Pure Python NMS baseline.
|
|
||||||
Copyright (c) 2015 Microsoft
|
|
||||||
Licensed under The MIT License
|
|
||||||
Written by Ross Girshick
|
|
||||||
"""
|
|
||||||
x1 = dets[:, 0]
|
|
||||||
y1 = dets[:, 1]
|
|
||||||
x2 = dets[:, 2]
|
|
||||||
y2 = dets[:, 3]
|
|
||||||
scores = dets[:, 4]
|
|
||||||
|
|
||||||
areas = (x2 - x1 + 1) * (y2 - y1 + 1)
|
|
||||||
order = scores.argsort()[::-1]
|
|
||||||
|
|
||||||
keep = []
|
|
||||||
while order.size > 0:
|
|
||||||
i = order[0]
|
|
||||||
keep.append(i)
|
|
||||||
xx1 = np.maximum(x1[i], x1[order[1:]])
|
|
||||||
yy1 = np.maximum(y1[i], y1[order[1:]])
|
|
||||||
xx2 = np.minimum(x2[i], x2[order[1:]])
|
|
||||||
yy2 = np.minimum(y2[i], y2[order[1:]])
|
|
||||||
|
|
||||||
w = np.maximum(0.0, xx2 - xx1 + 1)
|
|
||||||
h = np.maximum(0.0, yy2 - yy1 + 1)
|
|
||||||
inter = w * h
|
|
||||||
ovr = inter / (areas[i] + areas[order[1:]] - inter)
|
|
||||||
|
|
||||||
inds = np.where(ovr <= thresh)[0]
|
|
||||||
order = order[inds + 1]
|
|
||||||
|
|
||||||
return keep
|
|
||||||
|
|
||||||
def decode(loc, priors, variances):
|
|
||||||
data = (
|
|
||||||
priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
|
|
||||||
priors[:, 2:] * np.exp(loc[:, 2:] * variances[1])
|
|
||||||
)
|
|
||||||
boxes = np.concatenate(data, 1)
|
|
||||||
boxes[:, :2] -= boxes[:, 2:] / 2
|
|
||||||
boxes[:, 2:] += boxes[:, :2]
|
|
||||||
return boxes
|
|
||||||
|
|
||||||
def worker_thread(rfd, frame):
|
|
||||||
results = rfd.detect_retina(frame, is_background=True)
|
|
||||||
rfd.results.put(results, False)
|
|
||||||
rfd.finished = True
|
|
||||||
rfd.running = False
|
|
||||||
|
|
||||||
class RetinaFaceDetector():
|
|
||||||
def __init__(self, model_path="models/retinaface_640x640_opt.onnx", json_path="models/priorbox_640x640.json", threads=4, min_conf=0.4, nms_threshold=0.4, top_k=1, res=(640, 640)):
|
|
||||||
options = onnxruntime.SessionOptions()
|
|
||||||
options.inter_op_num_threads = 1
|
|
||||||
options.intra_op_num_threads = threads
|
|
||||||
options.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL
|
|
||||||
options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
|
|
||||||
options.log_severity_level = 3
|
|
||||||
providersList = onnxruntime.capi._pybind_state.get_available_providers()
|
|
||||||
self.session = onnxruntime.InferenceSession(model_path, sess_options=options, providers=providersList)
|
|
||||||
self.res_w, self.res_h = res
|
|
||||||
with open(json_path, "r") as prior_file:
|
|
||||||
self.priorbox = np.array(json.loads(prior_file.read()))
|
|
||||||
self.min_conf = min_conf
|
|
||||||
self.nms_threshold = nms_threshold
|
|
||||||
self.top_k = top_k
|
|
||||||
self.finished = False
|
|
||||||
self.running = False
|
|
||||||
self.results = queue.Queue()
|
|
||||||
|
|
||||||
def detect_retina(self, frame, is_background=False):
|
|
||||||
h, w, _ = frame.shape
|
|
||||||
im = None
|
|
||||||
im = cv2.resize(frame, (self.res_w, self.res_h), interpolation=cv2.INTER_LINEAR)
|
|
||||||
resize_w = w / self.res_w
|
|
||||||
resize_w = 1 / resize_w
|
|
||||||
resize_h = h / self.res_h
|
|
||||||
resize_h = 1 / resize_h
|
|
||||||
im = np.float32(im)
|
|
||||||
scale = np.array((self.res_w / resize_w, self.res_h / resize_h, self.res_w / resize_w, self.res_h / resize_h))
|
|
||||||
im -= (104, 117, 123)
|
|
||||||
im = im.transpose(2, 0, 1)
|
|
||||||
im = np.expand_dims(im, 0)
|
|
||||||
output = self.session.run([], {"input0": im})
|
|
||||||
loc, conf = output[0][0], output[1][0]
|
|
||||||
boxes = decode(loc, self.priorbox, [0.1, 0.2])
|
|
||||||
boxes = boxes * scale
|
|
||||||
scores = conf[:, 1]
|
|
||||||
|
|
||||||
inds = np.where(scores > self.min_conf)[0]
|
|
||||||
boxes = boxes[inds]
|
|
||||||
scores = scores[inds]
|
|
||||||
|
|
||||||
dets = np.hstack((boxes, scores[:, np.newaxis])).astype(np.float32, copy=False)
|
|
||||||
keep = py_cpu_nms(dets, self.nms_threshold)
|
|
||||||
dets = dets[keep, :]
|
|
||||||
dets = dets[:self.top_k, 0:4]
|
|
||||||
dets[:, 2:4] = dets[:, 2:4] - dets[:, 0:2]
|
|
||||||
|
|
||||||
if True:#is_background:
|
|
||||||
upsize = dets[:, 2:4] * np.array([[0.15, 0.2]])
|
|
||||||
dets[:, 0:2] -= upsize
|
|
||||||
dets[:, 2:4] += upsize * 2
|
|
||||||
|
|
||||||
return list(map(tuple, dets))
|
|
||||||
|
|
||||||
def background_detect(self, frame):
|
|
||||||
if self.running or self.finished:
|
|
||||||
return
|
|
||||||
self.running = True
|
|
||||||
im = copy.copy(frame)
|
|
||||||
thread = threading.Thread(target=worker_thread, args=(self, im))
|
|
||||||
thread.start()
|
|
||||||
|
|
||||||
def get_results(self):
|
|
||||||
if self.finished:
|
|
||||||
results = []
|
|
||||||
try:
|
|
||||||
while True:
|
|
||||||
detection = self.results.get(False)
|
|
||||||
results.append(detection)
|
|
||||||
except:
|
|
||||||
"No error"
|
|
||||||
self.finished = False
|
|
||||||
return list(*results)
|
|
||||||
else:
|
|
||||||
return []
|
|
||||||
|
|
||||||
if __name__== "__main__":
|
|
||||||
retina = RetinaFaceDetector(top_k=40, min_conf=0.2)
|
|
||||||
im = cv2.imread(sys.argv[1], cv2.IMREAD_COLOR)
|
|
||||||
start = time.perf_counter()
|
|
||||||
faces = retina.detect_retina(im)
|
|
||||||
end = 1000 * (time.perf_counter() - start)
|
|
||||||
print(f"Runtime: {end:.3f}ms")
|
|
||||||
for (x,y,w,h) in faces:
|
|
||||||
im = cv2.rectangle(im, (int(x),int(y)), (int(x+w),int(y+w)), (0,0,255), 1)
|
|
||||||
cv2.imshow("Frame", im)
|
|
||||||
while cv2.waitKey(0) & 0xff != ord('q'):
|
|
||||||
""
|
|
||||||
File diff suppressed because it is too large
Load diff
Loading…
Add table
Add a link
Reference in a new issue