add image save function

app.py

@@ -12,6 +12,11 @@ from utils import SCENARIOS, PROMPTS, IMAGES, OPTIONS, T2I_MODELS, INSTRUCTION,
 import spaces #[uncomment to use ZeroGPU]
 import transformers
 import gspread
+from googleapiclient.discovery import build
+from googleapiclient.http import MediaFileUpload
+from googleapiclient.errors import HttpError
+from google.oauth2.service_account import Credentials
+
 
 CLIP_MODEL = "ViT-H-14"
 PRETRAINED_CLIP = "laion2b_s32b_b79k"
@@ -36,13 +41,14 @@ DEFAULT_SCENARIO = "Product advertisement"
 METHODS = ["Baseline", "Experimental"]
 MAX_ROUND = 5
 
-# intermittent memory
 counter1, counter2 = 1, 1
 responses_memory = {}
 assigned_scenarios = list(SCENARIOS.keys())[:2]
 current_task1, current_task2 = METHODS # current task 1 (tab 1)
 task1_success, task2_success = False, False
 enable_submit1, enable_submit2 = False, False
+scopes = ['https://www.googleapis.com/auth/spreadsheets', 'https://www.googleapis.com/auth/drive']
+
 
 ########################################################################################################
 # Generating images with two methods
@@ -293,9 +299,20 @@ def redesign(participant, scenario, prompt, sim_radio, like_radio, dislike_radio
     else:
         return {submit_btn1: gr.skip()} if active_tab == "Task A" else {submit_btn2: gr.skip()}
 
+def save_image(creds, like_image, dislike_image, name):
+    try:
+        service = build("drive", "v3", credentials=creds)
+        for image_path, suffix in zip([like_image, dislike_image], ["satisfied", "unsatisfied"]):
+            filename = f"{name}_{suffix}"
+            file_metadata = {"name": filename, "parents": ["1ru3-QbbzyVSk-1kBfVv4nhElFqYh3ITj"]}
+            media = MediaFileUpload(image_path, mimetype="image/png")
+            uploaded_file = service.files().create(body=file_metadata, media_body=media, fields="id").execute()
+    
+    except HttpError as error:
+        print(f"An error occurred: {error}")    
 
-def save_response(participant, scenario, prompt, sim_radio, like_radio, dislike_radio, active_tab):
-    global current_task1, current_task2 # not change
+def save_response(participant, scenario, prompt, sim_radio, like_radio, dislike_radio, like_image, dislike_image, active_tab):
+    global current_task1, current_task2, scopes # not change
     global task1_success, task2_success, counter1, counter2, enable_submit1, enable_submit2, responses_memory, assigned_scenarios # will change
 
     method = current_task1 if active_tab == "Task A" else current_task2
@@ -310,11 +327,16 @@ def save_response(participant, scenario, prompt, sim_radio, like_radio, dislike_
         responses_memory[participant][method][counter]["unsatisfied_img"] = f"round {counter}, {dislike_radio}"
         
         try:
+            # save responses in google sheet
             gc = gspread.service_account(filename='credentials.json')
             sheet = gc.open("DiverseGen-phase3").sheet1 
-
+            
             for i, entry in responses_memory[participant][method].items():
                 sheet.append_row([participant, scenario, f"{active_tab}, {method}", i, entry["prompt"], entry["sim_radio"], entry["response"], entry["satisfied_img"], entry["unsatisfied_img"]])
+            
+            # save images in google drive
+            creds = Credentials.from_service_account_file('credentials.json',scopes=scopes)
+            save_image(creds, like_image, dislike_image, f"{participant}_{scenario}_{active_tab}_{method}")
 
             display_info_message("✅ Your answer is saved!")
 
@@ -330,7 +352,6 @@ def save_response(participant, scenario, prompt, sim_radio, like_radio, dislike_
             else:
                 task2_success = True
 
-
             # decide if change scenario
             # if scenario == assigned_scenarios[0]:
             #     next_scenario = assigned_scenarios[1] if task1_success and task2_success else assigned_scenarios[0]
@@ -426,11 +447,11 @@ with gr.Blocks(theme=gr.themes.Soft(font=[gr.themes.GoogleFont("Inconsolata"), "
         
             with gr.Row(elem_id="compact-row"):
                 with gr.Column(elem_id="col-container"):
-                    images_method1 = gr.Gallery(label="Images", columns=[4], rows=[1], height=200, elem_id="gallery")
-                    history_images1 = gr.Gallery(label="History Images", columns=[4], rows=[1], elem_id="gallery")
+                    images_method1 = gr.Gallery(label="Images", columns=[4], rows=[1], height=200, elem_id="gallery", type="png")
+                    history_images1 = gr.Gallery(label="History Images", columns=[4], rows=[1], elem_id="gallery", type="png")
                 with gr.Column(elem_id="col-container3"):
-                    like_image1 = gr.Image(label="Satisfied Image", width=200, height=200, sources='upload', type="pil")
-                    dislike_image1 = gr.Image(label="Unsatisfied Image", width=200, height=200, sources='upload', type="pil")
+                    like_image1 = gr.Image(label="Satisfied Image", width=200, height=200, sources='upload', format="png", type="filepath")
+                    dislike_image1 = gr.Image(label="Unsatisfied Image", width=200, height=200, sources='upload', format="png", type="filepath")
             with gr.Column(elem_id="col-container2"):
                 gr.Markdown("### 📝 Evaluation")               
                 sim_radio1 = gr.Radio(
@@ -481,11 +502,11 @@ with gr.Blocks(theme=gr.themes.Soft(font=[gr.themes.GoogleFont("Inconsolata"), "
 
             with gr.Row(elem_id="compact-row"):
                 with gr.Column(elem_id="col-container"):
-                    images_method2 = gr.Gallery(label="Images", columns=[4], rows=[1], height=200, elem_id="gallery")
-                    history_images2 = gr.Gallery(label="History Images", columns=[4], rows=[1], elem_id="gallery")
+                    images_method2 = gr.Gallery(label="Images", columns=[4], rows=[1], height=200, elem_id="gallery", type="png")
+                    history_images2 = gr.Gallery(label="History Images", columns=[4], rows=[1], elem_id="gallery", type="png")
                 with gr.Column(elem_id="col-container3"):
-                    like_image2 = gr.Image(label="Satisfied Image", width=200, height=200, sources='upload', type="pil")
-                    dislike_image2 = gr.Image(label="Unsatisfied Image", width=200, height=200, sources='upload', type="pil")
+                    like_image2 = gr.Image(label="Satisfied Image", width=200, height=200, sources='upload', format="png", type="filepath")
+                    dislike_image2 = gr.Image(label="Unsatisfied Image", width=200, height=200, sources='upload', format="png", type="filepath")
 
             with gr.Column(elem_id="col-container2"):
                 gr.Markdown("### 📝 Evaluation")
@@ -552,11 +573,11 @@ with gr.Blocks(theme=gr.themes.Soft(font=[gr.themes.GoogleFont("Inconsolata"), "
         outputs=[sim_radio2, dislike_radio2, like_radio2, images_method2, history_images2, example2.dataset, prompt2, next_btn2, redesign_btn2, submit_btn2]
     )
     submit_btn1.click(fn=save_response, 
-        inputs=[participant, scenario, prompt1, sim_radio1, like_radio1, dislike_radio1, active_tab], 
+        inputs=[participant, scenario, prompt1, sim_radio1, like_radio1, dislike_radio1, like_image1, dislike_image1, active_tab], 
         outputs=[sim_radio1, dislike_radio1, like_radio1, prompt1, next_btn1, redesign_btn1, submit_btn1, tabs])
     
     submit_btn2.click(fn=save_response, 
-        inputs=[participant, scenario, prompt2, sim_radio2, like_radio2, dislike_radio2, active_tab], 
+        inputs=[participant, scenario, prompt2, sim_radio2, like_radio2, dislike_radio2, like_image2, dislike_image2, active_tab], 
         outputs=[sim_radio2,  dislike_radio2, like_radio2, prompt2, next_btn2, redesign_btn2, submit_btn2, tabs])
 
 

requirements.txt

@@ -10,4 +10,7 @@ sentence-transformers
 ftfy
 mediapy
 open-clip-torch==2.24.0
-gspread
+gspread
+google-api-python-client
+google-auth-httplib2 
+google-auth-oauthlib