Update app.py

app.py

@@ -152,14 +152,19 @@ def run_showui(image, query, session_id, iterations=2):
 
     return images_during_iterations, str(click_xy)
 
-def save_and_upload_data(image_path, query, session_id, is_example_image, votes=None):
+def save_and_upload_data(image, query, session_id, is_example_image, votes=None):
     """Save the data to a JSON file and upload to S3."""
     if is_example_image == "True":
         return
 
     votes = votes or {"upvotes": 0, "downvotes": 0}
+
+    # Save image locally
+    image_file_name = f"{session_id}.png"
+    image.save(image_file_name)
+
     data = {
-        "image_path": image_path,
+        "image_path": image_file_name,
         "query": query,
         "votes": votes,
         "timestamp": datetime.now().isoformat()
@@ -170,8 +175,9 @@ def save_and_upload_data(image_path, query, session_id, is_example_image, votes=
     with open(local_file_name, "w") as f:
         json.dump(data, f)
     
+    # Upload JSON and image to S3
     upload_to_s3(local_file_name, 'altair.storage', object_name=f"ootb/{local_file_name}")
-    upload_to_s3(image_path, 'altair.storage', object_name=f"ootb/{os.path.basename(image_path)}")
+    upload_to_s3(image_file_name, 'altair.storage', object_name=f"ootb/{image_file_name}")
 
     return data