tracker.py

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # Suppress TensorFlow warnings

import random
import threading
import time
import tkinter as tk
from collections import deque
import numpy as np
from PIL import Image, ImageTk
import cv2
import mediapipe as mp
from playsound import playsound
from scipy.io import wavfile
from gradio_client import Client, handle_file
import uuid
import requests
import re
import json
import pandas as pd

screen_width = 1600
screen_height = 1080

# Face tracking setup
mp_face_mesh = mp.solutions.face_mesh
face_mesh = mp_face_mesh.FaceMesh(refine_landmarks=True)


# Camera setup
cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)



# Smoother for coordinates
class Smoother:
    def __init__(self, buffer_size=5):
        self.buffer = deque(maxlen=buffer_size)

    def smooth(self, x, y):
        self.buffer.append((x, y))
        avg_x = sum(p[0] for p in self.buffer) / len(self.buffer)
        avg_y = sum(p[1] for p in self.buffer) / len(self.buffer)
        return int(avg_x), int(avg_y)


smoother = Smoother()


def get_position(image):
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    results = face_mesh.process(image_rgb)

    if not results.multi_face_landmarks:
        return None, image

    landmarks = results.multi_face_landmarks[0].landmark
    left_eye = landmarks[133]  # left eye inner corner
    right_eye = landmarks[362]  # right eye inner corner

    mid_x = (left_eye.x + right_eye.x) / 2
    mid_y = (left_eye.y + right_eye.y) / 2

    # Mirror x-axis and boost y sensitivity
    x = int((1 - mid_x) * screen_width)
    y = int(mid_y * screen_height * 1.2)

    for idx in [1, 33, 263, 133, 362, 61]:
        lx = int(landmarks[idx].x * image.shape[1])
        ly = int(landmarks[idx].y * image.shape[0])
        cv2.circle(image, (lx, ly), 2, (0, 255, 0), -1)

    return smoother.smooth(x, y), image

class ProgressBar:
    def __init__(self, canvas, total_steps, margin=80, height=40):
        self.canvas = canvas
        self.total_steps = total_steps
        self.current_step = 0
        self.margin = margin
        self.height = height
        self.bg_bar = None
        self.fg_bar = None
        self.text = None

    def draw(self):
        bar_width = screen_width - 2 * self.margin
        x1 = self.margin
        x2 = screen_width - self.margin
        y1 = screen_height - self.height - 40
        y2 = y1 + self.height

        # Background rounded bar
        self.bg_bar = self.canvas.create_rectangle(
            x1, y1, x2, y2, fill="#E0E0E0", outline="", width=0
        )

        # Foreground animated bar
        self.fg_bar = self.canvas.create_rectangle(
            x1, y1, x1, y2, fill="#4CAF50", outline="", width=0
        )

        # Percentage text
        self.text = self.canvas.create_text(
            screen_width // 2,
            y1 - 10,
            text="0%",
            font=("Helvetica", 16, "bold"),
            fill="black",
        )

    def update(self, step):
        self.current_step = step
        if self.fg_bar is None:
            return

        x1 = self.margin
        y1 = screen_height - self.height - 40
        y2 = y1 + self.height
        percent = step / self.total_steps
        x2 = x1 + int((screen_width - 2 * self.margin) * percent)

        # Animate fill to the new width
        self.canvas.coords(self.fg_bar, x1, y1, x2, y2)
        self.canvas.itemconfig(self.text, text=f"{int(percent * 100)}%")



