Update app.py

app.py

@@ -62,7 +62,7 @@ def array_to_image_path(image_array, session_id):
     if image_array is None:
         raise ValueError("No image provided. Please upload an image before submitting.")
     img = Image.fromarray(np.uint8(image_array))
-    filename = f"image_{session_id}.png"
+    filename = f"{session_id}.png"
     img.save(filename)
     return os.path.abspath(filename)
 
@@ -126,8 +126,12 @@ def run_showui(image, query, session_id):
     result_image = draw_point(image_path, click_xy, radius=10)
     return result_image, str(click_xy), image_path
 
-def save_and_upload_data(image_path, query, session_id, votes=None):
+def save_and_upload_data(image_path, query, session_id, is_example_image, votes=None):
     """Save the data to a JSON file and upload to S3."""
+    if is_example_image:
+        print("Example image used. Skipping upload.")
+        return
+
     votes = votes or {"upvotes": 0, "downvotes": 0}
     data = {
         "image_path": image_path,
@@ -136,18 +140,23 @@ def save_and_upload_data(image_path, query, session_id, votes=None):
         "timestamp": datetime.now().isoformat()
     }
     
-    local_file_name = f"data_{session_id}.json"
+    local_file_name = f"{session_id}.json"
     
     with open(local_file_name, "w") as f:
         json.dump(data, f)
     
     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)}")
 
     return data
 
-def update_vote(vote_type, session_id):
+def update_vote(vote_type, session_id, is_example_image):
     """Update the vote count and re-upload the JSON file."""
-    local_file_name = f"data_{session_id}.json"
+    if is_example_image:
+        print("Example image used. Skipping vote update.")
+        return "Example image used. No vote recorded."
+
+    local_file_name = f"{session_id}.json"
     
     with open(local_file_name, "r") as f:
         data = json.load(f)
@@ -171,6 +180,7 @@ def build_demo(embed_mode, concurrency_count=1):
     with gr.Blocks(title="ShowUI Demo", theme=gr.themes.Default()) as demo:
         state_image_path = gr.State(value=None)
         state_session_id = gr.State(value=None)
+        state_is_example_image = gr.State(value=False)
 
         if not embed_mode:
             gr.HTML(
@@ -240,35 +250,37 @@ def build_demo(embed_mode, concurrency_count=1):
                     raise ValueError("No image provided. Please upload an image before submitting.")
                 
                 session_id = datetime.now().strftime("%Y%m%d_%H%M%S")
+                is_example_image = isinstance(image, str) and image.startswith("./examples/")
+                
                 result_image, click_coords, image_path = run_showui(image, query, session_id)
                 
-                save_and_upload_data(image_path, query, session_id)
+                save_and_upload_data(image_path, query, session_id, is_example_image)
                 
-                return result_image, click_coords, image_path, session_id
+                return result_image, click_coords, image_path, session_id, is_example_image
 
             submit_btn.click(
                 on_submit,
                 [imagebox, textbox],
-                [output_img, output_coords, state_image_path, state_session_id],
+                [output_img, output_coords, state_image_path, state_session_id, state_is_example_image],
             )
 
             clear_btn.click(
                 lambda: (None, None, None, None, None),
                 inputs=None,
-                outputs=[imagebox, textbox, output_img, output_coords, state_image_path, state_session_id],
+                outputs=[imagebox, textbox, output_img, output_coords, state_image_path, state_session_id, state_is_example_image],
                 queue=False
             )
 
             upvote_btn.click(
-                lambda session_id: update_vote("upvote", session_id),
-                inputs=state_session_id,
+                lambda session_id, is_example_image: update_vote("upvote", session_id, is_example_image),
+                inputs=[state_session_id, state_is_example_image],
                 outputs=[],
                 queue=False
             )
 
             downvote_btn.click(
-                lambda session_id: update_vote("downvote", session_id),
-                inputs=state_session_id,
+                lambda session_id, is_example_image: update_vote("downvote", session_id, is_example_image),
+                inputs=[state_session_id, state_is_example_image],
                 outputs=[],
                 queue=False
             )