app.py

import os
import json
import tempfile
import time
import uuid
import cv2
import numpy as np
import gradio as gr
import boto3
import torch
import subprocess
from botocore.exceptions import ClientError
from ultralytics import YOLO
from concurrent.futures import ThreadPoolExecutor
import threading
import sys
import traceback
from dotenv import load_dotenv
import datetime
import requests
import random

load_dotenv()

# force immediate flush
sys.stdout.reconfigure(line_buffering=True)
sys.stderr.reconfigure(line_buffering=True)

print("✅ DEBUG: Logging is now forced to flush immediately.")
print(f"DEBUG: PyTorch is using CUDA: {torch.cuda.is_available()}")
print(f"DEBUG: Available GPUs: {torch.cuda.device_count()}")
if torch.cuda.is_available():
    print(f"DEBUG: Current GPU: {torch.cuda.get_device_name(0)}")

AWS_ACCESS_KEY    = os.getenv('AWS_ACCESS_KEY_ID')
AWS_SECRET_KEY    = os.getenv('AWS_SECRET_ACCESS_KEY')
S3_BUCKET_NAME    = os.getenv('S3_BUCKET_NAME')
AWS_REGION        = os.getenv('AWS_REGION', 'us-east-1')
SQS_QUEUE_URL     = os.getenv("SQS_QUEUE_URL")

print("DEBUG: Checking environment variables:")
print(f"  - AWS_ACCESS_KEY_ID is {'SET' if AWS_ACCESS_KEY else 'NOT SET'}")
print(f"  - AWS_SECRET_ACCESS_KEY is {'SET' if AWS_SECRET_KEY else 'NOT SET'}")
print(f"  - S3_BUCKET_NAME = {S3_BUCKET_NAME}")
print(f"  - AWS_REGION     = {AWS_REGION}")
print(f"  - SQS_QUEUE_URL  = {SQS_QUEUE_URL}")

try:
    s3_client = boto3.client(
        's3',
        aws_access_key_id=AWS_ACCESS_KEY,
        aws_secret_access_key=AWS_SECRET_KEY,
        region_name=AWS_REGION
    )
    print("DEBUG: Successfully created boto3 S3 client.")
except Exception as e:
    print(f"DEBUG: Error creating boto3 S3 client: {e}")
    s3_client = None

try:
    sqs_client = boto3.client(
        'sqs',
        aws_access_key_id=AWS_ACCESS_KEY,
        aws_secret_access_key=AWS_SECRET_KEY,
        region_name=AWS_REGION
    )
    print("DEBUG: Successfully created boto3 SQS client.")
except Exception as e:
    print(f"DEBUG: Error creating boto3 SQS client: {e}")
    sqs_client = None


