app.py

import streamlit as st
import cv2
import numpy as np
import tempfile
from PIL import Image
import torch
from torchvision import transforms, models
import time
import plotly.graph_objects as go
from PIL import Image, ImageDraw
import base64
from io import BytesIO
import pandas as pd

# Set page config
st.set_page_config(
    page_title="Dog Language Understanding",
    page_icon="🐕",
    layout="wide"
)

class DogBehaviorAnalyzer:
    def __init__(self):
        # Initialize model (using pretrained ResNet for this example)
        self.model = models.resnet50(pretrained=True)
        self.model.eval()
        
        # Define image transformations
        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])
        ])
        
        # Enhanced behavior mappings with emotions and tips
        self.behaviors = {
            'tail_wagging': {
                'emotion': 'Happy and excited',
                'description': 'Your dog is expressing joy and enthusiasm!',
                'tips': [
                    'This is a great time for positive reinforcement training',
                    'Consider engaging in play or exercise',
                    'Use this excitement for teaching new tricks'
                ]
            },
            'barking': {
                'emotion': 'Alert or communicative',
                'description': 'Your dog is trying to communicate or alert you.',
                'tips': [
                    'Check what triggered the barking',
                    'Use positive reinforcement for quiet behavior',
                    'Consider training "quiet" and "speak" commands'
                ]
            },
            'ears_perked': {
                'emotion': 'Alert and interested',
                'description': 'Your dog is focused and attentive.',
                'tips': [
                    'Great moment for training exercises',
                    'Consider mental stimulation activities',
                    'Use this attention for bonding exercises'
                ]
            },
            'lying_down': {
                'emotion': 'Relaxed and comfortable',
                'description': 'Your dog is calm and at ease.',
                'tips': [
                    'Perfect time for gentle petting',
                    'Maintain a peaceful environment',
                    'Consider quiet bonding activities'
                ]
            },
            'jumping': {
                'emotion': 'Excited and playful',
                'description': 'Your dog is energetic and seeking interaction!',
                'tips': [
                    'Channel energy into structured play',
                    'Practice "four paws on floor" training',
                    'Consider agility exercises'
                ]
            }
        }

    def analyze_frame(self, frame):
        # Convert frame to PIL Image
        image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        
        # Transform image
        input_tensor = self.transform(image)
        input_batch = input_tensor.unsqueeze(0)
        
        # Simulate behavior detection
        # In a real implementation, you'd use a properly trained model
        behaviors = []
        confidence_scores = np.random.random(len(self.behaviors))
        
        for behavior, score in zip(self.behaviors.keys(), confidence_scores):
            if score > 0.7:  # Threshold for detection
                behaviors.append(behavior)
        
        return behaviors

    def create_animation(self, behavior):
        """Create simple animations for behaviors"""
        # Create a simple animation frame
        img = Image.new('RGB', (200, 200), 'white')
        draw = ImageDraw.Draw(img)
        
        if behavior == 'tail_wagging':
            # Draw a simple tail wagging animation
            draw.arc([50, 50, 150, 150], 0, 180, fill='black', width=2)
        elif behavior == 'barking':
            # Draw speech-bubble like shapes
            draw.ellipse([50, 50, 150, 150], outline='black', width=2)
        
        # Convert to base64 for display
        buffered = BytesIO()
        img.save(buffered, format="PNG")
        return base64.b64encode(buffered.getvalue()).decode()

def main():
    st.title("🐕 Dog Language Understanding")
    st.write("Upload a video of your dog to analyze their behavior and emotions!")

    # Initialize analyzer
    analyzer = DogBehaviorAnalyzer()

    # 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")
            video_placeholder = st.empty()
            
        with col2:
            st.subheader("Real-time Analysis")
            analysis_placeholder = st.empty()
            emotion_placeholder = st.empty()
            
        # Progress bar
        progress_bar = st.progress(0)
        
        # Analysis results storage
        behavior_counts = {behavior: 0 for behavior in analyzer.behaviors.keys()}
        current_emotions = set()
        
        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
            video_placeholder.image(
                cv2.cvtColor(frame, cv2.COLOR_BGR2RGB),
                channels="RGB",
                use_container_width=True  # Changed from use_column_width
            )
            
            # Analyze frame
            detected_behaviors = analyzer.analyze_frame(frame)
            for behavior in detected_behaviors:
                behavior_counts[behavior] += 1
                current_emotions.add(behavior)
            
            # Update analysis display with enhanced information
            analysis_text = "Detected Behaviors & Emotions:\n\n"
            for behavior, count in behavior_counts.items():
                if count > 0:
                    behavior_info = analyzer.behaviors[behavior]
                    analysis_text += f"• {behavior.replace('_', ' ').title()}: {count} times\n"
                    analysis_text += f"  Emotion: {behavior_info['emotion']}\n"
                    analysis_text += f"  {behavior_info['description']}\n\n"
                    
                    # Add animation for the behavior
                    animation_data = analyzer.create_animation(behavior)
                    st.image(f"data:image/png;base64,{animation_data}", 
                            width=100, 
                            caption=f"{behavior.replace('_', ' ').title()} Animation")
            
            analysis_placeholder.text_area(
                "Analysis Results",
                analysis_text,
                height=300
            )
            
            # Display training tips
            if len(current_emotions) > 0:
                st.subheader("Training Tips")
                for behavior in current_emotions:
                    tips = analyzer.behaviors[behavior]['tips']
                    for tip in tips:
                        st.info(tip)
            
            time.sleep(0.1)
        
        cap.release()
        
        # Final summary with enhanced visualizations
        st.subheader("Analysis Summary")
        
        # Create emotion timeline
        emotions_df = pd.DataFrame(list(current_emotions), columns=['Emotion'])
        fig = go.Figure(data=[go.Scatter(x=emotions_df.index, 
                                       y=emotions_df['Emotion'],
                                       mode='lines+markers')])
        fig.update_layout(title='Emotional Journey',
                         xaxis_title='Time',
                         yaxis_title='Emotion')
        st.plotly_chart(fig)
        
        # Final summary
        st.subheader("Analysis Summary")
        st.write("Overall behavior analysis of your dog:")
        
        # Create summary 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())
            
        with col2:
            total_behaviors = sum(behavior_counts.values())
            st.metric("Total Behaviors Detected", 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")
        
        # Enhanced recommendations
        st.subheader("Personalized Recommendations")
        if total_behaviors > 0:
            dominant_behavior = max(behavior_counts.items(), key=lambda x: x[1])[0]
            st.write(f"""
            Based on the analysis, here are personalized recommendations for {dominant_behavior}:
            
            {' '.join(analyzer.behaviors[dominant_behavior]['tips'])}
            
            General 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
            """)
        else:
            st.write("No behaviors detected. Try uploading a different video with clearer dog movements.")

if __name__ == "__main__":
    main()