Rewrite Game State handling

app.py

@@ -1,9 +1,7 @@
 import os
 import random
 from dataclasses import dataclass
-from threading import Lock
 from typing import List, Optional
-
 import gradio as gr
 import numpy as np
 import torch
@@ -25,8 +23,6 @@ class GameResults:
 
 class GameState:
     def __init__(self):
-        # Lock only used for round transitions
-        self.round_lock = Lock()
         self.reset()
     
     def reset(self):
@@ -55,25 +51,23 @@ class GameState:
         )
     
     def start_submission(self) -> bool:
-        with self.round_lock:
-            if not self.can_submit_guess():
-                return False
-            self.processing_submission = True
-            return True
+        if not self.can_submit_guess():
+            return False
+        self.processing_submission = True
+        return True
     
     def finish_submission(self, results: GameResults):
-        with self.round_lock:
-            if results.user_guess == results.correct_answer:
-                self.user_score += 1
-            if results.model_guess == results.correct_answer:
-                self.model_score += 1
-            
-            self.last_results = results
-            self.current_round += 1
-            self.processing_submission = False
-            
-            if self.current_round >= self.total_rounds:
-                self.is_game_active = False
+        if results.user_guess == results.correct_answer:
+            self.user_score += 1
+        if results.model_guess == results.correct_answer:
+            self.model_score += 1
+        
+        self.last_results = results
+        self.current_round += 1
+        self.processing_submission = False
+        
+        if self.current_round >= self.total_rounds:
+            self.is_game_active = False
     
     def get_current_image(self) -> Optional[str]:
         if not self.is_game_active or self.current_round >= len(self.game_images):
