v1

app.py

@@ -4,92 +4,156 @@ import numpy as np
 import tempfile
 from PIL import Image
 import torch
+import torch.nn as nn
 from torchvision import transforms, models
 import time
+from collections import deque
+import yaml
+from ultralytics import YOLO
 
 # Set page config
 st.set_page_config(
-    page_title="Dog Language Understanding",
+    page_title="Advanced Dog Language Understanding",
     page_icon="🐕",
     layout="wide"
 )
 
+class BehaviorDetector(nn.Module):
+    def __init__(self, num_classes):
+        super(BehaviorDetector, self).__init__()
+        # Use EfficientNet as base model (better performance than ResNet)
+        self.base_model = models.efficientnet_b0(pretrained=True)
+        # Replace classifier
+        num_features = self.base_model.classifier[1].in_features
+        self.base_model.classifier = nn.Sequential(
+            nn.Dropout(p=0.2),
+            nn.Linear(num_features, num_classes)
+        )
+        
+    def forward(self, x):
+        return torch.sigmoid(self.base_model(x))
+
 class DogBehaviorAnalyzer:
-    def __init__(self):
-        # Initialize model (using pretrained ResNet for this example)
-        self.model = models.resnet50(pretrained=True)
-        self.model.eval()
+    def __init__(self, model_path='best.pt'):
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        
+        # Initialize YOLO model for dog detection
+        self.yolo_model = YOLO(model_path) if model_path else None
         
-        # Define image transformations
+        # Initialize behavior classifier
+        self.num_behaviors = 5
+        self.behavior_model = BehaviorDetector(self.num_behaviors).to(self.device)
+        self.behavior_model.eval()
+        
+        # Define sophisticated transforms
         self.transform = transforms.Compose([
             transforms.Resize((224, 224)),
             transforms.ToTensor(),
             transforms.Normalize(mean=[0.485, 0.456, 0.406],
-                              std=[0.229, 0.224, 0.225])
+                              std=[0.229, 0.224, 0.225]),
+            transforms.RandomHorizontalFlip(p=0.3),
+            transforms.ColorJitter(brightness=0.2, contrast=0.2)
         ])
         
-        # Define behavior mappings
+        # Behavior definitions with confidence thresholds
         self.behaviors = {
-            'tail_wagging': 'Your dog is happy and excited!',
-            'barking': 'Your dog is trying to communicate or alert you.',
-            'ears_perked': 'Your dog is alert and interested.',
-            'lying_down': 'Your dog is relaxed and comfortable.',
-            'jumping': 'Your dog is energetic and playful!'
+            'tail_wagging': {'threshold': 0.75, 'description': 'Your dog is displaying happiness and excitement!'},
+            'barking': {'threshold': 0.8, 'description': 'Your dog is trying to communicate or alert you.'},
+            'ears_perked': {'threshold': 0.7, 'description': 'Your dog is alert and interested in something.'},
+            'lying_down': {'threshold': 0.85, 'description': 'Your dog is relaxed and comfortable.'},
+            'jumping': {'threshold': 0.8, 'description': 'Your dog is energetic and playful!'}
         }
+        
+        # Temporal smoothing using sliding window
+        self.behavior_history = {behavior: deque(maxlen=5) for behavior in self.behaviors.keys()}
+        
+    def preprocess_frame(self, frame):
+        """Advanced frame preprocessing"""
+        # Convert BGR to RGB
+        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        
+        # Apply adaptive histogram equalization
+        lab = cv2.cvtColor(rgb_frame, cv2.COLOR_RGB2LAB)
+        l, a, b = cv2.split(lab)
+        clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
+        cl = clahe.apply(l)
+        enhanced = cv2.merge((cl,a,b))
+        enhanced = cv2.cvtColor(enhanced, cv2.COLOR_LAB2RGB)
+        
+        return Image.fromarray(enhanced)
+
+    def detect_dog(self, frame):
+        """Detect dog in frame using YOLO"""
+        if self.yolo_model is None:
+            return True  # Skip detection if no model
+            
+        results = self.yolo_model(frame)
+        return len(results) > 0 and any(result.boxes for result in results)
 
     def analyze_frame(self, frame):
