Update app.py

app.py

@@ -1,15 +1,17 @@
 """
-Simplified Dog Tracking for Training  Dataset Collection
+Simplified Dog Tracking for Training Dataset Collection
 - Process video with adjustable threshold
 - Temporary storage with discard option
-- Manual validation (grid of thumbnails with per-image checkboxes)
+- Manual validation with checkbox selection per image
 - Export to folder structure for fine-tuning
+- Download to laptop as ZIP
 - Automatic HuggingFace backup/restore
 """
 import os
 os.environ["OMP_NUM_THREADS"] = "1"
 
 import zipfile
+import tempfile
 import gradio as gr
 import cv2
 import numpy as np
@@ -53,8 +55,8 @@ class DatasetCollectionApp:
         self.current_video_path = None
         self.is_processing = False
 
-        # Validation state: list of (temp_id, [checkbox components])
-        self.validation_checkboxes: List[tuple] = []
+        # Validation state: stores checkbox states for each temp_id
+        self.validation_data = {}  # {temp_id: [bool, bool, ...]}
 
         print("Dataset Collection App initialized")
         print(f"Database has {len(self.db.get_all_dogs())} dogs")
@@ -73,7 +75,7 @@ class DatasetCollectionApp:
         self.tracker.reset()
         self.reid.reset_session()
         self.current_video_path = None
-        self.validation_checkboxes = []
+        self.validation_data = {}
 
         gc.collect()
         if torch.cuda.is_available():
@@ -83,7 +85,8 @@ class DatasetCollectionApp:
             None,  # Clear video input
             "<p style='text-align:center; color:#868e96;'>Session cleared. Upload a new video to start.</p>",
             "",
-            ""
+            "",
+            gr.update(visible=False)  # Hide validation area
         )
 
     def discard_session(self):
@@ -92,17 +95,16 @@ class DatasetCollectionApp:
         self.temp_session.clear()
         self.tracker.reset()
         self.reid.reset_session()
-        self.validation_checkboxes = []
+        self.validation_data = {}
 
         gc.collect()
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
 