def run_ffmpeg_with_fallback(input_path, output_path, resize_filter, fps, codec_name):
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required for this operation")
    if not os.path.exists(input_path):
        raise RuntimeError(f"Input file does not exist: {input_path}")

    clamp = '-vf "scale=min(2032\\,iw):min(2032\\,ih)"'
    nvenc_cmd = (
        f'ffmpeg -y -i "{input_path}" {clamp} '
        f'-pix_fmt yuv420p -color_range tv '
        f'-color_primaries bt709 -color_trc bt709 -colorspace bt709 '
        f'-c:v h264_nvenc -preset p7 -tune hq -rc:v vbr_hq '
        f'-cq:v 19 -qmin:v 19 -qmax:v 21 -b:v 0 '
        f'-maxrate:v 130M -bufsize:v 130M '
        f'-profile:v high -r {fps} -c:a aac -b:a 192k '
        f'"{output_path}"'
    )
    print("Running GPU encoding:\n", nvenc_cmd)
    res = subprocess.run(nvenc_cmd, shell=True,
                         stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    print("NVENC STDERR:", res.stderr)
    if res.returncode != 0:
        raise RuntimeError(f"GPU encoding failed (exit {res.returncode}): {res.stderr.strip()}")
    if not os.path.exists(output_path) or os.path.getsize(output_path) == 0:
        raise RuntimeError("GPU encoding produced no output or output is empty")

    fixed = output_path.replace(".mp4", "_fixed.mp4")
    bsf_cmd = (
        f'ffmpeg -y -i "{output_path}" -c copy '
        f'-bsf:v h264_metadata=video_full_range_flag=0:'
        f'colour_primaries=1:transfer_characteristics=1:matrix_coefficients=1 '
        f'"{fixed}"'
    )
    print("Patching VUI with bitstream filter:\n", bsf_cmd)
    res2 = subprocess.run(bsf_cmd, shell=True,
                          stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    print("BSF STDERR:", res2.stderr)
    if res2.returncode != 0:
        raise RuntimeError(f"Failed to patch metadata (exit {res2.returncode}): {res2.stderr.strip()}")

    os.replace(fixed, output_path)
    print("✅ GPU encoding succeeded and Rec.709 VUI injected")


def get_video_properties(video_path):
    cmd = [
        'ffprobe', '-v', 'error', '-select_streams', 'v:0',
        '-show_entries', 'stream=codec_name,width,height,r_frame_rate,bit_rate',
        '-of', 'json', video_path
    ]
    result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    info = json.loads(result.stdout)
    stream = info['streams'][0]
    codec_name = stream['codec_name']
    width      = stream['width']
    height     = stream['height']
    # avoid eval; parse fractional fps safely
    r = stream['r_frame_rate']
    if isinstance(r, str) and '/' in r:
        num, den = r.split('/')
        num = int(num) if num else 0
        den = int(den) if den else 1
        fps = num / den if den else float(num)
    else:
        fps = float(r)
    bit_rate   = stream.get('bit_rate', None)
    return codec_name, fps, width, height, int(bit_rate) if bit_rate else None


def upload_to_s3(file_path, s3_key, extra_metadata=None):
    try:
        metadata = {str(k).replace(" ", "-").replace("_", "-").lower(): str(v)
                    for k, v in (extra_metadata or {}).items()} or None
        extra_args = {"Metadata": metadata} if metadata else {}
        print(f"DEBUG: Uploading '{file_path}' to S3 as '{s3_key}'...")
        s3_client.upload_file(file_path, S3_BUCKET_NAME, s3_key, ExtraArgs=extra_args)
        url = f"https://{S3_BUCKET_NAME}.s3.{AWS_REGION}.amazonaws.com/{s3_key}"
        print("DEBUG: S3 upload successful.")
        return url
    except Exception as e:
        print(f"ERROR: S3 upload failed: {e}")
        traceback.print_exc()
        return None


def s3_upload_with_timeout(file_path, s3_key, extra_metadata, timeout=60):
    with ThreadPoolExecutor(max_workers=1) as executor:
        future = executor.submit(upload_to_s3, file_path, s3_key, extra_metadata)
        try:
            return future.result(timeout=timeout)
        except Exception as e:
            print("ERROR: S3 upload timed out:", e)
            return None


def load_model():
    print("DEBUG: Loading YOLO model...")
    return YOLO("bests.pt").to("cuda" if torch.cuda.is_available() else "cpu")


def download_file_from_url(url, local_path):
    try:
        print(f"DEBUG: Downloading from URL: {url}")
        r = requests.get(url, stream=True)
        r.raise_for_status()
        with open(local_path, 'wb') as f:
            for chunk in r.iter_content(chunk_size=8192):
                f.write(chunk)
        print(f"✅ DEBUG: Download complete: {local_path}")
        return True
    except Exception as e:
        print(f"❌ ERROR: Failed to download: {e}")
        return False


def detect_ppe_video(video_file, original_key):
    """
    Reads frames, tracks with YOLO, dedupes static objects, draws annotations,
    pipes through FFmpeg, uploads video + JSON, and returns URLs + summary.
    """
    if not video_file:
        return None, "No video file provided."
    if "annotated" in os.path.basename(original_key).lower():
        return None, "Already annotated"
    if not torch.cuda.is_available():
        return None, "CUDA GPU is required"

    video_path = video_file.name
    model      = load_model()
    names      = model.names  # e.g., {0: 'Rebars', 1: 'Graafwerken'} — depends on your model
    start_time = time.time()

    cap    = cv2.VideoCapture(video_path)
    fps_cv = cap.get(cv2.CAP_PROP_FPS)
    w_cv   = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    h_cv   = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    # Prepare output paths
    parts = original_key.split('/')
    parts[0] = "progress"
    if "default" in parts:
        parts[parts.index("default")] = "annotated"
    parts[-1] = "annotated.mp4"
    processed_s3_key  = "/".join(parts)
    final_output_path = "/tmp/annotated.mp4"

    # Colors & counters (focus on Graafwerken / Rebars, fallback color for anything else)
    preferred_colors = {
        'Graafwerken': (0, 255, 255),
        'Rebars':      (0, 255, 0),
    }

    def color_for(label: str):
        return preferred_colors.get(
            label,
            (random.randint(20, 255), random.randint(20, 255), random.randint(20, 255))
        )

    counts      = {k: 0 for k in preferred_colors}  # only count the two classes of interest
    annotations = []
    seen        = set()

    with tempfile.TemporaryDirectory() as td:
        pipe_path = os.path.join(td, "pipe.mp4")

        # Start FFmpeg pipe for intermediate
        vf_chain = "format=yuv420p,scale=in_range=full:out_range=limited"
        pipe_cmd = (
            f'ffmpeg -y -f rawvideo -pix_fmt bgr24 -s {w_cv}x{h_cv} '
            f'-r {fps_cv} -i pipe: -vf "{vf_chain}" '
            f'-c:v libx264 -preset ultrafast -pix_fmt yuv420p "{pipe_path}"'
        )
        proc = subprocess.Popen(pipe_cmd, stdin=subprocess.PIPE,
                                stdout=subprocess.DEVNULL, stderr=subprocess.PIPE,
                                shell=True, bufsize=10**8)
        threading.Thread(
            target=lambda p: [print(l.decode().rstrip()) for l in iter(p.stderr.readline, b'')],
            args=(proc,), daemon=True
        ).start()

        def process_batch(frames, start_frame_index):
            # Run tracker for this batch
            results = model.track(frames, conf=0.4, iou=0.3, persist=True)
            for i, out_frame in enumerate(frames):
                frame_no = start_frame_index + i
                if results is None or len(results) <= i or results[i] is None or results[i].boxes is None:
                    proc.stdin.write(out_frame.tobytes())
                    continue

                boxes = results[i].boxes
                xyxy = boxes.xyxy.cpu().numpy()
                cls  = boxes.cls.cpu().numpy().astype(int)

                ids_t = getattr(boxes, "id", None)
                if ids_t is None:
                    tids = [None] * len(xyxy)
                else:
                    tids = ids_t.cpu().numpy().astype(int).tolist()

                for (box, cls_idx, tid) in zip(xyxy, cls, tids):
                    x1, y1, x2, y2 = map(int, box)
                    lbl = names.get(cls_idx, str(cls_idx))
                    cx, cy = (x1 + x2) / 2, (y1 + y2) / 2

                    # dedupe: use track ID if present, else rounded center
                    key = (lbl, tid if tid is not None else (round(cx), round(cy)))
                    if key not in seen:
                        seen.add(key)
                        annotations.append({
                            "frame":  frame_no,
                            "label":  lbl,
                            "id":     tid,
                            "bbox":   [x1, y1, x2, y2],
                            "center": [cx, cy]
                        })
                        if lbl in counts:
                            counts[lbl] += 1

                    col = color_for(lbl)
                    cv2.rectangle(out_frame, (x1, y1), (x2, y2), col, 2)
                    cv2.putText(out_frame, lbl, (x1, y1 - 5),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2)

                proc.stdin.write(out_frame.tobytes())

        # Read frames and process in batches, including the tail
        batch, frame_count = [], 0
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            frame_count += 1
            batch.append(frame)
            if len(batch) == 5:
                process_batch(batch, start_frame_index=frame_count - len(batch) + 1)
                batch.clear()

        if batch:
            process_batch(batch, start_frame_index=frame_count - len(batch) + 1)

        proc.stdin.close()
        proc.wait(timeout=30)
        cap.release()

        # final GPU pass + Rec.709 injection (left as in your original)
        codec_name, fps, w, h, _ = get_video_properties(video_path)
        rf = ''
        if w > 2032 or h > 2032:
            rf = '-vf scale=min(2032\\,iw):min(2032\\,ih)'
        try:
            run_ffmpeg_with_fallback(pipe_path, final_output_path, rf, fps, codec_name)
        except Exception as e:
            traceback.print_exc()
            return None, f"FFmpeg processing failed: {e}"

        # dump JSON and upload
        json_path = os.path.join(td, "annotations.json")
        with open(json_path, "w") as jf:
            json.dump({"counts": counts, "annotations": annotations}, jf, indent=2)
        json_s3_key = processed_s3_key.rsplit(".", 1)[0] + ".json"
        s3_json_url = s3_upload_with_timeout(json_path, json_s3_key, None)

    # upload annotated video
    s3_url = s3_upload_with_timeout(final_output_path, processed_s3_key, counts)
    if not s3_url:
        return None, "Upload failed"

    summary = "\n".join(f"{c}: {n}" for c, n in counts.items() if n > 0) or "No detections"
    print(f"✅ Uploaded video: {s3_url}")
    print(f"✅ Uploaded JSON:  {s3_json_url}")
    print(f"Processing completed in {time.time()-start_time:.2f}s")

    return {"video_url": s3_url, "json_url": s3_json_url, "summary": summary}


def submit_job(video_file, original_key):
    with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(video_file.name)[1]) as tmp:
        tmp.write(video_file.read())
        tmp_path = tmp.name

    s3_input_key = f"input/{os.path.basename(tmp_path)}"
    s3_url = upload_to_s3(tmp_path, s3_input_key)
    if not s3_url:
        return "S3 upload failed."

    job_message = {"bucket": S3_BUCKET_NAME, "object_key": s3_input_key, "presigned_url": s3_url}
    response = sqs_client.send_message(QueueUrl=SQS_QUEUE_URL, MessageBody=json.dumps(job_message))
    print("DEBUG: Sent job to SQS:", response)
    return "Job submitted successfully!"


def sqs_worker():
    while True:
        resp = sqs_client.receive_message(
            QueueUrl=SQS_QUEUE_URL, MaxNumberOfMessages=10, WaitTimeSeconds=20
        )
        messages = resp.get('Messages', [])
        if not messages:
            print("DEBUG: No messages in queue right now.")
        else:
            for message in messages:
                try:
                    job = json.loads(message['Body'])
                    print("DEBUG: Received job payload:", job)
                    bucket = job.get("bucket", S3_BUCKET_NAME)
                    key    = job.get("object_key") or (job.get("presigned_url") or "").split("amazonaws.com/")[-1].split("?")[0]
                    if not key:
                        print("ERROR: missing object_key, skipping.")
                        continue
                    local_path = f"/tmp/{os.path.basename(key)}"
                    if download_file_from_url(job.get("presigned_url"), local_path):
                        with open(local_path, "rb") as f:
                            result = detect_ppe_video(f, key)
                        print("DEBUG: detect_ppe_video result:", result)
                    else:
                        print("ERROR: failed to download", job.get("presigned_url"))
                except Exception as e:
                    print("❌ ERROR processing job:", e)
                    traceback.print_exc()
                finally:
                    sqs_client.delete_message(
                        QueueUrl=SQS_QUEUE_URL, ReceiptHandle=message['ReceiptHandle']
                    )
        time.sleep(10)


with gr.Blocks() as demo:
    gr.Markdown("# YOLOv8 PPE Detection with S3 Metadata")
    video_input        = gr.File(label="Upload Video", file_types=[".mp4", ".avi", ".mov"])
    original_key_input = gr.Textbox(label="S3 Key (original path)")
    detect_button      = gr.Button("Submit Job (Fire-and-Forget)")
    job_id_output      = gr.Textbox(label="Job ID")
    detect_button.click(fn=submit_job, inputs=[	video_input, original_key_input], outputs=[job_id_output])

if __name__ == "__main__":
    threading.Thread(target=sqs_worker, daemon=True).start()
    demo.launch(server_name="0.0.0.0", server_port=7860)