-        # Convert frame to PIL Image
-        image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+        """Analyze frame with temporal smoothing and confidence thresholds"""
+        # First detect if dog is present
+        if not self.detect_dog(frame):
+            return []
+            
+        # Preprocess frame
+        processed_frame = self.preprocess_frame(frame)
+        input_tensor = self.transform(processed_frame).unsqueeze(0).to(self.device)
         
-        # Transform image
-        input_tensor = self.transform(image)
-        input_batch = input_tensor.unsqueeze(0)
+        with torch.no_grad():
+            predictions = self.behavior_model(input_tensor).squeeze().cpu().numpy()
         
-        # Simulate behavior detection
-        # In a real implementation, you'd use a properly trained model
-        behaviors = []
-        confidence_scores = np.random.random(len(self.behaviors))
+        # Update behavior history
+        for behavior, pred in zip(self.behaviors.keys(), predictions):
+            self.behavior_history[behavior].append(pred)
         
-        for behavior, score in zip(self.behaviors.keys(), confidence_scores):
-            if score > 0.7:  # Threshold for detection
-                behaviors.append(behavior)
+        # Apply temporal smoothing and thresholds
+        detected_behaviors = []
+        for behavior, history in self.behavior_history.items():
+            if len(history) > 0:
+                avg_conf = sum(history) / len(history)
+                if avg_conf > self.behaviors[behavior]['threshold']:
+                    detected_behaviors.append((behavior, avg_conf))
         
-        return behaviors
+        return detected_behaviors
 
 def main():
-    st.title("🐕 Dog Language Understanding")
-    st.write("Upload a video of your dog to analyze their behavior!")
+    st.title("🐕 Advanced Dog Language Understanding")
+    st.write("Upload a video of your dog for detailed behavior analysis!")
 
-    # Initialize analyzer
     analyzer = DogBehaviorAnalyzer()
+    
+    # Add model confidence control
+    confidence_threshold = st.sidebar.slider(
+        "Detection Confidence Threshold",
+        min_value=0.5,
+        max_value=0.95,
+        value=0.7,
+        step=0.05
+    )
 
-    # File uploader
     video_file = st.file_uploader("Upload Video", type=['mp4', 'avi', 'mov'])
 
     if video_file is not None:
-        # Save uploaded file temporarily
         tfile = tempfile.NamedTemporaryFile(delete=False)
         tfile.write(video_file.read())
 
-        # Video analysis
         cap = cv2.VideoCapture(tfile.name)
         
-        # Create columns for layout
         col1, col2 = st.columns(2)
         
         with col1:
-            st.subheader("Video Preview")
+            st.subheader("Video Analysis")
             video_placeholder = st.empty()
             
         with col2:
-            st.subheader("Real-time Analysis")
+            st.subheader("Real-time Behavior Detection")
             analysis_placeholder = st.empty()
             
-            # Progress bar
-            progress_bar = st.progress(0)
-            
-        # Analysis results storage
+        progress_bar = st.progress(0)
+        
         behavior_counts = {behavior: 0 for behavior in analyzer.behaviors.keys()}
+        confidence_history = {behavior: [] for behavior in analyzer.behaviors.keys()}
         
         frame_count = 0
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
@@ -103,67 +167,91 @@ def main():
             progress = frame_count / total_frames
             progress_bar.progress(progress)
             
-            # Update video preview
+            # Update video preview with annotations
+            annotated_frame = frame.copy()
+            detected_behaviors = analyzer.analyze_frame(frame)
+            
+            # Draw behavior labels on frame
+            y_pos = 30
+            for behavior, conf in detected_behaviors:
+                if conf > confidence_threshold:
+                    behavior_counts[behavior] += 1
+                    confidence_history[behavior].append(conf)
+                    cv2.putText(annotated_frame, 
+                              f"{behavior.replace('_', ' ').title()}: {conf:.2f}",
+                              (10, y_pos), 
+                              cv2.FONT_HERSHEY_SIMPLEX, 
+                              0.7, 
+                              (0, 255, 0), 
+                              2)
+                    y_pos += 30
+            
             video_placeholder.image(
-                cv2.cvtColor(frame, cv2.COLOR_BGR2RGB),
+                cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB),
                 channels="RGB",