@@ -106,6 +100,40 @@ class GameState:
             </div>
             """
 
+def create_score_html(game_state: GameState):
+    results_html = ""
+    if game_state.last_results:
+        results_html = f"""
+        <div style='margin-top: 1rem; padding: 1rem; background-color: #e0e0e0; border-radius: 8px; color: #333;'>
+            <h4 style='color: #333; margin-bottom: 0.5rem;'>Last Round Results:</h4>
+            <p style='color: #333;'>Your guess: {game_state.last_results.user_guess}</p>
+            <p style='color: #333;'>Model's guess: {game_state.last_results.model_guess}</p>
+            <p style='color: #333;'>Correct answer: {game_state.last_results.correct_answer}</p>
+        </div>
+        """
+
+    current_display_round = min(game_state.current_round + 1, game_state.total_rounds)
+    
+    return f"""
+    <div style='padding: 1rem; background-color: #f0f0f0; border-radius: 8px; color: #333;'>
+        <h3 style='margin-bottom: 1rem; color: #333;'>Score Board</h3>
+        <div style='display: flex; justify-content: space-around;'>
+            <div>
+                <h4 style='color: #333;'>You</h4>
+                <p style='font-size: 1.5rem; color: #333;'>{game_state.user_score}</p>
+            </div>
+            <div>
+                <h4 style='color: #333;'>AI Model</h4>
+                <p style='font-size: 1.5rem; color: #333;'>{game_state.model_score}</p>
+            </div>
+        </div>
+        <div style='margin-top: 1rem;'>
+            <p style='color: #333;'>Round: {current_display_round}/{game_state.total_rounds}</p>
+        </div>
+        {results_html}
+    </div>
+    """
+
 def process_image(image):
     # Set Parameters
     top_k_patches = 5
@@ -159,8 +187,6 @@ def process_image(image):
 
     return sorted_probs
 
-game_state = GameState()
-
 def load_images() -> List[str]:
     real_image_folder = "real_images"
     fake_image_folder = "fake_images"
@@ -170,69 +196,37 @@ def load_images() -> List[str]:
     random.shuffle(selected_images)
     return selected_images
 
-def create_score_html() -> str:
-    results_html = ""
-    if game_state.last_results:
-        results_html = f"""
-        <div style='margin-top: 1rem; padding: 1rem; background-color: #e0e0e0; border-radius: 8px; color: #333;'>
-            <h4 style='color: #333; margin-bottom: 0.5rem;'>Last Round Results:</h4>
-            <p style='color: #333;'>Your guess: {game_state.last_results.user_guess}</p>
-            <p style='color: #333;'>Model's guess: {game_state.last_results.model_guess}</p>
-            <p style='color: #333;'>Correct answer: {game_state.last_results.correct_answer}</p>
-        </div>
-        """
-
-    current_display_round = min(game_state.current_round + 1, game_state.total_rounds)
+def start_game(state: Optional[GameState]):
+    # Initialize new game state if none exists
+    if state is None:
+        state = GameState()
     
-    return f"""
-    <div style='padding: 1rem; background-color: #f0f0f0; border-radius: 8px; color: #333;'>
-        <h3 style='margin-bottom: 1rem; color: #333;'>Score Board</h3>
-        <div style='display: flex; justify-content: space-around;'>
-            <div>
-                <h4 style='color: #333;'>You</h4>
-                <p style='font-size: 1.5rem; color: #333;'>{game_state.user_score}</p>
-            </div>
-            <div>
-                <h4 style='color: #333;'>AI Model</h4>
-                <p style='font-size: 1.5rem; color: #333;'>{game_state.model_score}</p>
-            </div>
-        </div>
-        <div style='margin-top: 1rem;'>
-            <p style='color: #333;'>Round: {current_display_round}/{game_state.total_rounds}</p>
-        </div>
-        {results_html}
-    </div>
-    """
-
-def start_game():
-    if not game_state.start_new_game():
-        return [gr.update()] * 6
+    if not state.start_new_game():
+        return [state] + [gr.update()] * 6
     
-    current_image = Image.open(game_state.get_current_image())
+    current_image = Image.open(state.get_current_image())
     
-    return (
+    return [
+        state,
         gr.update(value=current_image, visible=True),
         gr.update(visible=False),
         gr.update(visible=True, interactive=True),
         gr.update(visible=True, interactive=True),
-        create_score_html(),
-        gr.update(visible=False)
-    )
+        create_score_html(state),
+        gr.update(visible=False),
+    ]
 
-def submit_guess(user_guess: str):
-    # Early return if we can't submit a guess
-    if not game_state.can_submit_guess():
-        return [gr.update()] * 6
+def submit_guess(user_guess: str, state: GameState):
+    if not state.can_submit_guess():
+        return [state] + [gr.update()] * 6
     
-    # Mark submission as being processed
-    if not game_state.start_submission():
-        return [gr.update()] * 6
+    if not state.start_submission():
+        return [state] + [gr.update()] * 6
     
     try:
-        # Get current image and process it
-        current_image_path = game_state.get_current_image()
+        current_image_path = state.get_current_image()
         if not current_image_path:
-            return [gr.update()] * 6
+            return [state] + [gr.update()] * 6
         
         current_image = Image.open(current_image_path)
         model_prediction = process_image(current_image)
@@ -241,37 +235,38 @@ def submit_guess(user_guess: str):
         
         # Update game state with results
         results = GameResults(user_guess, model_guess, correct_answer)
-        game_state.finish_submission(results)
+        state.finish_submission(results)
         
         # Check if game is over
-        if not game_state.is_game_active:
-            return (
+        if not state.is_game_active:
+            return [
+                state,
                 gr.update(value=None, visible=False),
                 gr.update(visible=True),
                 gr.update(visible=False),
                 gr.update(visible=False),
-                create_score_html(),
-                gr.update(visible=True, value=game_state.get_game_over_message())
-            )
+                create_score_html(state),
+                gr.update(visible=True, value=state.get_game_over_message())
+            ]
         
-        # Get next image for the next round
-        next_image_path = game_state.get_current_image()
+        # Get next image
+        next_image_path = state.get_current_image()
         if not next_image_path:
-            return [gr.update()] * 6
-        
+            return [state] + [gr.update()] * 6
         next_image = Image.open(next_image_path)
         
-        return (
+        return [
+            state,
             gr.update(value=next_image, visible=True),
             gr.update(visible=False),
             gr.update(visible=True, interactive=True),
             gr.update(visible=True, interactive=True),
-            create_score_html(),
+            create_score_html(state),
             gr.update(visible=False)
-        )
+        ]
+
     except Exception as e:
-        # If any error occurs, reset the processing flag
-        game_state.processing_submission = False
+        state.processing_submission = False
         raise e
 
 # Custom CSS
@@ -303,6 +298,9 @@ custom_css = """
 
 # Define Gradio interface
 with gr.Blocks(css=custom_css) as iface:
+    # State variable for the game
+    state = gr.State(None)
+    
     with gr.Column(elem_id="game-container"):
         gr.HTML("""
             <table style="border-collapse: collapse; border: none; padding: 20px;">
@@ -359,7 +357,9 @@ with gr.Blocks(css=custom_css) as iface:
         # Event handlers
         start_button.click(
             fn=start_game,
+            inputs=[state],
             outputs=[
+                state,
                 image_display,
                 start_button,
                 real_button,
@@ -370,28 +370,33 @@ with gr.Blocks(css=custom_css) as iface:
         )
         
         real_button.click(
-            fn=lambda: submit_guess("Real"),
+            fn=lambda state: submit_guess("Real", state),
+            inputs=[state],
             outputs=[
+                state,
                 image_display,
                 start_button,
                 real_button,
                 fake_button,
                 score_display,
-                feedback_display
-            ]
+                feedback_display,
+            ],
         )
-        
+
         fake_button.click(
-            fn=lambda: submit_guess("Fake"),
+            fn=lambda state: submit_guess("Fake", state),
+            inputs=[state],
             outputs=[
+                state,
                 image_display,
                 start_button,
                 real_button,
                 fake_button,
                 score_display,
-                feedback_display
-            ]
+                feedback_display,
+            ],
         )
 
+
 # Launch the interface
 iface.launch()