diff --git a/config.py b/config.py index 52ff09c..cc189ef 100644 --- a/config.py +++ b/config.py @@ -26,14 +26,14 @@ class Config: self.output_formatting() with open("./config.json", "w") as outfile: json.dump(self.__dict__, outfile) - return self + return self.__dict__ else: self.receive_data_path() self.set_pg_config() self.output_formatting() with open("./config.json", "w") as outfile: json.dump(self.__dict__, outfile) - return self + return self.__dict__ # if a config file already exists, offer to use it instead def handle_prev_config(self): @@ -80,15 +80,16 @@ class Config: try: dsn = f"dbname={dbname} user={user} password={password}" if host: - dsn = dsn + host + dsn = dsn + " host=" + host if port: - dsn = dsn + port + dsn = dsn + + " port=" + port conn = psycopg2.connect(f"dbname={dbname} user={user} password={password}") conn.close() print("Connection successful!") + self.pg_dsn = dsn except psycopg2.Error as e: print("There was an error connecting: ") diff --git a/formatresult.py b/formatresult.py index 008bb46..c82cf2f 100644 --- a/formatresult.py +++ b/formatresult.py @@ -1,4 +1,5 @@ import json, os, shutil +import psycopg2 from pgutil import PGUTIL from config import Config @@ -9,7 +10,8 @@ def format_result(app_config: Config, json_path): data_path = app_config['data_path'] # if pg_config is not None, run the postgres prediction[0] of this code - pgutil = PGUTIL(app_config, json_path) + if app_config['pg_dsn']: + pgutil = PGUTIL(app_config, json_path) # if this is True, run the prediction[0] "for line in contents:" below sort_by_match_strength = app_config['sort_by_match_strength'] @@ -39,11 +41,20 @@ def format_result(app_config: Config, json_path): match_strength = 'weak' # container for data to ship to pg - pgutil.insert_data({ - "filename": img_path, - "topprediction": guess_label, - "matchaccuracy": match_accuracy - }) + if pgutil: + conn = psycopg2.connect(app_config['pg_dsn']) + cur = conn.cursor() + + pgutil.insert_data(cur, { + "filename": img_path, + "topprediction": guess_label, + "matchaccuracy": match_accuracy + }) + + conn.commit() + cur.close() + conn.close() + print("Database insertions successful!") # assign value for match strength if float(match_accuracy) > 0.9: diff --git a/main.py b/main.py index 263bcaf..1ecaeb1 100644 --- a/main.py +++ b/main.py @@ -62,7 +62,7 @@ for file in files: json_path = "./predictions/predictions" + cur_time + ".json" -print("Writing analysis results to " + json_path + "\n\n") +print("\nWriting analysis results to " + json_path + "\n\n") # convert object to JSON and write to JSON file with open(json_path, "w") as outfile: diff --git a/pgutil.py b/pgutil.py index c8a91fa..4875400 100644 --- a/pgutil.py +++ b/pgutil.py @@ -5,38 +5,35 @@ class PGUTIL: def __init__(self, config_file: Config = None, json_path: str = None): if config_file is None: load_config = Config() - config_file = load_config.run() + self.config_file = load_config.run() - pg_dsn = config_file['pg_dsn'] + self.pg_dsn = config_file['pg_dsn'] - if pg_dsn is None: - raise Exception("Insufficient data to establish PostgreSQL connection.") if json_path: self.json_path = json_path - self.config_file = config_file + if config_file['pg_dsn'] is None: + raise Exception("Insufficient data to establish PostgreSQL connection.") def create_tables(self): - pg_dsn = self.config_file['pg_dsn'] - print("Connecting to PostgreSQL...") # establish connection and create cursor - conn = psycopg2.connect(pg_dsn) + conn = psycopg2.connect(self.__dict__['pg_dsn']) cur = conn.cursor() # create base tables create_label_table = """ CREATE TABLE IF NOT EXISTS label ( - id INT PRIMARY KEY, - name varchar + id INT GENERATED ALWAYS AS IDENTITY PRIMARY KEY, + name varchar UNIQUE ); """ create_photo_table = """ CREATE TABLE IF NOT EXISTS photo ( - id INT PRIMARY KEY, - path varchar, - label varchar, + id INT GENERATED ALWAYS AS IDENTITY PRIMARY KEY, + filename varchar, + guesslabel varchar, matchstrength decimal, labelid INT REFERENCES label(id) ); @@ -60,12 +57,29 @@ class PGUTIL: print("Connection closed.") - def insert_data(self, analysis): - # if self.json_path is None: - # return + def insert_data(self, cur, analysis): + self.create_tables() - # establish pg connection - conn = psycopg2.connect(self['config_file']['dsn']) - cur = conn.cursor() + insert_label = """ + INSERT INTO label (name) VALUES (%s) + ON CONFLICT (name) DO NOTHING; + """ - \ No newline at end of file + filename = analysis['filename'] + guesslabel = analysis['topprediction'] + matchstrength = analysis['matchaccuracy'] + + insert_photo = """ + INSERT INTO photo (filename, guesslabel, matchstrength, labelid) VALUES ( + %s, %s, %s, ( + SELECT id FROM label + WHERE name = %s + ) + ); + """ + + try: + cur.execute(insert_label, [guesslabel]) + cur.execute(insert_photo, [filename, guesslabel, matchstrength, guesslabel]) + except psycopg2.Error: + raise Exception("Problem inserting data for photo " + analysis['filename']) \ No newline at end of file