Update app.py

app.py

@@ -4,26 +4,29 @@ import tempfile
 import os
 import base64
 from gtts import gTTS
-import speech_recognition as sr
-import torch
-from transformers import CLIPProcessor, CLIPModel
-from PIL import Image
+from ultralytics import YOLO
+from PIL import Image, ImageDraw, ImageColor
 import io
 import numpy as np
 import json
 from datetime import datetime
+import random
 
 # ============================================
-# 1. LOAD PDF (NO UPLOAD NEEDED)
+# 1. LOAD PDF AND YOLO MODEL
 # ============================================
 
-PDF_PATH = "2VBMAPP.pdf"  # Your pre-loaded PDF file
+PDF_PATH = "2VBMAPP.pdf"
 
-class ABATutor:  # FIXED: No space in class name
+class ABATutor:
     def __init__(self):
-        # Load CLIP for computer vision
-        self.clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
-        self.clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+        # Load YOLOv8 model
+        try:
+            self.yolo_model = YOLO("yolov8n.pt")
+            print("✅ YOLOv8 model loaded")
+        except Exception as e:
+            print(f"⚠️ YOLO error: {e}")
+            self.yolo_model = None
         
         # Load PDF
         try:
@@ -31,227 +34,234 @@ class ABATutor:  # FIXED: No space in class name
             self.total_pages = len(self.doc)
             print(f"✅ Loaded PDF with {self.total_pages} pages")
         except Exception as e:
-            print(f"⚠️ Could not load PDF: {e}")
+            print(f"⚠️ PDF error: {e}")
             self.doc = None
             self.total_pages = 1
         
-        self.current_page = 6  # Start at first Tact page (page 7 in PDF, index 6)
-        
-        # Session data
+        self.current_page = 6
         self.session_data = []
-        self.current_trial = {}
+        self.current_target = None  # Current object the AI is asking for
+        self.current_detections = []  # Current YOLO detections on this page
+        self.current_image = None  # Store current image for annotations
         