-                use_container_width=True  # Changed from use_column_width
+                use_container_width=True
             )
             
-            # Analyze frame
-            detected_behaviors = analyzer.analyze_frame(frame)
-            for behavior in detected_behaviors:
-                behavior_counts[behavior] += 1
-            
-            # Update analysis display
+            # Update analysis display with confidence scores
             analysis_text = "Detected Behaviors:\n\n"
             for behavior, count in behavior_counts.items():
                 if count > 0:
-                    analysis_text += f"• {behavior.replace('_', ' ').title()}: {count} times\n"
-                    analysis_text += f"  {analyzer.behaviors[behavior]}\n\n"
+                    avg_conf = sum(confidence_history[behavior]) / len(confidence_history[behavior])
+                    analysis_text += f"• {behavior.replace('_', ' ').title()}:\n"
+                    analysis_text += f"  Count: {count} | Avg Confidence: {avg_conf:.2f}\n"
+                    analysis_text += f"  {analyzer.behaviors[behavior]['description']}\n\n"
             
             analysis_placeholder.text_area(
                 "Analysis Results",
                 analysis_text,
-                height=300,
-                key=f"analysis_{frame_count}"  # Added unique key for each frame
+                height=300
             )
             
-            time.sleep(0.1)  # Add small delay for visualization
+            time.sleep(0.05)
             
         cap.release()
         
-        # Final summary
-        st.subheader("Analysis Summary")
-        st.write("Overall behavior analysis of your dog:")
+        # Final analysis
+        st.subheader("Comprehensive Analysis")
         
-        # Create summary metrics
+        # Create detailed metrics
         col1, col2, col3 = st.columns(3)
         
         with col1:
             most_common = max(behavior_counts.items(), key=lambda x: x[1])[0]
-            st.metric("Most Common Behavior", most_common.replace('_', ' ').title())
+            st.metric("Primary Behavior", most_common.replace('_', ' ').title())
             
         with col2:
             total_behaviors = sum(behavior_counts.values())
-            st.metric("Total Behaviors Detected", total_behaviors)
+            st.metric("Total Behaviors", total_behaviors)
             
         with col3:
-            behavior_variety = len([b for b in behavior_counts.values() if b > 0])
-            st.metric("Behavior Variety", f"{behavior_variety} types")
+            avg_confidence = np.mean([np.mean(conf) for conf in confidence_history.values() if conf])
+            st.metric("Average Confidence", f"{avg_confidence:.2%}")
+        
+        # Behavior distribution chart
+        st.subheader("Behavior Distribution")
+        behavior_data = {k.replace('_', ' ').title(): v for k, v in behavior_counts.items() if v > 0}
+        st.bar_chart(behavior_data)
         
-        # Recommendations
-        st.subheader("Recommendations")
+        # Recommendations based on analysis
+        st.subheader("Personalized Recommendations")
         if total_behaviors > 0:
-            st.write("""
-            Based on the analysis, here are some recommendations:
-            - Maintain regular exercise routines
-            - Provide mental stimulation through toys and training
-            - Continue positive reinforcement training
-            - Monitor your dog's body language for better communication
-            """)
+            st.write("Based on the detailed analysis, here are tailored recommendations:")
+            
+            # Generate specific recommendations based on detected behaviors
+            recommendations = []
+            if behavior_counts['tail_wagging'] > total_behaviors * 0.3:
+                recommendations.append("• Your dog shows frequent happiness - great time for training!")
+            if behavior_counts['barking'] > total_behaviors * 0.2:
+                recommendations.append("• Consider quiet command training to manage barking")
+            if behavior_counts['jumping'] > total_behaviors * 0.25:
+                recommendations.append("• Focus on calm behavior reinforcement")
+                
+            for rec in recommendations:
+                st.write(rec)
         else:
-            st.write("No behaviors detected. Try uploading a different video with clearer dog movements.")
+            st.write("No clear behaviors detected. Try recording with better lighting and closer camera angle.")
 
 if __name__ == "__main__":
-    main()
+    main()