Add Lock to handle User Guesses

app.py

@@ -1,11 +1,15 @@
+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
-import random
 from PIL import Image
 from skimage.feature import graycomatrix, graycoprops
 from torchvision import transforms
-import os
 
 NUM_ROUNDS = 10
 PROB_THRESHOLD = 0.3
@@ -13,6 +17,99 @@ PROB_THRESHOLD = 0.3
 # Load the model
 model = torch.jit.load("SuSy.pt")
 
+@dataclass
+class GameResults:
+    user_guess: str
+    model_guess: str
+    correct_answer: str
+
+class GameState:
+    def __init__(self):
+        self.lock = Lock()
+        self.reset()
+    
+    def reset(self):
+        with self.lock:
+            self.user_score = 0
+            self.model_score = 0
+            self.current_round = 0
+            self.total_rounds = NUM_ROUNDS
+            self.game_images: List[str] = []
+            self.is_game_active = False
+            self.last_results: Optional[GameResults] = None
+            self.processing_submission = False
+    
+    def start_new_game(self) -> bool:
+        with self.lock:
+            if self.is_game_active:
+                return False
+            self.reset()
+            self.game_images = load_images()
+            self.is_game_active = True
+            return True
+    
+    def can_submit_guess(self) -> bool:
+        with self.lock:
+            return (
+                self.is_game_active and 
+                not self.processing_submission and 
+                self.current_round < self.total_rounds
+            )
+    
+    def start_submission(self) -> bool:
+        with self.lock:
+            if not self.can_submit_guess():
+                return False
+            self.processing_submission = True
+            return True
+    
+    def finish_submission(self, results: GameResults):
+        with self.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
+    
+    def get_current_image(self) -> Optional[str]:
+        with self.lock:
+            if not self.is_game_active or self.current_round >= len(self.game_images):
+                return None
+            return self.game_images[self.current_round]
+    
+    def get_game_over_message(self) -> str:
+        with self.lock:
+            if self.user_score > self.model_score:
+                return """
+                <div style='text-align: center; margin-top: 20px; font-size: 1.2em;'>
+                    🎉 Congratulations! You won! 🎉<br>
+                    You've outperformed SuSy in detecting AI-generated images.<br>
+                    Click 'Start New Game' to play again.
+                </div>
+                """
+            elif self.user_score < self.model_score:
+                return """
+                <div style='text-align: center; margin-top: 20px; font-size: 1.2em;'>
+                    Better luck next time! SuSy won this round.<br>
+                    Keep practicing to improve your detection skills.<br>
+                    Click 'Start New Game' to try again.
+                </div>
+                """
+            else:
+                return """
+                <div style='text-align: center; margin-top: 20px; font-size: 1.2em;'>
+                    It's a tie! You matched SuSy's performance!<br>
+                    You're getting good at this.<br>
+                    Click 'Start New Game' to play again.
+                </div>
+                """
+
 def process_image(image):
     # Set Parameters
     top_k_patches = 5
@@ -66,47 +163,6 @@ def process_image(image):
 
     return sorted_probs
 
-
-class GameState:
-    def __init__(self):
-        self.user_score = 0
-        self.model_score = 0
-        self.current_round = 0
-        self.total_rounds = NUM_ROUNDS
-        self.game_images = []
-        self.is_game_active = False
-        self.last_results = None
-        self.waiting_for_input = True
-
-    def reset(self):
-        self.__init__()
-    
-    def get_game_over_message(self):
-        if self.user_score > self.model_score:
-            return """
-            <div style='text-align: center; margin-top: 20px; font-size: 1.2em;'>
-                🎉 Congratulations! You won! 🎉<br>
-                You've outperformed SuSy in detecting AI-generated images.<br>
-                Click 'Start New Game' to play again.
-            </div>
-            """
-        elif self.user_score < self.model_score:
-            return """
-            <div style='text-align: center; margin-top: 20px; font-size: 1.2em;'>
-                Better luck next time! SuSy won this round.<br>
-                Keep practicing to improve your detection skills.<br>
-                Click 'Start New Game' to try again.
-            </div>
-            """
-        else:
-            return """
-            <div style='text-align: center; margin-top: 20px; font-size: 1.2em;'>
-                It's a tie! You matched SuSy's performance!<br>
-                You're getting good at this.<br>
-                Click 'Start New Game' to play again.
-            </div>
-            """
-
 game_state = GameState()
 
 def load_images():
