app.py

import streamlit as st
import cv2
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="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, model_path=None):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        
        # Initialize YOLO model for dog detection (optional)
        try:
            self.yolo_model = YOLO(model_path) if model_path else None
        except Exception as e:
            st.warning("YOLO model not found. Running without dog detection.")
            self.yolo_model = None
        
        # Initialize behavior classifier
        self.num_behaviors = 5
        try:
            self.behavior_model = BehaviorDetector(self.num_behaviors).to(self.device)
        except Exception as e:
            st.warning("Error loading behavior model. Using default classifier.")
            self.behavior_model = models.resnet18(pretrained=True)
            num_features = self.behavior_model.fc.in_features
            self.behavior_model.fc = nn.Linear(num_features, self.num_behaviors)
            self.behavior_model = self.behavior_model.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]),
            transforms.RandomHorizontalFlip(p=0.3),
            transforms.ColorJitter(brightness=0.2, contrast=0.2)
        ])
        
        # Behavior definitions with confidence thresholds
        self.behaviors = {
            '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):
        """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)
        
        with torch.no_grad():
            predictions = self.behavior_model(input_tensor).squeeze().cpu().numpy()
        
        # Update behavior history
        for behavior, pred in zip(self.behaviors.keys(), predictions):
            self.behavior_history[behavior].append(pred)
        
        # 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 detected_behaviors

def main():
    st.title("🐕 Advanced Dog Language Understanding")
    st.write("Upload a video of your dog for detailed behavior analysis!")

    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
    )

    video_file = st.file_uploader("Upload Video", type=['mp4', 'avi', 'mov'])

    if video_file is not None:
        tfile = tempfile.NamedTemporaryFile(delete=False)
        tfile.write(video_file.read())

        cap = cv2.VideoCapture(tfile.name)
        
        col1, col2 = st.columns(2)
        
        with col1:
            st.subheader("Video Analysis")
            video_placeholder = st.empty()
            
        with col2:
            st.subheader("Real-time Behavior Detection")
            analysis_placeholder = st.empty()
            
        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))
        
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break
                
            frame_count += 1
            progress = frame_count / total_frames
            progress_bar.progress(progress)
            
            # 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(annotated_frame, cv2.COLOR_BGR2RGB),
                channels="RGB",
                use_container_width=True
            )
            
            # Update analysis display with confidence scores
            analysis_text = "Detected Behaviors:\n\n"
            for behavior, count in behavior_counts.items():
                if count > 0:
                    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_text_{frame_count}"
            )
            
            time.sleep(0.05)
            
        cap.release()
        
        # Final analysis
        st.subheader("Comprehensive Analysis")
        
        # 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("Primary Behavior", most_common.replace('_', ' ').title())
            
        with col2:
            total_behaviors = sum(behavior_counts.values())
            st.metric("Total Behaviors", total_behaviors)
            
        with col3:
            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 based on analysis
        st.subheader("Personalized Recommendations")
        if total_behaviors > 0:
            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 clear behaviors detected. Try recording with better lighting and closer camera angle.")

if __name__ == "__main__":
    main()