mirror of
https://github.com/iperov/DeepFaceLive
synced 2025-08-19 21:13:21 -07:00
refactoring
This commit is contained in:
iperov
parent
8489949f2c
commit
30ba51edf7
24 changed files with 663 additions and 459 deletions
|
|
@ -1,19 +1,22 @@
|
|||
import pickle
|
||||
import sqlite3
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from typing import Generator, List, Union, Iterable
|
||||
from typing import Generator, Iterable, List, Union
|
||||
|
||||
import cv2
|
||||
import h5py
|
||||
import numpy as np
|
||||
|
||||
from .. import console as lib_con
|
||||
from .FMask import FMask
|
||||
from .UFaceMark import UFaceMark
|
||||
from .UImage import UImage
|
||||
from .UPerson import UPerson
|
||||
|
||||
|
||||
class Faceset:
|
||||
|
||||
def __init__(self, path = None):
|
||||
def __init__(self, path = None, write_access=False, recreate=False):
|
||||
"""
|
||||
Faceset is a class to store and manage face related data.
|
||||
|
||||
|
|
@ -21,205 +24,155 @@ class Faceset:
|
|||
|
||||
path path to faceset .dfs file
|
||||
|
||||
write_access
|
||||
|
||||
recreate
|
||||
|
||||
Can be pickled.
|
||||
"""
|
||||
self._f = None
|
||||
|
||||
self._path = path = Path(path)
|
||||
if path.suffix != '.dfs':
|
||||
raise ValueError('Path must be a .dfs file')
|
||||
|
||||
self._conn = conn = sqlite3.connect(path, isolation_level=None)
|
||||
self._cur = cur = conn.cursor()
|
||||
if path.exists():
|
||||
if write_access and recreate:
|
||||
path.unlink()
|
||||
elif not write_access:
|
||||
raise FileNotFoundError(f'File {path} not found.')
|
||||
|
||||
cur = self._get_cursor()
|
||||
cur.execute('BEGIN IMMEDIATE')
|
||||
if not self._is_table_exists('FacesetInfo'):
|
||||
self.recreate(shrink=False, _transaction=False)
|
||||
cur.execute('COMMIT')
|
||||
self.shrink()
|
||||
else:
|
||||
cur.execute('END')
|
||||
self._mode = 'a' if write_access else 'r'
|
||||
self._open()
|
||||
|
||||
def __del__(self):
|
||||
self.close()
|
||||
|
||||
def __getstate__(self):
|
||||
return {'_path' : self._path}
|
||||
return {'_path' : self._path, '_mode' : self._mode}
|
||||
|
||||
def __setstate__(self, d):
|
||||
self.__init__( d['_path'] )
|
||||
self._f = None
|
||||
self._path = d['_path']
|
||||
self._mode = d['_mode']
|
||||
self._open()
|
||||
|
||||
def __repr__(self): return self.__str__()
|
||||
def __str__(self):
|
||||
return f"Faceset. UImage:{self.get_UImage_count()} UFaceMark:{self.get_UFaceMark_count()} UPerson:{self.get_UPerson_count()}"
|
||||
|
||||
def _is_table_exists(self, name):
|
||||
return self._cur.execute(f"SELECT count(*) FROM sqlite_master WHERE type='table' AND name=?", [name]).fetchone()[0] != 0
|
||||
def _open(self):
|
||||
if self._f is None:
|
||||
self._f = f = h5py.File(self._path, mode=self._mode)
|
||||
self._UFaceMark_grp = f.require_group('UFaceMark')
|
||||
self._UImage_grp = f.require_group('UImage')
|
||||
self._UImage_image_data_grp = f.require_group('UImage_image_data')
|
||||
self._UPerson_grp = f.require_group('UPerson')
|
||||
|
||||
def _get_cursor(self) -> sqlite3.Cursor: return self._cur
|
||||
|
||||
def close(self):
|
||||
if self._cur is not None:
|
||||
self._cur.close()
|
||||
self._cur = None
|
||||
if self._f is not None:
|
||||
self._f.close()
|
||||
self._f = None
|
||||
|
||||
if self._conn is not None:
|
||||
self._conn.close()
|
||||
self._conn = None
|
||||
|
||||
def shrink(self):
|
||||
self._cur.execute('VACUUM')
|
||||
|
||||
def recreate(self, shrink=True, _transaction=True):
|
||||
def optimize(self, verbose=True):
|
||||
"""
|
||||
delete all data and recreate Faceset structure.
|
||||
recreate Faceset with optimized structure.
|
||||
"""
|
||||
cur = self._get_cursor()
|
||||
if verbose:
|
||||
print(f'Optimizing {self._path.name}...')
|
||||
|
||||
if _transaction:
|
||||
cur.execute('BEGIN IMMEDIATE')
|
||||
tmp_path = self._path.parent / (self._path.stem + '_optimizing' + self._path.suffix)
|
||||
|
||||
for table_name, in cur.execute("SELECT name from sqlite_master where type = 'table';").fetchall():
|
||||
cur.execute(f'DROP TABLE {table_name}')
|
||||
tmp_fs = Faceset(tmp_path, write_access=True, recreate=True)
|
||||
self._group_copy(tmp_fs._UFaceMark_grp, self._UFaceMark_grp, verbose=verbose)
|
||||
self._group_copy(tmp_fs._UPerson_grp, self._UPerson_grp, verbose=verbose)
|
||||
self._group_copy(tmp_fs._UImage_grp, self._UImage_grp, verbose=verbose)
|
||||
self._group_copy(tmp_fs._UImage_image_data_grp, self._UImage_image_data_grp, verbose=verbose)
|
||||
tmp_fs.close()
|
||||
|
||||
(cur.execute('CREATE TABLE FacesetInfo (version INT)')
|
||||
.execute('INSERT INTO FacesetInfo VALUES (1)')
|
||||
self.close()
|
||||
self._path.unlink()
|
||||
tmp_path.rename(self._path)
|
||||
self._open()
|
||||
|
||||
.execute('CREATE TABLE UImage (uuid BLOB, name TEXT, format TEXT, data BLOB)')
|
||||
.execute('CREATE TABLE UPerson (uuid BLOB, data BLOB)')
|
||||
.execute('CREATE TABLE UFaceMark (uuid BLOB, UImage_uuid BLOB, UPerson_uuid BLOB, data BLOB)')
|
||||
)
|
||||
def _group_copy(self, group_dst : h5py.Group, group_src : h5py.Group, verbose=True):
|
||||
for key, value in lib_con.progress_bar_iterator(group_src.items(), desc=f'Copying {group_src.name} -> {group_dst.name}', suppress_print=not verbose):
|
||||
d = group_dst.create_dataset(key, shape=value.shape, dtype=value.dtype )
|
||||
d[:] = value[:]
|
||||
for a_key, a_value in value.attrs.items():
|
||||
d.attrs[a_key] = a_value
|
||||
|
||||
if _transaction:
|
||||
cur.execute('COMMIT')
|
||||
def _group_read_bytes(self, group : h5py.Group, key : str, check_key=True) -> Union[bytes, None]:
|
||||
if check_key and key not in group:
|
||||
return None
|
||||
dataset = group[key]
|
||||
data_bytes = bytearray(len(dataset))
|
||||
dataset.read_direct(np.frombuffer(data_bytes, dtype=np.uint8))
|
||||
return data_bytes
|
||||
|
||||
if shrink:
|
||||
self.shrink()
|
||||
def _group_write_bytes(self, group : h5py.Group, key : str, data : bytes, update_existing=True) -> Union[h5py.Dataset, None]:
|
||||
if key in group:
|
||||
if not update_existing:
|
||||
return None
|
||||
del group[key]
|
||||
|
||||
return group.create_dataset(key, data=np.frombuffer(data, dtype=np.uint8) )
|
||||
|
||||
###################
|
||||
### UFaceMark
|
||||
###################
|
||||
def _UFaceMark_from_db_row(self, db_row) -> UFaceMark:
|
||||
uuid, UImage_uuid, UPerson_uuid, data = db_row
|
||||
|
||||
ufm = UFaceMark()
|
||||
ufm.restore_state(pickle.loads(data))
|
||||
return ufm
|
||||
|
||||
def add_UFaceMark(self, ufacemark_or_list : UFaceMark):
|
||||
def add_UFaceMark(self, ufacemark_or_list : UFaceMark, update_existing=True):
|
||||
"""
|
||||
add or update UFaceMark in DB
|
||||
"""
|
||||
if not isinstance(ufacemark_or_list, Iterable):
|
||||
ufacemark_or_list : List[UFaceMark] = [ufacemark_or_list]
|
||||
|
||||
cur = self._cur
|
||||
cur.execute('BEGIN IMMEDIATE')
|
||||
for ufm in ufacemark_or_list:
|
||||
uuid = ufm.get_uuid()
|
||||
UImage_uuid = ufm.get_UImage_uuid()
|
||||
UPerson_uuid = ufm.get_UPerson_uuid()
|
||||
|
||||
data = pickle.dumps(ufm.dump_state())
|
||||
|
||||
if cur.execute('SELECT COUNT(*) from UFaceMark where uuid=?', [uuid] ).fetchone()[0] != 0:
|
||||
cur.execute('UPDATE UFaceMark SET UImage_uuid=?, UPerson_uuid=?, data=? WHERE uuid=?',
|
||||
[UImage_uuid, UPerson_uuid, data, uuid])
|
||||
else:
|
||||
cur.execute('INSERT INTO UFaceMark VALUES (?, ?, ?, ?)', [uuid, UImage_uuid, UPerson_uuid, data])
|
||||
cur.execute('COMMIT')
|
||||
self._group_write_bytes(self._UFaceMark_grp, ufm.get_uuid().hex(), pickle.dumps(ufm.dump_state()), update_existing=update_existing )
|
||||
|
||||
def get_UFaceMark_count(self) -> int:
|
||||
return self._cur.execute('SELECT COUNT(*) FROM UFaceMark').fetchone()[0]
|
||||
return len(self._UFaceMark_grp.keys())
|
||||
|
||||
def get_all_UFaceMark(self) -> List[UFaceMark]:
|
||||
return [ self._UFaceMark_from_db_row(db_row) for db_row in self._cur.execute('SELECT * FROM UFaceMark').fetchall() ]
|
||||
|
||||
return [ UFaceMark.from_state(pickle.loads(self._group_read_bytes(self._UFaceMark_grp, key, check_key=False))) for key in self._UFaceMark_grp.keys() ]
|
||||
|
||||
def get_all_UFaceMark_uuids(self) -> List[bytes]:
|
||||
return [ uuid.UUID(key).bytes for key in self._UFaceMark_grp.keys() ]
|
||||
|
||||
def get_UFaceMark_by_uuid(self, uuid : bytes) -> Union[UFaceMark, None]:
|
||||
c = self._cur.execute('SELECT * FROM UFaceMark WHERE uuid=?', [uuid])
|
||||
db_row = c.fetchone()
|
||||
if db_row is None:
|
||||
data = self._group_read_bytes(self._UFaceMark_grp, uuid.hex())
|
||||
if data is None:
|
||||
return None
|
||||
return UFaceMark.from_state(pickle.loads(data))
|
||||
|
||||
return self._UFaceMark_from_db_row(db_row)
|
||||
|
||||
def delete_UFaceMark_by_uuid(self, uuid : bytes) -> bool:
|
||||
key = uuid.hex()
|
||||
if key in self._UFaceMark_grp:
|
||||
del self._UFaceMark_grp[key]
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def iter_UFaceMark(self) -> Generator[UFaceMark, None, None]:
|
||||
"""
|
||||
returns Generator of UFaceMark
|
||||
"""
|
||||
for db_row in self._cur.execute('SELECT * FROM UFaceMark').fetchall():
|
||||
yield self._UFaceMark_from_db_row(db_row)
|
||||
for key in self._UFaceMark_grp.keys():
|
||||
yield UFaceMark.from_state(pickle.loads(self._group_read_bytes(self._UFaceMark_grp, key, check_key=False)))
|
||||
|
||||
def delete_all_UFaceMark(self):
|
||||
"""
|
||||
deletes all UFaceMark from DB
|
||||
"""
|
||||
(self._cur.execute('BEGIN IMMEDIATE')
|
||||
.execute('DELETE FROM UFaceMark')
|
||||
.execute('COMMIT') )
|
||||
###################
|
||||
### UPerson
|
||||
###################
|
||||
def _UPerson_from_db_row(self, db_row) -> UPerson:
|
||||
uuid, data = db_row
|
||||
up = UPerson()
|
||||
up.restore_state(pickle.loads(data))
|
||||
return up
|
||||
|
||||
def add_UPerson(self, uperson_or_list : UPerson):
|
||||
"""
|
||||
add or update UPerson in DB
|
||||
"""
|
||||
if not isinstance(uperson_or_list, Iterable):
|
||||
uperson_or_list : List[UPerson] = [uperson_or_list]
|
||||
|
||||
cur = self._cur
|
||||
cur.execute('BEGIN IMMEDIATE')
|
||||
for uperson in uperson_or_list:
|
||||
uuid = uperson.get_uuid()
|
||||
|
||||
data = pickle.dumps(uperson.dump_state())
|
||||
|
||||
if cur.execute('SELECT COUNT(*) from UPerson where uuid=?', [uuid]).fetchone()[0] != 0:
|
||||
cur.execute('UPDATE UPerson SET data=? WHERE uuid=?', [data])
|
||||
else:
|
||||
cur.execute('INSERT INTO UPerson VALUES (?, ?)', [uuid, data])
|
||||
cur.execute('COMMIT')
|
||||
|
||||
def get_UPerson_count(self) -> int:
|
||||
return self._cur.execute('SELECT COUNT(*) FROM UPerson').fetchone()[0]
|
||||
|
||||
def get_all_UPerson(self) -> List[UPerson]:
|
||||
return [ self._UPerson_from_db_row(db_row) for db_row in self._cur.execute('SELECT * FROM UPerson').fetchall() ]
|
||||
|
||||
def iter_UPerson(self) -> Generator[UPerson, None, None]:
|
||||
"""
|
||||
iterator of all UPerson's
|
||||
"""
|
||||
for db_row in self._cur.execute('SELECT * FROM UPerson').fetchall():
|
||||
yield self._UPerson_from_db_row(db_row)
|
||||
|
||||
def delete_all_UPerson(self):
|
||||
"""
|
||||
deletes all UPerson from DB
|
||||
"""
|
||||
(self._cur.execute('BEGIN IMMEDIATE')
|
||||
.execute('DELETE FROM UPerson')
|
||||
.execute('COMMIT') )
|
||||
for key in self._UFaceMark_grp.keys():
|
||||
del self._UFaceMark_grp[key]
|
||||
|
||||
###################
|
||||
### UImage
|
||||
###################
|
||||
def _UImage_from_db_row(self, db_row) -> UImage:
|
||||
uuid, name, format, data_bytes = db_row
|
||||
img = cv2.imdecode(np.frombuffer(data_bytes, dtype=np.uint8), flags=cv2.IMREAD_UNCHANGED)
|
||||
|
||||
uimg = UImage()
|
||||
uimg.set_uuid(uuid)
|
||||
uimg.set_name(name)
|
||||
uimg.assign_image(img)
|
||||
return uimg
|
||||
|
||||
def add_UImage(self, uimage_or_list : UImage, format : str = 'webp', quality : int = 100):
|
||||
def add_UImage(self, uimage_or_list : UImage, format : str = 'png', quality : int = 100, update_existing=True):
|
||||
"""
|
||||
add or update UImage in DB
|
||||
|
||||
|
|
@ -239,9 +192,8 @@ class Faceset:
|
|||
raise ValueError('quality must be in range [0..100]')
|
||||
|
||||
if not isinstance(uimage_or_list, Iterable):
|
||||
uimage_or_list = [uimage_or_list]
|
||||
uimage_or_list : List[UImage] = [uimage_or_list]
|
||||
|
||||
uimage_datas = []
|
||||
for uimage in uimage_or_list:
|
||||
if format == 'webp':
|
||||
imencode_args = [int(cv2.IMWRITE_WEBP_QUALITY), quality]
|
||||
|
|
@ -251,44 +203,112 @@ class Faceset:
|
|||
imencode_args = [int(cv2.IMWRITE_JPEG2000_COMPRESSION_X1000), quality*10]
|
||||
else:
|
||||
imencode_args = []
|
||||
|
||||
ret, data_bytes = cv2.imencode( f'.{format}', uimage.get_image(), imencode_args)
|
||||
if not ret:
|
||||
raise Exception(f'Unable to encode image format {format}')
|
||||
uimage_datas.append(data_bytes.data)
|
||||
|
||||
cur = self._cur
|
||||
cur.execute('BEGIN IMMEDIATE')
|
||||
for uimage, data in zip(uimage_or_list, uimage_datas):
|
||||
uuid = uimage.get_uuid()
|
||||
if cur.execute('SELECT COUNT(*) from UImage where uuid=?', [uuid] ).fetchone()[0] != 0:
|
||||
cur.execute('UPDATE UImage SET name=?, format=?, data=? WHERE uuid=?', [uimage.get_name(), format, data, uuid])
|
||||
else:
|
||||
cur.execute('INSERT INTO UImage VALUES (?, ?, ?, ?)', [uuid, uimage.get_name(), format, data])
|
||||
cur.execute('COMMIT')
|
||||
key = uimage.get_uuid().hex()
|
||||
|
||||
def get_UImage_count(self) -> int: return self._cur.execute('SELECT COUNT(*) FROM UImage').fetchone()[0]
|
||||
def get_UImage_by_uuid(self, uuid : Union[bytes, None]) -> Union[UImage, None]:
|
||||
"""
|
||||
"""
|
||||
if uuid is None:
|
||||
self._group_write_bytes(self._UImage_grp, key, pickle.dumps(uimage.dump_state(exclude_image=True)), update_existing=update_existing )
|
||||
d = self._group_write_bytes(self._UImage_image_data_grp, key, data_bytes.data, update_existing=update_existing )
|
||||
d.attrs['format'] = format
|
||||
d.attrs['quality'] = quality
|
||||
|
||||
def get_UImage_count(self) -> int:
|
||||
return len(self._UImage_grp.keys())
|
||||
|
||||
def get_all_UImage(self) -> List[UImage]:
|
||||
return [ self._get_UImage_by_key(key) for key in self._UImage_grp.keys() ]
|
||||
|
||||
def get_all_UImage_uuids(self) -> List[bytes]:
|
||||
return [ uuid.UUID(key).bytes for key in self._UImage_grp.keys() ]
|
||||
|
||||
def _get_UImage_by_key(self, key, check_key=True) -> Union[UImage, None]:
|
||||
data = self._group_read_bytes(self._UImage_grp, key, check_key=check_key)
|
||||
if data is None:
|
||||
return None
|
||||
uimg = UImage.from_state(pickle.loads(data))
|
||||
|
||||
db_row = self._cur.execute('SELECT * FROM UImage where uuid=?', [uuid]).fetchone()
|
||||
if db_row is None:
|
||||
return None
|
||||
return self._UImage_from_db_row(db_row)
|
||||
image_data = self._group_read_bytes(self._UImage_image_data_grp, key, check_key=check_key)
|
||||
if image_data is not None:
|
||||
uimg.assign_image (cv2.imdecode(np.frombuffer(image_data, dtype=np.uint8), flags=cv2.IMREAD_UNCHANGED))
|
||||
|
||||
def iter_UImage(self) -> Generator[UImage, None, None]:
|
||||
return uimg
|
||||
|
||||
def get_UImage_by_uuid(self, uuid : bytes) -> Union[UImage, None]:
|
||||
return self._get_UImage_by_key(uuid.hex())
|
||||
|
||||
def delete_UImage_by_uuid(self, uuid : bytes):
|
||||
key = uuid.hex()
|
||||
if key in self._UImage_grp:
|
||||
del self._UImage_grp[key]
|
||||
if key in self._UImage_image_data_grp:
|
||||
del self._UImage_image_data_grp[key]
|
||||
|
||||
def iter_UImage(self, include_key=False) -> Generator[UImage, None, None]:
|
||||
"""
|
||||
iterator of all UImage's
|
||||
returns Generator of UImage
|
||||
"""
|
||||
for db_row in self._cur.execute('SELECT * FROM UImage').fetchall():
|
||||
yield self._UImage_from_db_row(db_row)
|
||||
for key in self._UImage_grp.keys():
|
||||
uimg = self._get_UImage_by_key(key, check_key=False)
|
||||
yield (uimg, key) if include_key else uimg
|
||||
|
||||
def delete_all_UImage(self):
|
||||
"""
|
||||
deletes all UImage from DB
|
||||
"""
|
||||
(self._cur.execute('BEGIN IMMEDIATE')
|
||||
.execute('DELETE FROM UImage')
|
||||
.execute('COMMIT') )
|
||||
for key in self._UImage_grp.keys():
|
||||
del self._UImage_grp[key]
|
||||
for key in self._UImage_image_data_grp.keys():
|
||||
del self._UImage_image_data_grp[key]
|
||||
|
||||
###################
|
||||
### UPerson
|
||||
###################
|
||||
|
||||
def add_UPerson(self, uperson_or_list : UPerson, update_existing=True):
|
||||
"""
|
||||
add or update UPerson in DB
|
||||
"""
|
||||
if not isinstance(uperson_or_list, Iterable):
|
||||
uperson_or_list : List[UPerson] = [uperson_or_list]
|
||||
|
||||
for uperson in uperson_or_list:
|
||||
self._group_write_bytes(self._UPerson_grp, uperson.get_uuid().hex(), pickle.dumps(uperson.dump_state()), update_existing=update_existing )
|
||||
|
||||
def get_UPerson_count(self) -> int:
|
||||
return len(self._UPerson_grp.keys())
|
||||
|
||||
def get_all_UPerson(self) -> List[UPerson]:
|
||||
return [ UPerson.from_state(pickle.loads(self._group_read_bytes(self._UPerson_grp, key, check_key=False))) for key in self._UPerson_grp.keys() ]
|
||||
|
||||
def get_all_UPerson_uuids(self) -> List[bytes]:
|
||||
return [ uuid.UUID(key).bytes for key in self._UPerson_grp.keys() ]
|
||||
|
||||
def get_UPerson_by_uuid(self, uuid : bytes) -> Union[UPerson, None]:
|
||||
data = self._group_read_bytes(self._UPerson_grp, uuid.hex())
|
||||
if data is None:
|
||||
return None
|
||||
return UPerson.from_state(pickle.loads(data))
|
||||
|
||||
def delete_UPerson_by_uuid(self, uuid : bytes) -> bool:
|
||||
key = uuid.hex()
|
||||
if key in self._UPerson_grp:
|
||||
del self._UPerson_grp[key]
|
||||
return True
|
||||
return False
|
||||
|
||||
def iter_UPerson(self) -> Generator[UPerson, None, None]:
|
||||
"""
|
||||
returns Generator of UPerson
|
||||
"""
|
||||
for key in self._UPerson_grp.keys():
|
||||
yield UPerson.from_state(pickle.loads(self._group_read_bytes(self._UPerson_grp, key, check_key=False)))
|
||||
|
||||
def delete_all_UPerson(self):
|
||||
"""
|
||||
deletes all UPerson from DB
|
||||
"""
|
||||
for key in self._UPerson_grp.keys():
|
||||
del self._UPerson_grp[key]
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue