Autonomous NBA Shot-Tracking Engine

import streamlit as st
import cv2
import tempfile
import os
import numpy as np
from collections import deque
from ultralytics import YOLO

st.set_page_config(page_title="NBA Vision (Final)", layout="wide", page_icon="🏀")
st.title("🏀 NBA Vision: Production Candidate")

# --- 1. UTILS ---
class RimMemory:
    def __init__(self):
        self.last_rim = None
        self.frames_without_rim = 0
    def update(self, raw_rim):
        if raw_rim is not None:
            self.last_rim = raw_rim
            self.frames_without_rim = 0
            return raw_rim
        elif self.last_rim is not None and self.frames_without_rim < 60:
            self.frames_without_rim += 1
            return self.last_rim
        return None

# --- GRAVITY KILLER (DUMB SMOOTHER) ---
class DumbSmoother:
    def __init__(self, history_size=3):
        self.history = deque(maxlen=history_size)

    def update(self, raw_box, rim_box=None):
        if raw_box is not None:
            self.history.append(raw_box)
            return np.mean(self.history, axis=0).astype(int).tolist()
        return None 

# --- 2. SHOT DETECTOR (HIGH & TIGHT) ---
class ShotDetector:
    def __init__(self):
        self.basket_made = False
        self.shot_in_progress = False 
        self.cooldown = 0
        
    def update(self, ball_box, rim_box):
        # Allow multiple shots in a clip (with cooldown) or stick to binary?
        # Sticking to binary "Did a Make Happen?" for this demo
        if self.basket_made: return True 
        if ball_box is None or rim_box is None: return False

        bx, by, bw, bh = ball_box
        rx, ry, rw, rh = rim_box
        
        ball_cy = int(by + bh // 2)
        ball_cx = int(bx + bw // 2)

        rim_center_x = int(rx + rw // 2)
        
        # --- ZONES (Hidden Logic) ---
        # 1. ENTRY (Wide): 50% width
        entry_width = int(rw * 0.5)
        ex1 = int(rim_center_x - entry_width // 2)
        ex2 = int(rim_center_x + entry_width // 2)
        entry_y = ry
        
        # 2. EXIT (High & Tight): 25% width at 50% depth
        exit_width = int(rw * 0.25)
        fx1 = int(rim_center_x - exit_width // 2)
        fx2 = int(rim_center_x + exit_width // 2)
        finish_line_y = ry + int(rh * 0.50) 

        # --- LOGIC ---
        
        # A. DETECT ENTRY
        if (ex1 < ball_cx < ex2) and (entry_y < ball_cy < entry_y + rh*0.4):
            self.shot_in_progress = True 

        # B. DETECT FINISH
        if self.shot_in_progress:
            if ball_cy > finish_line_y:
                # CHECK: Is it centered in the narrow neck?
                if (fx1 < ball_cx < fx2): 
                    self.basket_made = True
                    self.shot_in_progress = False 
                    return True
                else:
                    # REJECTED: Ball drifted wide (Parallax/Floater Miss)
                    self.shot_in_progress = False
            
            # C. ABORT (Popped up)
            if ball_cy < ry:
                self.shot_in_progress = False
        
        return False

# --- 3. LOAD MODELS ---
@st.cache_resource
def load_models():
    # Use the best weights
    ball_path = 'runs/detect/ball_model_blur_fix/weights/best.pt'
    hoop_path = 'runs/detect/nba_hoop_custom_v1/weights/best.pt'
    
    ball_model = YOLO(ball_path) if os.path.exists(ball_path) else YOLO('yolo11n.pt')
    hoop_model = YOLO(hoop_path) if os.path.exists(hoop_path) else YOLO('yolo11n.pt')
    
    return ball_model, hoop_model

ball_model, hoop_model = load_models()

# --- MAIN UI ---
st.sidebar.header("Configuration")
ball_conf = st.sidebar.slider("Ball Confidence", 0.1, 0.9, 0.25)
hoop_conf = st.sidebar.slider("Hoop Confidence", 0.3, 0.9, 0.5)

uploaded_file = st.file_uploader("Upload Clip", type=["mp4", "mov"])

if uploaded_file is not None:
    tfile = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
    tfile.write(uploaded_file.read())
    
    cap = cv2.VideoCapture(tfile.name)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    
    # OUTPUT
    output_path = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4').name
    try: fourcc = cv2.VideoWriter_fourcc(*'avc1')
    except: fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
    
    progress_bar = st.progress(0)
    frame_cnt = 0
    
    rim_memory = RimMemory()
    smoother = DumbSmoother() # NO GRAVITY
    scorer = ShotDetector()
    
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret: break
        
        # 1. AI
        raw_ball, raw_rim = None, None
        
        b_res = ball_model.track(frame, persist=True, verbose=False, conf=ball_conf)
        if b_res[0].boxes:
            b = b_res[0].boxes[0].xywh.cpu().numpy()[0]
            raw_ball = [b[0], b[1], b[2], b[3]]
            
        r_res = hoop_model(frame, verbose=False, conf=hoop_conf)
        if r_res[0].boxes:
            r = r_res[0].boxes[0].xywh.cpu().numpy()[0]
            raw_rim = [r[0]-r[2]/2, r[1]-r[3]/2, r[2], r[3]]

        # 2. PHYSICS
        stable_rim = rim_memory.update(raw_rim)
        final_ball = smoother.update(raw_ball) 
        
        # 3. LOGIC
        ball_box_score = None
        if final_ball:
            bx, by, bw, bh = final_ball
            ball_box_score = [int(bx-bw/2), int(by-bh/2), int(bw), int(bh)]
        
        scorer.update(ball_box_score, stable_rim)
        
        # 4. DRAW
        # Draw Rim (Clean Green)
        if stable_rim:
            rx, ry, rw, rh = map(int, stable_rim)
            cv2.rectangle(frame, (rx, ry), (rx+rw, ry+rh), (0, 255, 0), 2)
            cv2.putText(frame, "RIM", (rx, ry-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

        # Draw Ball (Clean Yellow)
        if ball_box_score:
            x, y, w, h = ball_box_score
            cv2.rectangle(frame, (x, y), (x+w, y+h), (0, 255, 255), 2)
            cv2.putText(frame, "BASKETBALL", (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 2)

        # Draw Status
        if scorer.basket_made:
             # Flash Green Border
             cv2.rectangle(frame, (0,0), (width, height), (0, 255, 0), 10)
             # Center Text
             cv2.putText(frame, "BASKET!", (width//2 - 100, 100), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 255, 0), 4)

        out.write(frame)
        frame_cnt += 1
        progress_bar.progress(min(frame_cnt / total_frames, 1.0))

    cap.release()
    out.release()
    progress_bar.empty()
    
    result = 1 if scorer.basket_made else 0
    if result == 1:
        st.success(f"### Made Basket: {result}")
    else:
        st.error(f"### Made Basket: {result}")

    with open(output_path, 'rb') as f:
        st.video(f.read())