import os
# Optimize OpenMP and MKL thread allocation for AMD Ryzen 5 6600H (6 Cores)
os.environ["OMP_NUM_THREADS"] = "6"
os.environ["MKL_NUM_THREADS"] = "6"

import cv2
import torch
from ultralytics import YOLO

# 1. Load the PyTorch YOLO segmentation model
model_path = "best.pt"
model = YOLO(model_path)

# Optimize PyTorch CPU thread pools for 6 physical cores to avoid SMT hyperthreading overhead
torch.set_num_threads(6)
print("Thread PyTorch diset ke 6 (Physical Cores) untuk optimalisasi CPU AMD Ryzen 5.")

# Auto-detect device
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Device inferensi diset ke: {device}")

# 2. Open the video file
video_path = r"D:\Belajar\Menghitung karung\0727.mp4"
cap = cv2.VideoCapture(video_path)

if not cap.isOpened():
    print(f"Error: Gagal membuka video di {video_path}")
    exit(1)

# Optimasi 1: Frame Stride (Frame Skipping)
# FRAME_STRIDE = 3 artinya memproses 1 dari setiap 3 frame (sangat berguna untuk video 60fps agar CPU tidak overload)
FRAME_STRIDE = 3
frame_idx = 0

print("=== Simple Predict (Optimized for AMD Ryzen) Running ===")
print("Tekan 'q' di jendela video untuk keluar.\n")

annotated_frame = None

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        print("Video selesai diputar atau tidak terbaca.")
        break

    frame_idx += 1

    # Hanya jalankan deteksi model pada frame tertentu berdasarkan STRIDE
    if FRAME_STRIDE <= 1 or frame_idx % FRAME_STRIDE == 0 or annotated_frame is None:
        # Optimasi 2: perkecil resolusi inferensi imgsz=320 untuk kecepatan maksimal
        # Optimasi 3: gunakan device yang sesuai (cpu)
        results = model(frame, conf=0.15, classes=[0], imgsz=320, device=device, verbose=False)
        
        # Optimasi 4: Gambar hasil deteksi (diset masks=False untuk kecepatan menggambar di CPU)
        annotated_frame = results[0].plot(masks=False)

    # 5. Tampilkan frame di jendela
    resized_frame = cv2.resize(annotated_frame, (960, 540))
    cv2.imshow("YOLO Live Predict - Karung (Optimized)", resized_frame)

    # Keluar jika tombol 'q' ditekan
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# 6. Bersihkan resource
cap.release()
cv2.destroyAllWindows()
print("Proses selesai.")