-        # Return UI updates for validation container + status + database display
         return (
-            gr.update(visible=False),  # hide validation container
+            gr.update(visible=False),  # Hide validation container
             f"Discarded {count} temporary dogs. Try different threshold.",
-            gr.update(visible=False)   # hide database display
+            gr.update(visible=False)   # Hide database display
         )
 
     def process_video(self, video_path: str, reid_threshold: float,
@@ -110,11 +112,12 @@ class DatasetCollectionApp:
         """Process video and store in temporary session"""
 
         if not video_path:
-            return None, "Please upload a video", ""
+            return None, "Please upload a video", "", gr.update(visible=False)
 
         self.is_processing = True
         self.current_video_path = video_path
         self.temp_session.clear()
+        self.validation_data = {}
 
         # Set threshold
         self.reid.set_threshold(reid_threshold)
@@ -127,7 +130,7 @@ class DatasetCollectionApp:
         try:
             cap = cv2.VideoCapture(video_path)
             if not cap.isOpened():
-                return None, "Cannot open video", ""
+                return None, "Cannot open video", "", gr.update(visible=False)
 
             total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
             fps = cap.get(cv2.CAP_PROP_FPS) or 30
@@ -248,6 +251,10 @@ class DatasetCollectionApp:
             # Store in temp session
             self.temp_session = temp_dogs
 
+            # Initialize validation data (all images selected by default)
+            for temp_id in temp_dogs.keys():
+                self.validation_data[temp_id] = [True] * len(temp_dogs[temp_id]['images'])
+
             # Generate summary
             summary = f"Processing complete!\n"
             summary += f"Detected {original_count} dogs initially\n"
@@ -261,9 +268,11 @@ class DatasetCollectionApp:
             if len(temp_dogs) == 0:
                 summary += "No dogs met the minimum requirement of 14 images.\n"
                 summary += "Try adjusting the ReID threshold or using a longer video."
+                show_validation = False
             else:
                 summary += "Results stored in TEMPORARY session\n"
-                summary += "Review and validate before saving to database"
+                summary += "Go to Tab 2 to review and select images before saving"
+                show_validation = True
 
             gallery_html = self._create_temp_gallery()
 
@@ -271,12 +280,17 @@ class DatasetCollectionApp:
             if torch.cuda.is_available():
                 torch.cuda.empty_cache()
 
-            return gallery_html, summary, "Ready for validation" if len(temp_dogs) > 0 else "No valid dogs"
+            return (
+                gallery_html, 
+                summary, 
+                "Ready for validation" if len(temp_dogs) > 0 else "No valid dogs",
+                gr.update(visible=show_validation)
+            )
 
         except Exception as e:
             import traceback
             error = f"Error: {str(e)}\n{traceback.format_exc()}"
-            return None, error, ""
+            return None, error, "", gr.update(visible=False)
         finally:
             self.is_processing = False
 
@@ -322,56 +336,442 @@ class DatasetCollectionApp:
         html += "</div></div>"
         return html
 
-    def create_validation_interface(self):
-        """Create simplified validation interface for visual review (fallback)
-           Note: This function returns HTML used by the simple load button.
-           We keep it for backward compatibility; the interactive grid is built
-           by render_validation in the Gradio Blocks UI below."""
+    def load_validation_interface(self):
+        """Load validation interface with checkbox selection"""
         if not self.temp_session:
-            return "<p>No temporary session to validate</p>"
-
+            return (
+                gr.update(visible=False),
+                "No temporary session to validate. Process a video first.",
+                ""
+            )
+
+        # Create components list for dynamic rendering
+        validation_components = []
+        
         html = "<div style='padding: 20px;'>"
-        html += "<h2 style='text-align:center;'>Validation - Visual Review</h2>"
-        html += "<p style='text-align:center; color:#666;'>Review images before saving. All images will be saved when you click 'Save to Database'</p>"
+        html += "<h2 style='text-align:center;'>Review and Select Images</h2>"
+        html += "<p style='text-align:center; color:#666;'>Check/uncheck images to keep/discard. All are selected by default.</p>"
+        html += "</div>"
 
-        for temp_id in sorted(self.temp_session.keys()):
-            dog_data = self.temp_session[temp_id]
-            images = dog_data['images']
+        status = f"Loaded {len(self.temp_session)} dogs for validation. Review and click 'Save Selected to Database' when ready."
+
+        return (
+            gr.update(visible=True),
+            status,
+            html
+        )
+
+    def save_validated_to_database(self, *checkbox_states):
+        """Save validated images to permanent database"""
+        if not self.temp_session:
+            return "No temporary session to save", gr.update()
+
+        try:
+            saved_count = 0
+            total_images_saved = 0
+            
+            # Collect checkbox states
+            checkbox_idx = 0
+            
+            for temp_id in sorted(self.temp_session.keys()):
+                dog_data = self.temp_session[temp_id]
+                num_images = len(dog_data['images'])
+                
+                # Get checkbox states for this dog
+                selected_indices = []
+                for i in range(num_images):
+                    if checkbox_idx < len(checkbox_states) and checkbox_states[checkbox_idx]:
+                        selected_indices.append(i)
+                    checkbox_idx += 1
+                
+                # Skip if no images selected
+                if not selected_indices:
+                    continue
+                
+                # Add dog to database
+                dog_id = self.db.add_dog(
+                    name=f"Dog_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{temp_id}"
+                )
+                
+                # Add only selected images
+                for idx in selected_indices:
+                    self.db.add_dog_image(
+                        dog_id=dog_id,
+                        image=dog_data['images'][idx],
+                        timestamp=dog_data['timestamps'][idx],
+                        confidence=dog_data['confidences'][idx],
+                        bbox=dog_data['bboxes'][idx]
+                    )
+                    total_images_saved += 1
+                
+                saved_count += 1
+            
+            # Clear temporary session after saving
+            self.temp_session.clear()
+            self.validation_data = {}
+            
+            # Backup to HuggingFace
+            self._backup_database()
+            
+            # Show updated database
+            db_html = self._show_database()
+            
+            summary = f"✅ Successfully saved {saved_count} dogs with {total_images_saved} selected images to permanent database!"
+            
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            
+            return summary, gr.update(value=db_html, visible=True)
+            
+        except Exception as e:
+            import traceback
+            error = f"Error saving: {str(e)}\n{traceback.format_exc()}"
+            return error, gr.update()
+
+    def _backup_database(self):
+        """Backup database to HuggingFace"""
+        try:
+            from huggingface_hub import HfApi
+            
+            hf_token = os.getenv('HF_TOKEN')
+            if not hf_token:
+                print("Warning: HF_TOKEN not found, skipping backup")
+                return
+            
+            api = HfApi()
+            repo_id = "mustafa2ak/dog-dataset-backup"
+            
+            # Upload database file
+            api.upload_file(
+                path_or_fileobj='dog_monitoring.db',
+                path_in_repo='dog_monitoring.db',
+                repo_id=repo_id,
+                repo_type='dataset',
+                token=hf_token
+            )
+            
+            print(f"✅ Database backed up to {repo_id}")
+            
+        except Exception as e:
+            print(f"Backup failed: {str(e)}")
 
+    def _restore_database(self):
+        """Restore database from HuggingFace"""
+        try:
+            from huggingface_hub import hf_hub_download
+            
+            hf_token = os.getenv('HF_TOKEN')
+            if not hf_token:
+                print("No HF_TOKEN found, starting with fresh database")
+                return
+            
+            repo_id = "mustafa2ak/dog-dataset-backup"
+            
+            # Download database
+            db_path = hf_hub_download(
+                repo_id=repo_id,
+                filename='dog_monitoring.db',
+                repo_type='dataset',
+                token=hf_token
+            )
+            
+            # Copy to current directory
+            import shutil
+            shutil.copy(db_path, 'dog_monitoring.db')
+            
+            print(f"✅ Database restored from {repo_id}")
+            
+        except Exception as e:
+            print(f"No backup found or restore failed: {str(e)}")
+
+    def _show_database(self) -> str:
+        """Show current database contents"""
+        dogs = self.db.get_all_dogs()
+        
+        if not dogs:
+            return "<p style='text-align:center; color:#868e96;'>No dogs in database yet</p>"
+        
+        html = "<div style='padding: 20px;'>"
+        html += f"<h2 style='text-align:center; color:#228be6;'>Permanent Database ({len(dogs)} dogs)</h2>"
+        html += "<div style='display: grid; grid-template-columns: repeat(auto-fit, minmax(300px, 1fr)); gap: 20px;'>"
+        
+        for _, dog in dogs.iterrows():
+            images = self.db.get_dog_images(dog['dog_id'])
+            display_count = min(6, len(images))
+            
             html += f"""
-            <div style='border: 2px solid #495057; border-radius: 10px; 
-                        padding: 15px; margin: 20px 0; background: #f8f9fa;'>
-                <h3 style='margin: 0 0 15px 0;'>Temp Dog #{temp_id} - {len(images)} images</h3>
-                <div style='display: grid; grid-template-columns: repeat(6, 1fr); gap: 10px;'>
+            <div style='border: 2px solid #228be6; border-radius: 10px; 
+                        padding: 15px; background: #e7f5ff;'>
+                <h3 style='margin: 0 0 10px 0; color:#1971c2;'>{dog['name']}</h3>
+                <p style='color: #666; margin: 5px 0;'>ID: {dog['dog_id']}</p>
+                <p style='color: #666; margin: 5px 0;'>Images: {len(images)}</p>
+                <p style='color: #666; margin: 5px 0; font-size: 12px;'>
+                    First seen: {dog['first_seen']}
+                </p>
+                <div style='display: grid; grid-template-columns: repeat(3, 1fr); gap: 5px; margin-top: 10px;'>
             """
-
-            for idx, img in enumerate(images):
+            
+            for img_data in images[:display_count]:
+                img = img_data['image']
                 img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                 img_base64 = self._img_to_base64(img_rgb)
-
                 html += f"""
-                <div style='position: relative;'>
-                    <img src='data:image/jpeg;base64,{img_base64}'
-                         style='width: 100%; aspect-ratio: 1; object-fit: cover; 
-                                border-radius: 5px; border: 2px solid #dee2e6;'>
-                    <div style='position: absolute; bottom: 5px; right: 5px; 
-                                background: rgba(0,0,0,0.7); color: white; 
-                                padding: 2px 6px; border-radius: 3px; font-size: 10px;'>
-                        {idx+1}
-                    </div>
-                </div>
+                <img src='data:image/jpeg;base64,{img_base64}'
+                     style='width: 100%; aspect-ratio: 1; object-fit: cover; 
+                            border-radius: 5px;'>
                 """
-
-            html += """
-                </div>
-            </div>
-            """
-
-        html += "<p style='text-align:center; color:#868e96; margin-top: 30px;'>"
-        html += "If results look good, click 'Save to Database' below.<br>"
-        html += "If not satisfied, go back to Tab 1 and click 'Discard & Retry' with different threshold."
-        html += "</p>"
-        html += "</div>"
+            
+            html += "</div></div>"
+        
+        html += "</div></div>"
         return html
 
-    # (rest of code unchanged, includes save_validated_to_database, _backup_database, _restore_database, _show_database, export_dataset, _img_to_base64, create_interface, launch, __main__)
+    def export_dataset(self):
+        """Export dataset as downloadable ZIP file"""
+        try:
+            dogs = self.db.get_all_dogs()
+            
+            if dogs.empty:
+                return "No dogs in database to export", None
+            
+            # Create in-memory ZIP file
+            zip_buffer = BytesIO()
+            
+            with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zipf:
+                total_images = 0
+                export_info = []
+                
+                for _, dog in dogs.iterrows():
+                    dog_id = dog['dog_id']
+                    dog_name = dog['name'] or f"dog_{dog_id}"
+                    safe_name = "".join(c if c.isalnum() or c in ('_', '-') else '_' for c in dog_name)
+                    
+                    images = self.db.get_dog_images(dog_id)
+                    
+                    if not images:
+                        continue
+                    
+                    # Add each image to ZIP
+                    for idx, img_data in enumerate(images):
+                        image = img_data['image']
+                        
+                        # Convert to PIL Image
+                        img_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+                        pil_image = Image.fromarray(img_rgb)
+                        
+                        # Save to bytes
+                        img_buffer = BytesIO()
+                        pil_image.save(img_buffer, format='JPEG', quality=95)
+                        img_bytes = img_buffer.getvalue()
+                        
+                        # Add to ZIP
+                        filename = f"training_dataset/{safe_name}/image_{idx+1:04d}.jpg"
+                        zipf.writestr(filename, img_bytes)
+                        total_images += 1
+                    
+                    export_info.append({
+                        'dog_id': int(dog_id),
+                        'name': dog_name,
+                        'image_count': len(images)
+                    })
+                
+                # Add metadata
+                metadata = {
+                    'export_date': datetime.now().isoformat(),
+                    'total_dogs': len(dogs),
+                    'total_images': total_images,
+                    'dogs': export_info
+                }
+                
+                zipf.writestr('training_dataset/metadata.json', json.dumps(metadata, indent=2))
+            
+            # Save to temporary file
+            zip_buffer.seek(0)
+            temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.zip', prefix='dog_dataset_')
+            temp_file.write(zip_buffer.getvalue())
+            temp_file.close()
+            
+            summary = f"✅ Dataset exported successfully!\n\n"
+            summary += f"📦 Total dogs: {len(dogs)}\n"
+            summary += f"🖼️  Total images: {total_images}\n\n"
+            summary += "Click the download button below to save to your laptop."
+            
+            return summary, temp_file.name
+            
+        except Exception as e:
+            import traceback
+            error = f"Export error: {str(e)}\n{traceback.format_exc()}"
+            return error, None
+
+    def _img_to_base64(self, img_array: np.ndarray) -> str:
+        """Convert image array to base64 string"""
+        img_pil = Image.fromarray(img_array)
+        buffered = BytesIO()
+        img_pil.save(buffered, format="JPEG", quality=85)
+        return base64.b64encode(buffered.getvalue()).decode()
+
+    def create_interface(self):
+        """Create Gradio interface with validation checkboxes"""
+        
+        with gr.Blocks(title="Dog Dataset Collection", theme=gr.themes.Soft()) as app:
+            gr.Markdown("""
+            # 🐕 Dog Training Dataset Collection
+            **Process → Validate → Save → Export**
+            """)
+            
+            with gr.Tabs():
+                # TAB 1: Process Video
+                with gr.Tab("1. Process Video"):
+                    gr.Markdown("### Upload and process video to detect dogs")
+                    
+                    with gr.Row():
+                        with gr.Column():
+                            video_input = gr.Video(label="Upload Video")
+                            
+                            with gr.Row():
+                                reid_threshold = gr.Slider(
+                                    minimum=0.1, maximum=0.9, value=0.3, step=0.05,
+                                    label="ReID Threshold (lower = more dogs)"
+                                )
+                                sample_rate = gr.Slider(
+                                    minimum=1, maximum=10, value=3, step=1,
+                                    label="Frame Sampling Rate"
+                                )
+                            
+                            flip_camera = gr.Checkbox(label="Flip Camera Horizontally", value=False)
+                            
+                            with gr.Row():
+                                process_btn = gr.Button("🎬 Process Video", variant="primary", size="lg")
+                                stop_btn = gr.Button("⏹️ Stop", variant="stop")
+                                clear_btn = gr.Button("🗑️ Clear & Reset")
+                            
+                            progress_text = gr.Textbox(label="Progress", lines=1)
+                            status_text = gr.Textbox(label="Status", lines=8)
+                        
+                        with gr.Column():
+                            gallery_output = gr.HTML(label="Detection Results")
+                    
+                    with gr.Row():
+                        discard_btn = gr.Button("❌ Discard & Retry with Different Threshold", variant="secondary")
+                
+                # TAB 2: Validate & Save
+                with gr.Tab("2. Validate & Save"):
+                    gr.Markdown("### Review detected dogs and select images to keep")
+                    
+                    with gr.Column(visible=False) as validation_container:
+                        validation_status = gr.Textbox(label="Status", lines=2)
+                        
+                        load_btn = gr.Button("📋 Load Validation Interface", variant="primary", size="lg")
+                        
+                        # Dynamic validation area
+                        @gr.render(inputs=[], triggers=[load_btn.click])
+                        def render_validation():
+                            if not self.temp_session:
+                                gr.Markdown("No temporary session. Process a video first.")
+                                return
+                            
+                            checkboxes = []
+                            
+                            for temp_id in sorted(self.temp_session.keys()):
+                                dog_data = self.temp_session[temp_id]
+                                images = dog_data['images']
+                                
+                                with gr.Group():
+                                    gr.Markdown(f"### 🐕 Dog #{temp_id} - {len(images)} images")
+                                    
+                                    # Create grid of images with checkboxes
+                                    for i in range(0, len(images), 6):
+                                        with gr.Row():
+                                            for j in range(6):
+                                                if i + j < len(images):
+                                                    img = images[i + j]
+                                                    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+                                                    
+                                                    with gr.Column(scale=1, min_width=120):
+                                                        gr.Image(
+                                                            value=img_rgb,
+                                                            label=f"#{i+j+1}",
+                                                            interactive=False,
+                                                            height=150,
+                                                            show_download_button=False
+                                                        )
+                                                        cb = gr.Checkbox(
+                                                            label="Keep",
+                                                            value=True,
+                                                            elem_id=f"cb_{temp_id}_{i+j}"
+                                                        )
+                                                        checkboxes.append(cb)
+                            
+                            # Save button
+                            save_btn = gr.Button("💾 Save Selected to Database", variant="primary", size="lg")
+                            save_status = gr.Textbox(label="Save Status", lines=3)
+                            
+                            # Connect save button
+                            save_btn.click(
+                                fn=self.save_validated_to_database,
+                                inputs=checkboxes,
+                                outputs=[save_status, validation_container]
+                            )
+                
+                # TAB 3: Database & Export
+                with gr.Tab("3. Database & Export"):
+                    gr.Markdown("### View database and export for fine-tuning")
+                    
+                    refresh_db_btn = gr.Button("🔄 Refresh Database", variant="secondary")
+                    database_display = gr.HTML(label="Database Contents", visible=False)
+                    
+                    gr.Markdown("---")
+                    
+                    export_btn = gr.Button("📦 Export Dataset", variant="primary", size="lg")
+                    export_status = gr.Textbox(label="Export Status", lines=5)
+                    download_btn = gr.File(label="Download Exported Dataset", interactive=False)
+            
+            # Event handlers
+            process_btn.click(
+                fn=self.process_video,
+                inputs=[video_input, reid_threshold, flip_camera, sample_rate],
+                outputs=[gallery_output, status_text, progress_text, validation_container]
+            )
+            
+            stop_btn.click(
+                fn=self.stop_processing,
+                outputs=[status_text, progress_text, gallery_output]
+            )
+            
+            clear_btn.click(
+                fn=self.clear_reset,
+                outputs=[video_input, gallery_output, status_text, progress_text, validation_container]
+            )
+            
+            discard_btn.click(
+                fn=self.discard_session,
+                outputs=[validation_container, status_text, database_display]
+            )
+            
+            load_btn.click(
+                fn=self.load_validation_interface,
+                outputs=[validation_container, validation_status, gr.HTML()]
+            )
+            
+            refresh_db_btn.click(
+                fn=lambda: gr.update(value=self._show_database(), visible=True),
+                outputs=[database_display]
+            )
+            
+            export_btn.click(
+                fn=self.export_dataset,
+                outputs=[export_status, download_btn]
+            )
+        
+        return app
+
+    def launch(self):
+        """Launch the Gradio app"""
+        app = self.create_interface()
+        app.launch(share=False, server_name="0.0.0.0", server_port=7860)
+
+
+if __name__ == "__main__":
+    app = DatasetCollectionApp()
+    app.launch()