Update app.py

app.py

@@ -1,17 +1,35 @@
+import os
 import gradio as gr
 import pandas as pd
+from pydrive.auth import GoogleAuth
+from pydrive.drive import GoogleDrive
 
-# Function to handle the user's input and store it in the DataFrame
-def describe_image(image, desc1, desc2, desc3, csv_file_path):
-    # Check if the file exists, and if not, create it
-    if not pd.api.types.is_file_like(csv_file_path):
-        data = pd.DataFrame(columns=["Image", "Description 1", "Description 2", "Description 3"])
+# Load credentials from environment variables
+gauth = GoogleAuth()
+gauth.LocalWebserverAuth()
+
+# Initialize Google Drive
+drive = GoogleDrive(gauth)
+
+# Function to handle the user's input and store it in a Google Drive file
+def describe_image(image, desc1, desc2, desc3):
+    data = pd.DataFrame(columns=["Image", "Description 1", "Description 2", "Description 3"])
+    data = data.append({"Image": image, "Description 1": desc1, "Description 2": desc2, "Description 3": desc3}, ignore_index=True)
+
+    # Create or get a reference to the file on Google Drive
+    file_list = drive.ListFile({'q': "'" + os.environ.get("GDRIVE_FOLDER_ID") + "' in parents and trashed=false"}).GetList()
+    for file in file_list:
+        if file['title'] == os.environ.get("CSV_FILENAME"):
+            gdrive_file = file
+            break
     else:
-        data = pd.read_csv(csv_file_path)
+        gdrive_file = drive.CreateFile({'title': os.environ.get("CSV_FILENAME"), "parents": [{"kind": "drive#fileLink", "id": os.environ.get("GDRIVE_FOLDER_ID")}]})
+
+    # Write the data to the file and save it on Google Drive
+    data.to_csv('temp.csv', index=False)
+    gdrive_file.SetContentFile('temp.csv')
+    gdrive_file.Upload()
 
-    data = data.append({"Image": image, "Description 1": desc1, "Description 2": desc2, "Description 3": desc3},
-                       ignore_index=True)
-    data.to_csv(csv_file_path, index=False)  # Save the data to the specified CSV file
     return None  # Return None for no specific output
 
 iface = gr.Interface(
@@ -20,14 +38,17 @@ iface = gr.Interface(
         gr.Image(type="pil", label="Upload an Image"),
         "text",
         "text",
-        "text",
-        "text"  # Add an input field for specifying the CSV file path
+        "text"
     ],
-    outputs=gr.Info("Data saved successfully"),  # Display the success message
+    outputs=gr.Info("Data saved successfully on Google Drive"),  # Display the success message
     live=False,
     title="Image Description App",
     description="Upload an image and provide three descriptions for it.",
     allow_flagging=False,
+    theme=gr.themes.Base(primary_hue="green").set(
+        button_primary_background_fill="*primary_400",
+        button_primary_background_fill_hover="*primary_300",
+    )
 )
 
-iface.launch()
+iface.launch()