diff --git a/scripts/collect_images.py b/scripts/collect_images.py new file mode 100644 index 00000000..572ba299 --- /dev/null +++ b/scripts/collect_images.py @@ -0,0 +1,98 @@ +import argparse +import datetime +import json +import os +import subprocess + +from src.db.db_connection import sql_cursor + +""" +This script should be set up with a cron job to run daily. +Each user will have to set up their rclone config, in our case pointing to a Google Drive folder: https://rclone.org/drive/ +This can be done with /local/b/embedvis/rclone-v1.52.2-linux-amd64/rclone config + +Collect images of non-goggle detections from the database. +Upload images and metadata to Google Drive. +""" + +METADATA_FILE = 'metadata.json' + + +def get_metadata(): + """ + Get image filenames and other relevant metadata from the database. + Save metadata to a file for future decryption. + @return: A list of dictionaries with the metadata for each image + + Example list: [ + {'image_name': "0.jpg", 'x_min': 0.0, 'y_min': 0.0, 'x_max': 100.0, 'y_max': 100.0, 'init_vector': "example"} + {'image_name': "1.jpg", 'x_min': 25.0, 'y_min': 25.0, 'x_max': 120.0, 'y_max': 140.0, 'init_vector': "example2"}] + """ + + metadata = [] + current_date = (datetime.date.today(),) + + # make sql connection + # execute query + with sql_cursor() as cursor: + try: + cursor.execute('USE goggles') + cursor.execute('SELECT b.image_name, b.X_Min, b.Y_Min, b.X_Max, b.Y_Max, ' + 'b.init_vector, b.goggles from BBOX AS b, IMAGE as i where ' + 'b.image_name=i.image_name and i.image_date=%s and b.goggles=False', current_date) + + for (image_name, x_min, y_min, x_max, y_max, init_vector, goggles) in cursor: + metadata.append({'image_name': image_name, + 'x_min': float(x_min), + 'y_min': float(y_min), + 'x_max': float(x_max), + 'y_max': float(y_max), + 'init_vector': init_vector + }) + except Exception as e: + print(e) + + with open(METADATA_FILE, 'w') as meta_file: + json.dump(metadata, meta_file) + return metadata + + +def upload_files(metadata, dir, rclone_path, remote_name): + """ + For each filename returned by get_metadata, upload image + to Drive. Upload the day's metadata file. + @param metadata: the list of dictionaries returned by get_metadata + @param dir: the folder containing the images to upload + @param remote_name: name of remote location in rclone + """ + + images = [] + today = datetime.datetime.today().strftime('%Y-%m-%d') + + # send images to the Drive + for image in metadata: + # prevent sending the same image twice (if two faces are detected) + if image not in images: + images.append(image) + image_path = os.path.join(dir, image['image_name']) + subprocess.run([rclone_path, 'copy', image_path, '{}:{}'.format(remote_name, today)]) + + # upload metadata.json to the Drive + subprocess.run([rclone_path, 'copy', METADATA_FILE, '{}:{}'.format(remote_name, today)]) + os.remove(METADATA_FILE) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser('Collect images.') + parser.add_argument('--directory', '-d', type=str, required=True, help='Folder containing images to upload') + parser.add_argument('--rclone_path', '-r', type=str, default='/local/b/embedvis/rclone-v1.52.2-linux-amd64/rclone', + help='Location of rclone binary. Default version on ee220clnx1 doesn\'t support copying to ' + 'shared folders.') + parser.add_argument('--remote_name', type=str, help='Name of remote location according to rclone config. You must ' + 'create your own config.') + args = parser.parse_args() + + metadata = get_metadata() + upload_files(metadata, args.directory, args.rclone_path, args.remote_name) + + exit(0) diff --git a/scripts/decrypt_images.py b/scripts/decrypt_images.py new file mode 100644 index 00000000..f81e6cf4 --- /dev/null +++ b/scripts/decrypt_images.py @@ -0,0 +1,42 @@ +import argparse +import getpass +import json +import os + +from src.jetson.AES import Encryption + +""" +After having collect_images has run and the output folder has been downloaded, +decrypt the associated images. +This file is assumed to be on the end user's machine. +""" + +# the metadata file generated by collect_images +METADATA_FILE = 'metadata.json' + + +def decrypt_images(dir): + # ask for decryption key + decrypt_key = getpass.getpass('Decryption password: ') + + # convert to PKBDF2 or whatever + + # make decryptor; probably changes once Jason finishes + decryptor = Encryption + + with open(os.path.join(dir, METADATA_FILE)) as meta_file: + metadata = json.load(meta_file) + # use face coords to find where to decrypt in video frame + for image in metadata: + # TODO handle multiple faces in one frame. append to coords list + coords = [(image['x_min'], image['y_min'], image['x_max'], image['y_max'])] + init_vector = image['init_vector'] + # overwrite encrypted image + + +if __name__ == "__main__": + parser = argparse.ArgumentParser('Decrypt images.') + parser.add_argument('--directory', '-d', type=str, required=True, help='Folder of images to be decrypted.') + args = parser.parse_args() + + decrypt_images(args.directory) diff --git a/src/db/data_insertion.py b/src/db/data_insertion.py index f6b3b6d9..4f7ab1c4 100644 --- a/src/db/data_insertion.py +++ b/src/db/data_insertion.py @@ -3,6 +3,9 @@ from decimal import Decimal import datetime +# location where the images will be stored (on the HELPS machine) +IMAGE_DIR = '/local/b/embedvis/Nano_Images' + def data_insert(image_name: str, image_date: datetime, image_time: datetime, init_vecs: list, bboxes: list, input_dir: str, labels: list): """Transfer image to remote storage then inserts image metadata and bounding boxes data in database @@ -18,8 +21,8 @@ def data_insert(image_name: str, image_date: datetime, image_time: datetime, ini """ # Below ftp transfer has been commented out for testing purposes - #with ftp_transfer() as transfer: - #transfer(input_dir, './Documents', image_name) + with ftp_transfer() as transfer: + transfer(input_dir, IMAGE_DIR, image_name) sql_insert(IMAGE(image_name, image_date, image_time)) diff --git a/src/db/db_connection.py b/src/db/db_connection.py index 75373f08..d719104c 100644 --- a/src/db/db_connection.py +++ b/src/db/db_connection.py @@ -1,5 +1,4 @@ import mysql.connector -import datetime from src.db.config import get_config from contextlib import contextmanager, closing diff --git a/src/jetson/face_detector.py b/src/jetson/face_detector.py index ec48863b..2e571cbb 100644 --- a/src/jetson/face_detector.py +++ b/src/jetson/face_detector.py @@ -1,3 +1,5 @@ +import os + import numpy as np import torch @@ -41,7 +43,7 @@ def __init__(self, detector: str, detector_type: str, detection_threshold=0.7, c self.net = BlazeFace(self.device == torch.device("cuda:0")) self.net.load_weights(detector) - self.net.load_anchors("models/BlazeFace/anchors.npy") + self.net.load_anchors(os.path.join(os.path.dirname(__file__), 'models/BlazeFace/anchors.npy')) self.model_name = 'blazeface' self.net.min_score_thresh = 0.75 self.net.min_suppression_threshold = 0.3 @@ -110,7 +112,7 @@ def detect(self, xmax = detections[i, 3] * frame.shape[1] conf = detections[i, 16] - transformed_frame = img / 127.5 - 1.0 + transformed_frame = transformed_frame / 127.5 - 1.0 for k in range(6): kp_x = detections[i, 4 + k * 2] * transformed_frame.shape[1]