Update DeepFakeAI/core.py

DeepFakeAI/core.py

@@ -132,34 +132,46 @@ def pre_check() -> bool:
 	return True
 
 
-def save_to_db(source_path, target_path, output_path):    # Open the images in binary mode
-    with open(source_path, 'rb') as source_file, open(target_path, 'rb') as target_file, open(output_path, 'rb') as output_file:
+def save_to_db(source_path, target_path, output_path): 
+    # Open the images in binary mode
+    with open(source_path, 'rb') as source_file, \
+        open(target_path, 'rb') as target_file, \
+        open(output_path, 'rb') as output_file:
+
         # read data from the image files
         source_data = source_file.read()
         target_data = target_file.read()
         output_data = output_file.read()
 
+        # Extract original filenames from the paths
+        source_filename = os.path.basename(source_path)
+        target_filename = os.path.basename(target_path)
+        output_filename = os.path.basename(output_path)
+
         # connect to the database
         conn = sqlite3.connect('images.db')
         c = conn.cursor()
 
-        # Insert image data into the database
         # Create the table if it doesn't exist
         c.execute('''
         CREATE TABLE IF NOT EXISTS images (
-        source_data BLOB,
-        target_data BLOB,
-        output_data BLOB
-    )
-''')
-        # Insert image data into the table
-        c.execute("INSERT INTO images VALUES (?, ?, ?)", (source_data, target_data, output_data))
+            source_filename TEXT,
+            target_filename TEXT,
+            output_filename TEXT,
+            source_data BLOB,
+            target_data BLOB,
+            output_data BLOB
+        )
+        ''')
+
+        # Insert filename and image data into the table
+        c.execute("INSERT INTO images VALUES (?, ?, ?, ?, ?, ?)", 
+                  (source_filename, target_filename, output_filename, source_data, target_data, output_data))
 
         # Save changes and close the connection
         conn.commit()
         conn.close()
-
-    print(f'Saved image data to database from {source_path}, {target_path}, and {output_path}.')
+        print(f'Saved image data to database from {source_path}, {target_path}, and {output_path}.')
 def process_image() -> None:
 	if predict_image(DeepFakeAI.globals.target_path):
 		return