Init

GoogleDrive_API.py

@@ -0,0 +1,80 @@
+import os
+from google.oauth2 import service_account
+from googleapiclient.discovery import build
+from googleapiclient.http import MediaFileUpload, MediaIoBaseDownload
+
+
+class GoogleDrive_API:
+    def __init__(self):
+        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.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)
+
+        return service
+
+    def get_files(self):
+        # List all files in the folder
+        results = (
+            self.service.files()
+            .list(
+                q=f"'{self.PARENT_FOLDER_ID}' in parents and trashed=false",
+                fields="files(id, name)",
+            )
+            .execute()
+        )
+        return results.get("files", [])
+
+    def clear_files(self):
+        items = self.get_files()
+
+        # Delete each file
+        for item in items:
+            file_id = item["id"]
+            self.service.files().delete(fileId=file_id).execute()
+
+    def upload_file(self, file_name, file_path):
+        file_metadata = {
+            "name": file_name,
+            "parents": [self.PARENT_FOLDER_ID],
+        }
+        media = MediaFileUpload(file_path)
+        file = (
+            self.service.files()
+            .create(body=file_metadata, media_body=media, fields="id")
+            .execute()
+        )
+        print(rf"{file_path} uploaded. ID:", file.get("id"))
+
+    def download_file(self, file_name, file_path):
+        results = (
+            self.service.files()
+            .list(
+                q=f"'{self.PARENT_FOLDER_ID}' in parents and name='{file_name}' and trashed=false",
+                fields="files(id)",
+            )
+            .execute()
+        )
+        items = results.get("files", [])
+
+        if items:
+            # Get the file ID
+            file_id = items[0]["id"]
+
+            request = self.service.files().get_media(fileId=file_id)
+
+            with open(file_path, "wb") as file:
+                downloader = MediaIoBaseDownload(file, request)
+                done = False
+                while not done:
+                    status, done = downloader.next_chunk()
+                    print(f"Download {int(status.progress() * 100)}%.")
+        else:
+            print(rf"Fine {file_name} is Not Found")

Kaggle_API.py

@@ -0,0 +1,162 @@
+import os
+import shutil
+import subprocess
+import json
+import time
+from PIL import Image
+
+from GoogleDrive_API import GoogleDrive_API
+
+
+class API_Connection:
+    def __init__(self, kaggle_username: str = "", kaggle_key=""):
+        os.environ["KAGGLE_USERNAME"] = kaggle_username
+        os.environ["KAGGLE_KEY"] = kaggle_key
+
+        self.PROJECT_PATH = r""
+        self.NOTEBOOK_ID = "amirmoris/pix2pix"
+        self.DATASET_NAME = "dataset"
+
+    def execute_terminal_command(self, command: str):
+        try:
+            process = subprocess.Popen(
+                command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+            )
+            output, error = process.communicate()
+            output = str(output.decode("utf-8"))
+            error = str(error.decode("utf-8"))
+            print(rf"Command executed successfully: {command}")
+            return output
+        except Exception as e:
+            return None, str(e)
+
+    def correct_path(self, path: str):
+        return path[1:] if path.startswith("\\") else path
+
+    def create_folder(self, path: str):
+        path = self.correct_path(path)
+        if os.path.exists(path):
+            shutil.rmtree(path)
+        self.execute_terminal_command(rf"mkdir {path}")
+
+    def pull_kaggle_notebook(self, notebook_path: str):
+        command = rf"kaggle kernels pull {self.NOTEBOOK_ID} -p {notebook_path} -m"
+        return self.execute_terminal_command(command)
+
+    def push_kaggle_notebook(self, notebook_path: str):
+        command = rf"kaggle kernels push -p {notebook_path}"
+        return self.execute_terminal_command(command)
+
+    def get_notebook_status(self):
+        command = rf"kaggle kernels status {self.NOTEBOOK_ID}"
+        return self.execute_terminal_command(command)
+
+    def run(self, notebook_path: str):
+        notebook_path = self.correct_path(notebook_path)
+        self.pull_kaggle_notebook(notebook_path)
+        return self.push_kaggle_notebook(notebook_path)
+
+    def write_file(self, data: list, file_path: str, file_name: str = ""):
+        if len(file_name) > 0:
+            file_path = rf"{file_path}\{file_name}"
+
+        file_path = self.correct_path(file_path)
+        # Writing JSON data
+        with open(file_path, "w") as file:
+            for idx in range(len(data)):
+                json_string = json.dumps(data[idx]) + (
+                    "\n" if idx < len(data) - 1 else ""
+                )
+                file.write(json_string)
+
+    def read_image(self, image_path: str):
+        try:
+            image = Image.open(image_path)
+            return image
+        except IOError:
+            print("Unable to load image")
+            return None
+
+    def get_notebook_output(self, output_path: str):
+        output_path = self.correct_path(output_path)
+        command = rf"kaggle kernels output {self.NOTEBOOK_ID} -p {output_path}"
+        return self.execute_terminal_command(command)
+
+    def generate_image(
+        self,
+        input_image_name: str,
+        edit_instruction: str,
+        output_image_name: str,
+    ):
+        if len(input_image_name) == 0 or len(edit_instruction) == 0:
+            return False, rf"Missing Input"
+
+        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")
+
+        self.create_folder(dataset_path)
+
+        # copy image to the dataset
+        shutil.copyfile(
+            rf"local_dataset\{input_image_name}",
+            rf"{dataset_path}\{input_image_name}",
+        )
+
+        data = [
+            {
+                "edit": edit_instruction,
+                "input_image_path": input_image_name,
+                "output_image_path": output_image_name,
+            }
+        ]
+
+        self.write_file(data, dataset_path, "data.jsonl")
+        # update dataset
+        GoogleDrive_connection.upload_file(
+            "data.jsonl", rf"{self.DATASET_NAME}\data.jsonl"
+        )
+        GoogleDrive_connection.upload_file(
+            input_image_name, rf"{self.DATASET_NAME}\{input_image_name}"
+        )
+
+        # run notebook
+        print(self.run(notebook_path))
+
+        number_of_checks = 0
+        while True:
+            status = str(self.get_notebook_status()).replace("\n", "")
+            print(rf"- status no #{number_of_checks} : {status}")
+            number_of_checks += 1
+            if "complete" in status:
+                break
+
+            if "error" in status:
+                return False, "notebook status error"
+            if "cancelAcknowledged" in status:
+                return False, "notebook status cancelAcknowledged"
+            time.sleep(120)
+
+        # get output
+        GoogleDrive_connection.download_file(
+            output_image_name, rf"{dataset_path}\{output_image_name}"
+        )
+        output_image = self.read_image(rf"{dataset_path}\{output_image_name}")
+        # clear input and output
+
+        if output_image is None:
+            return False, "An error occured while running, no output image found"
+
+        return True, output_image
+
+
+def main():
+    pass
+
+
+if __name__ == "__main__":
+    main()

