-
Notifications
You must be signed in to change notification settings - Fork 2
/
create_db_visualization.py
39 lines (29 loc) · 1012 Bytes
/
create_db_visualization.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
import cv2
import shutil
from tqdm import tqdm
from pathlib import Path
import db
from config import ROOTPATH
from tools import reid_img_revert
def main():
# Create directory.
vis_path = Path.joinpath(ROOTPATH, "visualisation")
if vis_path.exists():
shutil.rmtree(vis_path)
vis_path.mkdir(exist_ok=True)
# Read images from DB.
all_images_data = db.select_all_images_data()
for image_data in tqdm(all_images_data):
write_image(image_data, vis_path=vis_path)
def write_image(image_data: dict, vis_path: Path) -> None:
""" Convert and write image. """
# Create person directory.
image_path = Path.joinpath(vis_path, str(image_data["person_id"]))
if not image_path.exists():
image_path.mkdir()
# Convert image to normal format and write it.
image = reid_img_revert(image_data["image"])
image_filename = f"{image_data['date']}.jpg"
cv2.imwrite(f"{image_path}/{image_filename}", image)
if __name__ == "__main__":
main()