# Dwell menu screen using canvas
class DwellMenu:
    def __init__(self, canvas, items, dwell_time=2.0, padding=20, height=120, title=""):
        self.canvas = canvas
        self.items = items
        self.dwell_time = dwell_time
        self.padding = padding
        self.height = height
        self.selected = None
        self.hover_start = None
        self.title_text = None
        self.timer_circle = None
        self.rects = []

        self.timer_x = 50  # Example x-coordinate for the timer
        self.timer_y = screen_height - 150  # Example y-coordinate for the timer
        self.timer_radius = 0
        self.timer_max_radius = 60  # Maximum radius for the timer circle


        num_items = len(items)



        #always draw title if there is one
        if title:
            font_size = 60 if title == "Let's Pick a Genre For Your Song..." else 90
            self.title_text = self.canvas.create_text(
                screen_width // 2,
                100,
                text = title,
                font=("Helvetica MS", font_size, "bold"),
                fill="black",
            )

        #then check if there are 8 items and draw a grid
        if num_items == 8: #4 columns, 2 rows grid layout
            columns = 4
            rows = 2
            item_size = 200
            total_width = columns * item_size + (columns + 1) * self.padding
            start_x = (screen_width - total_width) // 2 + self.padding
            start_y = 220

            for idx, item in enumerate(items):
                row = idx // columns
                col = idx % columns
                x1 = start_x + col * (item_size + self.padding)
                y1 = start_y + row * (item_size + self.padding)
                x2 = x1 + item_size
                y2 = y1 + item_size

                color = f"#{random.randint(0x55, 0xFF):02x}{random.randint(0x55, 0xFF):02x}{random.randint(0x55, 0xFF):02x}"
                rect = self.canvas.create_rectangle(x1, y1, x2, y2, fill=color, outline="black", width=3)
                text = self.canvas.create_text(
                    (x1 + x2) // 2,
                    (y1 + y2) // 2,
                    text=item,
                    font=("Helvetica MS", 18, "bold"),
                    fill="white",
                )
                self.rects.append((rect, text, item, x1, y1, x2, y2))
            return  # Prevent vertical layout from being drawn on top

        if num_items == 1 and "Summary" in title or "Song Summary" in title:
            summary = items[0]
            self.title_text = self.canvas.create_text(
                screen_width // 2,
                screen_height // 2,
                text=summary,
                font=("Helvetica MS", 36, "bold"),
                fill="black",
                width=1000  # Wrap text if it's long
            )
            return  # Don't draw anything else

        else:
            # Standard vertical stacking
            item_width = 400
            item_height = 150
            x = screen_width // 2 - item_width // 2
            y = 250


            for item in items:
                custom_colors = {
                    "Create": "#4A90E2",
                    "Explore": "#7ED957",
                    "Library": "#9B59B6"
                }
                color = custom_colors.get(item, "#AAAAAA")
                rect = self.canvas.create_rectangle(
                    x, y, x + item_width, y + self.height, fill=color, outline="black", width=3
                )
                text = self.canvas.create_text(
                    x + item_width // 2,
                    y + self.height // 2,
                    text=item,
                    font=("Helvetica MS", 40, "bold"),
                    fill="white",
                )
                self.rects.append((rect, text, item, x, y, x + item_width, y + self.height))
                y += item_height + self.padding



        self.timer_x = 50
        self.timer_y = screen_height - 150
        self.timer_radius = 0
        self.timer_max_radius = 60







    def update(self, x, y):
        hovering = False
        for rect, text, label, x1, y1, x2, y2 in self.rects:
            if x1 <= x <= x2 and y1 <= y <= y2:
                hovering = True
                if self.hover_start is None:
                    self.hover_start = time.time()
                elif time.time() - self.hover_start >= self.dwell_time:
                    self.selected = label
                    return label
                self.canvas.itemconfig(rect, outline="yellow", width=5)
            else:
                self.canvas.itemconfig(rect, outline="black", width=3)

        if hovering:
            progress = (
                (time.time() - self.hover_start) / self.dwell_time
                if self.hover_start
                else 0
            )
            radius = int(self.timer_max_radius * progress)
            if self.timer_circle:
                self.canvas.delete(self.timer_circle)
            self.timer_circle = self.canvas.create_oval(
                self.timer_x - radius,
                self.timer_y - radius,
                self.timer_x + radius,
                self.timer_y + radius,
                fill="orange",
                outline="red",
                width=3,
            )
        else:
            if self.timer_circle:
                self.canvas.delete(self.timer_circle)
                self.timer_circle = None
            self.hover_start = None

        return None

    def destroy(self):
        if self.title_text:
            self.canvas.delete(self.title_text)
        for rect, text, *_ in self.rects:
            self.canvas.delete(rect)
            self.canvas.delete(text)
        if self.timer_circle:
            self.canvas.delete(self.timer_circle)
            self.timer_circle = None