@@ -119,45 +175,45 @@ def load_images():
     return selected_images
 
 def create_score_html():
-    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>
-        """
+    with game_state.lock:
+        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>
+        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>
-                <h4 style='color: #333;'>AI Model</h4>
-                <p style='font-size: 1.5rem; color: #333;'>{game_state.model_score}</p>
+            <div style='margin-top: 1rem;'>
+                <p style='color: #333;'>Round: {current_display_round}/{game_state.total_rounds}</p>
             </div>
+            {results_html}
         </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():
-    game_state.reset()
-    game_state.game_images = load_images()
-    game_state.is_game_active = True
-    game_state.waiting_for_input = True
-    current_image = Image.open(game_state.game_images[0])
+    if not game_state.start_new_game():
+        return [gr.update()] * 6
+    
+    current_image = Image.open(game_state.get_current_image())
     
     return (
         gr.update(value=current_image, visible=True),
@@ -168,54 +224,60 @@ def start_game():
         gr.update(visible=False)
     )
 
-def submit_guess(user_guess):
-    if not game_state.is_game_active or not game_state.waiting_for_input:
+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
     
-    # Compute Model Guess
-    current_image = Image.open(game_state.game_images[game_state.current_round])
-    model_prediction = process_image(current_image)
-    model_guess = "Real" if model_prediction['Authentic'] > PROB_THRESHOLD else "Fake"
-    correct_answer = "Real" if "real_images" in game_state.game_images[game_state.current_round] else "Fake"
-    
-    # Update scores
-    if user_guess == correct_answer:
-        game_state.user_score += 1
-    if model_guess == correct_answer:
-        game_state.model_score += 1
-    
-    # Store last results for display
-    game_state.last_results = {
-        'user_guess': user_guess,
-        'model_guess': model_guess,
-        'correct_answer': correct_answer
-    }
-    
-    game_state.current_round += 1
-    game_state.waiting_for_input = True
+    # Mark submission as being processed
+    if not game_state.start_submission():
+        return [gr.update()] * 6
     
-    # Check if game is over
-    if game_state.current_round >= game_state.total_rounds:
-        game_state.is_game_active = False
+    try:
+        # Get current image and process it
+        current_image_path = game_state.get_current_image()
+        if not current_image_path:
+            return [gr.update()] * 6
+        
+        current_image = Image.open(current_image_path)
+        model_prediction = process_image(current_image)
+        model_guess = "Real" if model_prediction['Authentic'] > PROB_THRESHOLD else "Fake"
+        correct_answer = "Real" if "real_images" in current_image_path else "Fake"
+        
+        # Update game state with results
+        results = GameResults(user_guess, model_guess, correct_answer)
+        game_state.finish_submission(results)
+        
+        # Check if game is over
+        if not game_state.is_game_active:
+            return (
+                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())
+            )
+        
+        # Get next image for the next round
+        next_image_path = game_state.get_current_image()
+        if not next_image_path:
+            return [gr.update()] * 6
+        
+        next_image = Image.open(next_image_path)
+        
         return (
-            gr.update(value=None, visible=False),
-            gr.update(visible=True),
-            gr.update(visible=False),
+            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(),
-            gr.update(visible=True, value=game_state.get_game_over_message())
+            gr.update(visible=False)
         )
-    
-    next_image = Image.open(game_state.game_images[game_state.current_round])
-    
-    return (
-        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(),
-        gr.update(visible=False)
-    )
+    except Exception as e:
+        # If any error occurs, reset the processing flag
+        game_state.processing_submission = False
+        raise e
 
 # Custom CSS
 custom_css = """
@@ -238,7 +300,7 @@ custom_css = """
     min-width: 120px;
 }
 .image-container img {
-    max-height: 768px !important;
+    max-height: 640px !important;
     width: auto !important;
     object-fit: contain !important;
 }