app.py

@@ -0,0 +1,89 @@
+import os
+import random
+from PIL import Image
+import gradio as gr
+
+from Kaggle_API import API_Connection
+
+
+def generate_button_clicked(input_image: Image = None, edit_instruction: str = ""):
+    if input_image is None:
+        raise gr.Error("Missing Input: input_image")
+    if edit_instruction == 0:
+        raise gr.Error("Missing Input: edit_instruction")
+
+    # set kaggle-api variables
+    kaggle_username = os.environ["kaggle_username"]
+    kaggle_key = os.environ["kaggle_key"]
+
+    api_connection = API_Connection(kaggle_username, kaggle_key)
+    api_connection.create_folder(rf"{api_connection.PROJECT_PATH}\local_dataset")
+
+    image_ID = ""
+    while len(image_ID) < 4:
+        image_ID += str(random.randint(0, 9))
+
+    input_image_name = rf"input_image_{image_ID}.png"
+    output_image_name = rf"output_image_{image_ID}.png"
+
+    input_image.save(rf"local_dataset\{input_image_name}")
+
+    status, img = api_connection.generate_image(
+        input_image_name, edit_instruction, output_image_name
+    )
+
+    if status == False:
+        raise gr.Error(img)  # img represents the error
+
+    return img
+
+
+def reset_button_clicked():
+    return None, "", None
+
+
+SCOPES = ["https://www.googleapis.com/auth/drive"]
+SERVICE_ACCOUNT_FILE = "service_account.json"
+PARENT_FOLDER_ID = "1r-MlnEpWHx3b1fxHDnHcZ2-Wh_Y89676"
+
+
+def main():
+    with gr.Blocks(css="footer {visibility: hidden}") as demo:
+        with gr.Row():
+            input_image = gr.Image(label="Input Image", type="pil", interactive=True)
+            edited_image = gr.Image(
+                label=f"Edited Image", type="pil", interactive=False
+            )
+
+        with gr.Row():
+            with gr.Column(scale=3):
+                instruction = gr.Textbox(
+                    lines=1, label="Edit Instruction", interactive=True
+                )
+            with gr.Column(scale=1, min_width=100):
+                with gr.Row():
+                    generate_button = gr.Button("Generate")
+                with gr.Row():
+                    reset_button = gr.Button("Reset")
+
+        generate_button.click(
+            fn=generate_button_clicked,
+            inputs=[
+                input_image,
+                instruction,
+            ],
+            outputs=edited_image,
+        )
+        reset_button.click(
+            fn=reset_button_clicked,
+            outputs=[input_image, instruction, edited_image],
+        )
+
+    # Launch Gradio interface
+
+    demo.queue(max_size=1)
+    demo.launch(share=True)
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,3 @@
+gradio
+kaggle
+google-api-python-client