# Tkinter UI
class HeadPoseApp:
    def __init__(self, root):
        self.selected_genre = ""
        self.selected_vibe = ""
        self.selected_bpm = ""
        self.root = root
        self.root.title("Head Tracker")

        try:
            # Initialize camera
            self.cap = cv2.VideoCapture(1)
            if not self.cap.isOpened():
                print("Error: Could not open camera")
                self.root.destroy()
                return

            # Create UI elements
            self.canvas = tk.Canvas(
                root, width=screen_width, height=screen_height, bg="white"
            )
            self.canvas.pack()

            self.dot_radius = 10
            self.dot = tk.Frame(self.canvas, width=20, height=20, bg="blue")
            self.dot.place(x=screen_width // 2, y=screen_height // 2)

            self.video_label = tk.Label(self.canvas)
            self.video_label.place(x=10, y=10)

            self.menus = [
                ("OpenOctave", ["Create", "Explore", "Library"]),
                ("Let's Pick a Genre For Your Song...", ["Pop", "Hip-Hop", "Jazz", "Rock", "Electronic", "Classical", "R&B", "Country"]),
                ("Lastly, Pick Your Vibe", ["Sad", "Happy", "Romantic", "Energetic", "Chill", "Angry", "Mysterious", "Dreamy"]),
                ("Pick the BPM", ["0-60", "60-80", "80-100", "100-120", "120-140", "140-160", "160-180", "180+ BPM"]), 
                ("Summary", [""], "Show it To Me"), 
                ("Listen to your thoughts...", ["Visualizer"]),  
            ] 

            self.progress_bar = ProgressBar(self.canvas, total_steps=len(self.menus) - 1)
            self.current_menu_index = 0
            self.dwell_menu = None
            self.create_next_menu()

            self.running = True
            self.update_thread = threading.Thread(target=self.pose_loop)
            self.update_thread.daemon = True
            self.update_thread.start()

            self.root.protocol("WM_DELETE_WINDOW", self.on_close)

        except Exception as e:
            print(f"Error initializing application: {str(e)}")
            self.root.destroy()
            return

        self.canvas = tk.Canvas(
            root, width=screen_width, height=screen_height, bg="white"
        )
        self.canvas.pack()

        self.dot_radius = 10
        self.dot = tk.Frame(self.canvas, width=20, height=20, bg="blue")
        self.dot.place(x=screen_width // 2, y=screen_height // 2)

        self.video_label = tk.Label(self.canvas)
        self.video_label.place(x=10, y=10)

        self.menus = [
            ("OpenOctave", ["Create", "Explore", "Library"]),
            ("Let's Pick a Genre For Your Song...", ["Pop", "Hip-Hop", "Jazz", "Rock", "Electronic", "Classical", "R&B", "Country"]),
            ("Lastly, Pick Your Vibe", ["Sad", "Happy", "Romantic", "Energetic", "Chill", "Angry", "Mysterious", "Dreamy"]),
            ("Pick the BPM", ["0-60", "60-80", "80-100", "100-120", "120-140", "140-160", "160-180", "180+ BPM"]),
            ("Summary", [""], "Show it To Me"), #This will be upgraded dynamically
            ("Listen to your thoughts...", ["Visualizer"]),  # NEW men
        ]

        self.progress_bar = ProgressBar(self.canvas, total_steps=len(self.menus) - 1)
        self.current_menu_index = 0
        self.dwell_menu = None
        self.create_next_menu()

        self.running = True
        self.update_thread = threading.Thread(target=self.pose_loop)
        self.update_thread.daemon = True
        self.update_thread.start()

        self.root.protocol("WM_DELETE_WINDOW", self.on_close)


    def runTimeSeries(self):
        switch = True
        timeSeries = pd.DataFrame("Excitement", "Engagement", "Stress")
        while switch:
            

    def generate_default_audio(self, prompt):
        print("Generating BROOOOOO..." + prompt)
        url = "http://127.0.0.1:7860/gradio_api/queue/join"
        session_hash = str(uuid.uuid4())
        payload = {
            "data": [
                prompt,
                "", 
                0, 
                47, 
                7, 
                100, 
                0, 
                "-1", 
                "dpmpp-3m-sde", 
                0.01, 
                100, 
                1, 
                0, 
                1, 
                0, 
                "wav", 
                "output.wav", 
                True, 
                None, 
                0.1, 
                10, 
                47, 
                None
            ],
            "event_data": None,
            "fn_index": 1,
            "trigger_id": 7,
            "session_hash": session_hash
        }

        # Make the POST request and receive the response
        response = requests.post(url, json=payload)

        url2 = f"http://127.0.0.1:7860/gradio_api/queue/data?session_hash={session_hash}"

        response2 = requests.get(url2)
        content = response2.content.decode('utf-8')
        
        # Split the content by data: and process each line
        lines = content.split('data: ')
        for line in lines:
            if line.strip():
                try:
                    # Remove any trailing newlines and parse as JSON
                    json_data = json.loads(line.strip())
                    if 'output' in json_data and 'data' in json_data['output']:
                        # The audio URL is in the first element of the data array
                        audio_url = json_data['output']['data'][0]['url']
                        print(f"Found audio URL: {audio_url}")
                        return audio_url
                except json.JSONDecodeError:
                    continue

        print("No audio URL found in response")
        return None

    def play_audio_background(self, audio_url):
        """Play audio file in the background"""
        if audio_url:
            # Download the audio file
            response = requests.get(audio_url)
            if response.status_code == 200:
                # Save to temporary file
                temp_file = "temp_audio.wav"
                with open(temp_file, 'wb') as f:
                    f.write(response.content)
                
                # Play the audio in a separate thread
                threading.Thread(target=lambda: playsound(temp_file), daemon=True).start()

    def pose_loop(self):
        while self.running:
            ret, frame = self.cap.read()
            if not ret:
                continue

            point, annotated_frame = get_position(frame)
            if point:
                x, y = point
                x = max(0, min(screen_width, x))
                y = max(0, min(screen_height, y))
                self.move_dot(x, y)
                selected = self.dwell_menu.update(x, y)
                if selected:
                    print("Selected:", selected)
                    if self.current_menu_index == 1:  # Genre menu
                        self.selected_genre = selected
                    elif self.current_menu_index == 2:  # Vibe menu
                        self.selected_vibe = selected
                    elif self.current_menu_index == 3:  # BPM menu
                        self.selected_bpm = selected

                    self.dwell_menu.destroy()
                    self.current_menu_index += 1

                    if self.current_menu_index == 4:
                        print("Generating audio... AGAIN")
                        summary_text = f"You generated a song that is {self.selected_genre} with a {self.selected_vibe} vibe and a tempo of {self.selected_bpm}. Here's what we created for you."
                        self.menus[4] = ("Song Summary", [summary_text])
                        
                        # Generate and play audio
                        audio_url = self.generate_default_audio(summary_text)
                        self.play_audio_background(audio_url)
                    self.create_next_menu()


            # Resize and convert frame for Tkinter
            display_frame = cv2.resize(annotated_frame, (200, 150))
            img = cv2.cvtColor(display_frame, cv2.COLOR_BGR2RGB)
            imgtk = ImageTk.PhotoImage(image=Image.fromarray(img))
            self.video_label.imgtk = imgtk
            self.video_label.configure(image=imgtk)

            time.sleep(0.01)

    def move_dot(self, x, y):
        self.dot.place(x=x, y=y)

    def create_next_menu(self):
        if self.current_menu_index < len(self.menus):
            title, items = self.menus[self.current_menu_index]
            self.dwell_menu = DwellMenu(self.canvas, items, dwell_time=2.0, title=title)

            if self.current_menu_index > 0:
                self.progress_bar.draw()
                self.progress_bar.update(self.current_menu_index)
        else:
            print("All menus complete.")

    def on_close(self):
        self.running = False
        if hasattr(self, 'cap'):
            self.cap.release()
        self.root.destroy()

if __name__ == "__main__":
    root = tk.Tk()
    app = HeadPoseApp(root)
    root.mainloop()