feat: time & status to data.jsonl

GoogleDrive_API.py

@@ -5,19 +5,39 @@ from googleapiclient.http import MediaFileUpload, MediaIoBaseDownload
 
 
 class GoogleDrive_API:
-    def __init__(self):
+    def __init__(self, SERVICE_ACCOUNT_FILE_path: str = ""):
         self.SCOPES = ["https://www.googleapis.com/auth/drive"]
-        self.SERVICE_ACCOUNT_FILE = "service_account.json"
         self.PARENT_FOLDER_ID = "1r-MlnEpWHx3b1fxHDnHcZ2-Wh_Y89676"
-        self.service = self.authenticate()
+        self.service = self.authenticate(SERVICE_ACCOUNT_FILE_path)
         self.clear_files()
 
-    def authenticate(self):
-        credentials = service_account.Credentials.from_service_account_file(
-            self.SERVICE_ACCOUNT_FILE, scopes=self.SCOPES
-        )
-        service = build("drive", "v3", credentials=credentials)
+    def authenticate(self, SERVICE_ACCOUNT_FILE):
+        if len(SERVICE_ACCOUNT_FILE) > 0:
+            if not os.path.exists(SERVICE_ACCOUNT_FILE):
+                raise rf"Google SERVICE_ACCOUNT_FILE Not Found, Path = {SERVICE_ACCOUNT_FILE}"
 
+            credentials = service_account.Credentials.from_service_account_file(
+                SERVICE_ACCOUNT_FILE, scopes=self.SCOPES
+            )
+        else:
+            service_account_info = {
+                "type": "service_account",
+                "project_id": os.environ["project_id"],
+                "private_key_id": os.environ["private_key_id"],
+                "private_key": os.environ["private_key"],
+                "client_email": os.environ["client_email"],
+                "client_id": os.environ["client_id"],
+                "auth_uri": "https://accounts.google.com/o/oauth2/auth",
+                "token_uri": "https://oauth2.googleapis.com/token",
+                "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
+                "client_x509_cert_url": os.environ["client_x509_cert_url"],
+                "universe_domain": "googleapis.com",
+            }
+            credentials = service_account.Credentials.from_service_account_info(
+                service_account_info, scopes=self.SCOPES
+            )
+
+        service = build("drive", "v3", credentials=credentials)
         return service
 
     def get_files(self):
@@ -40,7 +60,7 @@ class GoogleDrive_API:
             file_id = item["id"]
             self.service.files().delete(fileId=file_id).execute()
 
-    def upload_file(self, file_name, file_path):
+    def upload_file(self, file_name: str, file_path: str):
         file_metadata = {
             "name": file_name,
             "parents": [self.PARENT_FOLDER_ID],
@@ -53,7 +73,7 @@ class GoogleDrive_API:
         )
         print(rf"{file_path} uploaded. ID:", file.get("id"))
 
-    def download_file(self, file_name, file_path):
+    def download_file(self, file_name: str, file_path: str):
         results = (
             self.service.files()
             .list(

Kaggle_API.py

@@ -3,16 +3,17 @@ import shutil
 import subprocess
 import json
 import time
+from datetime import datetime
 from PIL import Image
 
-from GoogleDrive_API import GoogleDrive_API
-
 
 class API_Connection:
-    def __init__(self, kaggle_username: str = "", kaggle_key=""):
+    def __init__(self, GD_connection, kaggle_username: str = "", kaggle_key: str = ""):
         os.environ["KAGGLE_USERNAME"] = kaggle_username
         os.environ["KAGGLE_KEY"] = kaggle_key
 
+        self.GoogleDrive_connection = GD_connection
+
         self.PROJECT_PATH = r""
         self.NOTEBOOK_ID = "amirmoris/pix2pix"
         self.DATASET_NAME = "dataset"
@@ -94,8 +95,6 @@ class API_Connection:
         if len(output_image_name) == 0:
             return False, rf"Missing Output"
 
-        GoogleDrive_connection = GoogleDrive_API()
-
         dataset_path = self.correct_path(rf"{self.PROJECT_PATH}\{self.DATASET_NAME}")
         notebook_path = self.correct_path(rf"{self.PROJECT_PATH}\notebook")
 
@@ -109,6 +108,8 @@ class API_Connection:
 
         data = [
             {
+                "status": "IDLE",
+                "time": self.get_current_time(),
                 "edit": edit_instruction,
                 "input_image_path": input_image_name,
                 "output_image_path": output_image_name,
@@ -117,10 +118,10 @@ class API_Connection:
 
         self.write_file(data, dataset_path, "data.jsonl")
         # update dataset
-        GoogleDrive_connection.upload_file(
+        self.GoogleDrive_connection.upload_file(
             "data.jsonl", rf"{self.DATASET_NAME}\data.jsonl"
         )
-        GoogleDrive_connection.upload_file(
+        self.GoogleDrive_connection.upload_file(
             input_image_name, rf"{self.DATASET_NAME}\{input_image_name}"
         )
 
@@ -142,7 +143,7 @@ class API_Connection:
             time.sleep(120)
 
         # get output
-        GoogleDrive_connection.download_file(
+        self.GoogleDrive_connection.download_file(
             output_image_name, rf"{dataset_path}\{output_image_name}"
         )
         output_image = self.read_image(rf"{dataset_path}\{output_image_name}")
@@ -153,6 +154,9 @@ class API_Connection:
 
         return True, output_image
 
+    def get_current_time(self):
+        return str(datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
+
 
 def main():
     pass

app.py

@@ -4,6 +4,7 @@ from PIL import Image
 import gradio as gr
 
 from Kaggle_API import API_Connection
+from GoogleDrive_API import GoogleDrive_API
 
 
 def generate_button_clicked(input_image: Image = None, edit_instruction: str = ""):
@@ -16,7 +17,9 @@ def generate_button_clicked(input_image: Image = None, edit_instruction: str = "
     kaggle_username = os.environ["kaggle_username"]
     kaggle_key = os.environ["kaggle_key"]
 
-    api_connection = API_Connection(kaggle_username, kaggle_key)
+    GoogleDrive_connection = GoogleDrive_API()
+    api_connection = API_Connection(GoogleDrive_connection, kaggle_username, kaggle_key)
+
     api_connection.create_folder(rf"{api_connection.PROJECT_PATH}\local_dataset")
 
     image_ID = ""