-        # Expected labels for each page (built from book content)
-        self.page_labels = self.build_label_map()
-        
-        print(f"✅ CLIP model ready")
-    
-    def build_label_map(self):
-        """Map page numbers to expected object labels based on VB-MAPP book"""
-        # These are the target objects from your book pages
-        return {
+        # Target objects for each page (from VB-MAPP book)
+        self.page_targets = {
             6: ["cat", "dog", "ball", "shoe", "car", "apple", "bird", "fish", "hat", "cup"],
             7: ["cat", "dog", "ball", "shoe", "car", "apple"],
             8: ["bird", "fish", "hat", "cup", "book", "chair"],
             9: ["train", "plane", "boat", "truck", "bus", "bike"],
             10: ["frog", "monkey", "elephant", "lion", "giraffe", "zebra"],
-            # Add more pages as needed
         }
     
-    def get_current_image(self):
-        """Extract current page as PIL Image for CV"""
+    def get_page_image_with_boxes(self, show_boxes=False):
+        """Extract current page and optionally draw YOLO detection boxes"""
         if self.doc is None:
-            # Return a blank image if PDF not loaded
             return Image.new('RGB', (800, 600), color='white')
         
         page = self.doc.load_page(self.current_page)
-        pix = page.get_pixmap(dpi=100)
+        pix = page.get_pixmap(dpi=120)
         img_data = pix.tobytes("png")
         pil_img = Image.open(io.BytesIO(img_data))
-        return pil_img
-    
-    def get_current_image_html(self):
-        """Get HTML image for display"""
-        if self.doc is None:
-            return '<p style="text-align:center;padding:50px;">⚠️ PDF not loaded. Please ensure 2VBMAPP.pdf is in the Space.</p>'
         
-        page = self.doc.load_page(self.current_page)
-        pix = page.get_pixmap(dpi=100)
-        img_data = pix.tobytes("png")
-        b64 = base64.b64encode(img_data).decode()
-        return f'<img src="data:image/png;base64,{b64}" style="max-width:100%; border-radius:10px; box-shadow:0 4px 6px rgba(0,0,0,0.1);" />'
-    
-    def cv_detect_object(self, spoken_word, image=None):
-        """
-        Use CLIP to check if spoken word matches objects in image
-        Returns confidence score and whether correct
-        """
-        if image is None:
-            image = self.get_current_image()
-        
-        # Get expected labels for this page
-        expected = self.page_labels.get(self.current_page, ["object"])
-        
-        # Build candidate labels (expected + spoken word)
-        candidates = expected + [spoken_word]
-        candidates = list(set(candidates))  # Remove duplicates
-        
-        try:
-            # Run CLIP
-            inputs = self.clip_processor(text=candidates, images=image, return_tensors="pt", padding=True)
-            with torch.no_grad():
-                outputs = self.clip_model(**inputs)
-                logits_per_image = outputs.logits_per_image
-                probs = logits_per_image.softmax(dim=1).numpy()
+        # Run YOLO to get detections
+        if self.yolo_model:
+            results = self.yolo_model(pil_img, conf=0.25, verbose=False)
+            self.current_detections = []
             
-            # Get best match
-            best_idx = probs[0].argmax()
-            best_label = candidates[best_idx]
-            confidence = float(probs[0][best_idx])
-        except Exception as e:
-            print(f"CLIP error: {e}")
-            best_label = spoken_word
-            confidence = 0.5
+            for r in results:
+                if r.boxes is not None:
+                    for box in r.boxes:
+                        class_id = int(box.cls[0])
+                        class_name = self.yolo_model.names[class_id].lower()
+                        confidence = float(box.conf[0])
+                        x1, y1, x2, y2 = box.xyxy[0].tolist()
+                        
+                        self.current_detections.append({
+                            "class": class_name,
+                            "confidence": confidence,
+                            "bbox": [int(x1), int(y1), int(x2), int(y2)],
+                            "center": [(int(x1)+int(x2))//2, (int(y1)+int(y2))//2]
+                        })
+            
+            # Draw boxes if requested
+            if show_boxes:
+                draw = ImageDraw.Draw(pil_img)
+                for det in self.current_detections:
+                    x1, y1, x2, y2 = det["bbox"]
+                    draw.rectangle([x1, y1, x2, y2], outline="red", width=3)
+                    draw.text((x1, y1-15), f"{det['class']} ({det['confidence']:.0%})", fill="red")
         
-        # Check if spoken word matches any expected label
-        spoken_lower = spoken_word.lower().strip()
-        is_correct = any(spoken_lower == exp.lower() for exp in expected)
+        return pil_img
+    
+    def get_current_image_html(self, show_boxes=False):
+        """Get HTML image with clickable areas"""
+        pil_img = self.get_page_image_with_boxes(show_boxes=show_boxes)
         
-        # Also check CLIP confidence for bonus
-        clip_confirms = best_label.lower() == spoken_lower and confidence > 0.3
+        # Convert to base64
+        buffered = io.BytesIO()
+        pil_img.save(buffered, format="PNG")
+        img_base64 = base64.b64encode(buffered.getvalue()).decode()
         
-        return {
-            "is_correct": is_correct or clip_confirms,
-            "confidence": confidence,
-            "detected_label": best_label,
-            "expected_labels": expected,
-            "spoken_word": spoken_word
-        }
+        # If we have detections, create an image map for clickable areas
+        if self.current_detections:
+            # Build clickable areas HTML
+            areas_html = ""
+            for i, det in enumerate(self.current_detections):
+                x1, y1, x2, y2 = det["bbox"]
+                areas_html += f'''
+                <area shape="rect" coords="{x1},{y1},{x2},{y2}" 
+                      data-class="{det['class']}" 
+                      data-index="{i}"
+                      onclick="selectObject(this)"
+                      style="cursor:pointer;"
+                      title="Click on {det['class']}">
+                '''
+            
+            # JavaScript for click handling
+            js_script = """
+            <script>
+            function selectObject(area) {
+                let className = area.getAttribute('data-class');
+                let index = area.getAttribute('data-index');
+                
+                // Send to Gradio
+                const event = new CustomEvent('gradio_click', {
+                    detail: { class: className, index: index }
+                });
+                window.dispatchEvent(event);
+                
+                // Also update a hidden input
+                let hiddenInput = document.getElementById('selected_object');
+                if (hiddenInput) {
+                    hiddenInput.value = className;
+                    hiddenInput.dispatchEvent(new Event('change'));
+                }
+                
+                // Visual feedback
+                area.style.outline = '3px solid green';
+                setTimeout(() => { area.style.outline = ''; }, 500);
+            }
+            </script>
+            """
+            
+            html = f'''
+            <div style="position: relative;">
+                <img src="data:image/png;base64,{img_base64}" 
+                     usemap="#objectmap" 
+                     style="max-width:100%; border-radius:10px; cursor:pointer;" />
+                <map name="objectmap" id="objectmap">
+                    {areas_html}
+                </map>
+                <input type="hidden" id="selected_object" value="">
+                {js_script}
+            </div>
+            <p style="font-size:12px; color:#666; margin-top:5px;">
+            💡 Click directly on any object in the image above to answer!
+            </p>
+            '''
+            return html
+        else:
+            # No detections found
+            return f'''
+            <img src="data:image/png;base64,{img_base64}" style="max-width:100%; border-radius:10px;" />
+            <p style="font-size:12px; color:#999; margin-top:5px;">
+            ⚠️ No objects detected on this page. Try a different page with clear pictures.
+            </p>
+            '''
+    
+    def get_clickable_image(self):
+        """Return clickable image for Gradio"""
+        return self.get_current_image_html(show_boxes=True)
     
     def speak(self, text):
-        """Convert text to speech and return audio file"""
+        """Convert text to speech"""
         try:
             tts = gTTS(text=text, lang="en", slow=False)
             audio_path = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3").name
             tts.save(audio_path)
             return audio_path
-        except Exception as e:
-            print(f"TTS error: {e}")
+        except:
             return None
     
-    def get_sd_prompt(self):
-        """Get the SD prompt for current page based on book content"""
-        # Pages 6-37 are Tact: "What's that?"
-        if 6 <= self.current_page <= 37:
-            return "What's that?"
-        # Pages 52-82 are Listener Responding: "Where is the ___?"
-        elif 52 <= self.current_page <= 82:
-            return "Where is the dog?"
-        # Pages 84-108 are VP-MTS: "Match"
-        elif 84 <= self.current_page <= 108:
-            return "Match the picture"
-        else:
-            return "What is this?"
-    
-    def listen(self):
-        """Capture and transcribe child's spoken response"""
-        recognizer = sr.Recognizer()
-        try:
-            with sr.Microphone() as source:
-                recognizer.adjust_for_ambient_noise(source, duration=0.5)
-                audio = recognizer.listen(source, timeout=5, phrase_time_limit=3)
-                text = recognizer.recognize_google(audio)
-                return text.lower().strip()
-        except sr.WaitTimeoutError:
-            return "[no response]"
-        except sr.UnknownValueError:
-            return "[could not understand]"
-        except Exception as e:
-            return f"[error: {str(e)}]"
-    
-    def run_trial(self):
-        """Run one complete assessment trial with CV validation"""
-        # Get SD prompt
-        prompt = self.get_sd_prompt()
-        prompt_audio = self.speak(prompt)
+    def ask_new_question(self):
+        """Generate a new question based on current page detections"""
+        if not self.current_detections:
+            return None, "No objects detected on this page. Please go to a page with clear pictures."
         
-        # Get current image
-        current_image = self.get_current_image()
+        # Filter detections that are in our target list for this page
+        valid_targets = self.page_targets.get(self.current_page, [])
+        eligible_detections = [d for d in self.current_detections if d["class"] in valid_targets]
         
-        # Listen to child
-        child_response = self.listen()
+        if not eligible_detections:
+            # Use any detection if none match the page targets
+            eligible_detections = self.current_detections
         
-        # Run CV detection
-        cv_result = self.cv_detect_object(child_response, current_image)
+        # Pick a random object to ask about
+        target = random.choice(eligible_detections)
+        self.current_target = target
         
-        # Generate feedback
-        if cv_result["is_correct"]:
-            feedback = f"✅ Great job! '{child_response}' is correct!"
+        # Create the question
+        question = f"Where is the {target['class']}? Click on it!"
+        question_audio = self.speak(question)
+        
+        return question, question_audio
+    
+    def check_answer(self, clicked_class):
+        """Check if the clicked object matches what was asked"""
+        if not self.current_target:
+            return "⚠️ Please click 'Ask Question' first!", False, 0
+        
+        is_correct = (clicked_class == self.current_target["class"])
+        
+        if is_correct:
+            feedback = f"✅ Correct! That IS the {self.current_target['class']}! Great job!"
             score = 1
-        elif child_response == "[no response]":
-            feedback = "⏸️ I didn't hear anything. Let's try again. " + prompt
-            score = 0
-        elif child_response == "[could not understand]":
-            feedback = "🎤 I couldn't understand. Can you say it clearly?"
-            score = 0
         else:
-            expected_show = cv_result["expected_labels"][0] if cv_result["expected_labels"] else "the picture"
-            feedback = f"❌ Not quite. I heard '{child_response}'. The answer is {expected_show}. Let's keep practicing!"
+            # Find where the target actually is
+            target_pos = ""
+            for det in self.current_detections:
+                if det["class"] == self.current_target["class"]:
+                    target_pos = f" (Look for the {det['class']} in the picture)"
+                    break
+            feedback = f"❌ Not quite. You clicked on {clicked_class}, but I asked for the {self.current_target['class']}.{target_pos} Try again!"
             score = 0
         
-        feedback_audio = self.speak(feedback)
-        
-        # Log trial
+        # Log the trial
         trial = {
             "timestamp": datetime.now().isoformat(),
             "page": self.current_page,
-            "prompt": prompt,
-            "child_response": child_response,
-            "cv_detected": cv_result["detected_label"],
-            "cv_confidence": cv_result["confidence"],
-            "expected": cv_result["expected_labels"],
-            "correct": cv_result["is_correct"],
-            "score": score
+            "question": f"Where is the {self.current_target['class']}?",
+            "asked_for": self.current_target["class"],
+            "clicked_on": clicked_class,
+            "correct": is_correct,
+            "score": score,
+            "detection_confidence": self.current_target["confidence"]
         }
         self.session_data.append(trial)
         
-        # Calculate totals
-        total_score = sum(t["score"] for t in self.session_data)
-        trials_count = len(self.session_data)
-        
-        return {
-            "prompt": prompt,
-            "prompt_audio": prompt_audio,
-            "child_response": child_response,
-            "cv_detected": cv_result["detected_label"],
-            "cv_confidence": f"{cv_result['confidence']:.2%}",
-            "feedback": feedback,
-            "feedback_audio": feedback_audio,
-            "score": score,
-            "total_score": total_score,
-            "trials_count": trials_count,
-            "page_html": self.get_current_image_html()
-        }
+        return feedback, score, is_correct
     
     def next_page(self):
-        """Move to next assessment page"""
+        """Move to next page and reset state"""
         if self.current_page < self.total_pages - 1:
             self.current_page += 1
-        return self.get_current_image_html()
+        self.current_target = None
+        return self.get_clickable_image(), f"Page {self.current_page + 1} / {self.total_pages}"
     
     def prev_page(self):
         """Move to previous page"""
         if self.current_page > 0:
             self.current_page -= 1
-        return self.get_current_image_html()
+        self.current_target = None
+        return self.get_clickable_image(), f"Page {self.current_page + 1} / {self.total_pages}"
     
     def get_session_report(self):
-        """Generate session summary JSON"""
+        """Generate session summary"""
         correct_count = sum(1 for t in self.session_data if t["correct"])
         total = len(self.session_data)
         accuracy = correct_count / total if total > 0 else 0
         
-        report = {
+        return {
             "session_summary": {
                 "total_trials": total,
                 "correct_responses": correct_count,
@@ -261,132 +271,168 @@ class ABATutor:  # FIXED: No space in class name
             "trials": self.session_data,
             "export_date": datetime.now().isoformat()
         }
-        return report
 
 # ============================================
 # 2. INITIALIZE TUTOR
 # ============================================
 
-tutor = ABATutor()  # FIXED: Class name without space
+tutor = ABATutor()
+selected_class_state = gr.State("")
 
 # ============================================
 # 3. GRADIO INTERFACE
 # ============================================
 
 custom_css = """
-.green-bg { background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); }
 .main-container { max-width: 1400px; margin: auto; }
 .page-card { background: white; border-radius: 20px; padding: 20px; box-shadow: 0 10px 40px rgba(0,0,0,0.1); }
-.score-badge { font-size: 36px; font-weight: bold; background: #4CAF50; color: white; border-radius: 50%; width: 80px; height: 80px; display: flex; align-items: center; justify-content: center; }
 .feedback-box { background: #f0f4ff; border-radius: 15px; padding: 15px; margin: 10px 0; }
+.question-card { background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); color: white; border-radius: 15px; padding: 20px; }
+.click-hint { background: #e8f5e9; padding: 10px; border-radius: 10px; text-align: center; font-size: 14px; }
 """
 
 with gr.Blocks(css=custom_css, title="ABA-AI-Tutor", theme=gr.themes.Soft()) as demo:
     gr.Markdown("""
-    # 🗣️🤖 ABA-AI-Tutor
+    # 🖱️🤖 ABA-AI-Tutor
     
-    ### AI-Powered VB-MAPP Level 2 Assessment
-    *The AI shows a page, asks the question, listens to your answer, and uses computer vision to check if you're correct.*
+    ### Click-to-Select VB-MAPP Assessment
+    *The AI asks "Where is the ___?" You **click on the object** in the picture to answer!*
     """)
     
     with gr.Row(elem_classes=["main-container"]):
-        # LEFT COLUMN: PDF Viewer
+        # LEFT COLUMN: Clickable PDF Viewer
         with gr.Column(scale=3):
             with gr.Group(elem_classes=["page-card"]):
-                pdf_display = gr.HTML(value=tutor.get_current_image_html(), label="Current Page")
+                pdf_display = gr.HTML(value=tutor.get_clickable_image())
                 with gr.Row():
                     prev_btn = gr.Button("◀ Previous Page", size="sm", variant="secondary")
-                    page_info = gr.Textbox(value=f"Page {tutor.current_page + 1} / {tutor.total_pages}", interactive=False, label="Page", container=False)
+                    page_info = gr.Textbox(value=f"Page {tutor.current_page + 1} / {tutor.total_pages}", interactive=False, container=False)
                     next_btn = gr.Button("Next Page ▶", size="sm", variant="secondary")
         
         # RIGHT COLUMN: Assessment Controls
         with gr.Column(scale=2):
-            gr.Markdown("### 🎯 Assessment Session")
+            gr.Markdown("### 🎯 Click Assessment")
+            
+            # Hidden input to capture clicks (workaround for Gradio)
+            clicked_class = gr.Textbox(visible=False, elem_id="selected_object", label="Clicked Object")
             
-            assess_btn = gr.Button("🎙️ Ask Question & Listen", variant="primary", size="lg")
+            with gr.Group(elem_classes=["question-card"]):
+                question_display = gr.Textbox(label="📢 AI Question", interactive=False, placeholder="Click 'Ask Question' to start...")
+                question_audio = gr.Audio(label="🔊 Question Audio", type="filepath", interactive=False)
+            
+            with gr.Row():
+                ask_btn = gr.Button("🎲 Ask New Question", variant="primary", size="lg")
+                check_btn = gr.Button("✅ Check My Click", variant="secondary", size="lg")
             
             with gr.Group(elem_classes=["feedback-box"]):
-                prompt_display = gr.Textbox(label="📢 Question Asked", interactive=False)
-                prompt_audio = gr.Audio(label="🔊 Question Audio", type="filepath", interactive=False)
-                
-                child_response = gr.Textbox(label="🎤 Child's Response", interactive=False, placeholder="The AI will transcribe your answer...")
-                cv_detected = gr.Textbox(label="🤖 CV Detected", interactive=False, placeholder="Computer vision sees...")
-                cv_confidence = gr.Textbox(label="📊 CV Confidence", interactive=False)
-                
                 feedback_msg = gr.Textbox(label="💬 AI Feedback", interactive=False)
                 feedback_audio = gr.Audio(label="🔊 Feedback Audio", type="filepath", interactive=False)
             
             with gr.Row():
-                trial_score = gr.Number(label="⭐ This Trial Score", interactive=False, value=0)
+                trial_score = gr.Number(label="⭐ Last Score", interactive=False, value=0)
                 total_score_display = gr.Number(label="🏆 Total Score", interactive=False, value=0)
                 trials_count = gr.Number(label="📋 Trials Completed", interactive=False, value=0)
+            
+            gr.Markdown("""
+            <div class="click-hint">
+            🖱️ <strong>How to play:</strong><br>
+            1. Click "Ask New Question"<br>
+            2. Listen to what the AI asks for<br>
+            3. <strong>Click directly on that object</strong> in the picture<br>
+            4. Click "Check My Click" to see if you're right!
+            </div>
+            """)
     
     # Session Report Tab
     with gr.Tabs():
         with gr.TabItem("📊 Session Report"):
             refresh_btn = gr.Button("Refresh Report")
             report_json = gr.JSON(label="Full Assessment Data", value=tutor.get_session_report())
-            gr.Markdown("### 🎓 Summary for Clinician")
             accuracy_display = gr.Markdown("**Accuracy:** 0%")
-        
-        with gr.TabItem("⚙️ Manual Navigation"):
-            gr.Markdown("## Jump to Specific Page")
-            page_slider = gr.Slider(minimum=0, maximum=tutor.total_pages-1, step=1, label="Page Number", value=tutor.current_page)
-            jump_btn = gr.Button("Go to Page")
-            page_preview = gr.HTML()
-            
-            def jump_to_page(page_num):
-                tutor.current_page = int(page_num)
-                return tutor.get_current_image_html()
-            
-            jump_btn.click(fn=jump_to_page, inputs=page_slider, outputs=page_preview)
     
     # ============================================
     # 4. EVENT HANDLERS
     # ============================================
     
-    def run_and_update():
-        result = tutor.run_trial()
-        return (
-            result["prompt"], result["prompt_audio"],
-            result["child_response"], result["cv_detected"], result["cv_confidence"],
-            result["feedback"], result["feedback_audio"],
-            result["score"], result["total_score"], result["trials_count"],
-            result["page_html"]
-        )
+    # Ask new question
+    def ask_question():
+        question, audio = tutor.ask_new_question()
+        if question:
+            return question, audio, "Waiting for your click...", None, 0, 0, 0
+        else:
+            return "No objects detected on this page. Try a different page.", None, "⚠️ No clickable objects found", None, 0, 0, 0
+    
+    ask_btn.click(
+        fn=ask_question,
+        outputs=[question_display, question_audio, feedback_msg, feedback_audio, trial_score, total_score_display, trials_count]
+    )
+    
+    # Check answer when child clicks
+    def process_click_and_check(clicked, current_question):
+        if not clicked or clicked == "":
+            return "Click on an object in the picture first!", None, 0, 0, 0
+        
+        feedback, score, is_correct = tutor.check_answer(clicked.lower())
+        feedback_audio = tutor.speak(feedback)
+        
+        total_score = sum(t["score"] for t in tutor.session_data)
+        trials = len(tutor.session_data)
+        
+        return feedback, feedback_audio, score, total_score, trials
     
-    assess_btn.click(
-        fn=run_and_update,
-        outputs=[
-            prompt_display, prompt_audio,
-            child_response, cv_detected, cv_confidence,
-            feedback_msg, feedback_audio,
-            trial_score, total_score_display, trials_count,
-            pdf_display
-        ]
+    check_btn.click(
+        fn=process_click_and_check,
+        inputs=[clicked_class, question_display],
+        outputs=[feedback_msg, feedback_audio, trial_score, total_score_display, trials_count]
     )
     
+    # Navigation
     def prev_page_update():
-        html = tutor.prev_page()
-        return html, f"Page {tutor.current_page + 1} / {tutor.total_pages}"
+        html, info = tutor.prev_page()
+        return html, info, "", None, "", None, 0, 0, 0
     
     def next_page_update():
-        html = tutor.next_page()
-        return html, f"Page {tutor.current_page + 1} / {tutor.total_pages}"
+        html, info = tutor.next_page()
+        return html, info, "", None, "", None, 0, 0, 0
     
-    prev_btn.click(fn=prev_page_update, outputs=[pdf_display, page_info])
-    next_btn.click(fn=next_page_update, outputs=[pdf_display, page_info])
+    prev_btn.click(
+        fn=prev_page_update,
+        outputs=[pdf_display, page_info, question_display, question_audio, feedback_msg, feedback_audio, trial_score, total_score_display, trials_count]
+    )
+    
+    next_btn.click(
+        fn=next_page_update,
+        outputs=[pdf_display, page_info, question_display, question_audio, feedback_msg, feedback_audio, trial_score, total_score_display, trials_count]
+    )
     
+    # Refresh report
     def refresh_report():
         report = tutor.get_session_report()
         accuracy = report["session_summary"]["accuracy"]
         return report, f"**Accuracy:** {accuracy}"
     
     refresh_btn.click(fn=refresh_report, outputs=[report_json, accuracy_display])
-
-# ============================================
-# 5. LAUNCH
-# ============================================
+    
+    # JavaScript to capture clicks and send to Gradio
+    demo.load(_js="""
+    function captureClicks() {
+        window.addEventListener('gradio_click', function(e) {
+            let hiddenInput = document.getElementById('selected_object');
+            if (hiddenInput) {
+                hiddenInput.value = e.detail.class;
+                hiddenInput.dispatchEvent(new Event('change', { bubbles: true }));
+                
+                // Also try to update Gradio's state
+                const checkBtn = document.querySelector('button[aria-label="Check My Click"]');
+                if (checkBtn) {
+                    checkBtn.style.animation = 'pulse 0.5s';
+                    setTimeout(() => { checkBtn.style.animation = ''; }, 500);
+                }
+            }
+        });
+    }
+    captureClicks();
+    """)
 
 if __name__ == "__main__":
     demo.launch()