Add files via upload

KernFerm · web-flow · commit 48fbecb56c7f · 2025-04-26T06:44:34.000-04:00
Signed-off-by: Bubbles The Dev &lt;152947339+KernFerm@users.noreply.github.com&gt;
diff --git a/Run-Application/config-launcher.bat b/Run-Application/config-launcher.bat
@@ -0,0 +1,18 @@
+@echo off
+REM Directly set the file path to config.py
+set scriptPath="config.py"
+
+REM Verify the entered file path exists
+if not exist %scriptPath% (
+    echo Error: The specified file does not exist. Please check the path.
+    pause >nul
+    exit /b 1
+)
+
+REM Open the config.py file in Notepad
+echo Opening config.py in Notepad...
+notepad %scriptPath%
+
+REM Keep the window open
+echo Press any key to exit...
+pause >nul
diff --git a/Run-Application/config.py b/Run-Application/config.py
@@ -0,0 +1,37 @@
+# Configuration for BetterCam Screen Capture and YOLO model
+
+# 🔍 Screen Capture Settings
+screenWidth = 640     # Width of the screen region to capture
+screenHeight = 640    # Height of the screen region to capture
+
+# 🎯 Object Detection Settings
+confidenceThreshold = 0.5  # Minimum confidence to draw detection
+nmsThreshold = 0.4         # Non-max suppression threshold
+
+# 📦 YOLO Model Selection
+# Options: 'torch' (YOLOv5/v8 .pt), 'onnx' (.onnx), or 'engine' (.engine for TensorRT)
+modelType = 'torch'
+
+# 🧠 Model File Paths
+torchModelPath = 'models/yolov8n.pt'           # PyTorch model (YOLOv5/YOLOv8)
+onnxModelPath = 'models/yolov8n.onnx'          # ONNX model path
+tensorrtModelPath = 'models/yolov8n.engine'    # TensorRT engine file
+
+# 🎥 BetterCam Capture Settings
+targetFPS = 60            # Target frames per second
+maxBufferLen = 512        # Max buffer size for frame storage
+region = None             # Set to None for full-screen capture
+monitorIdx = 0            # Use 0 for primary monitor
+
+# 🖍️ Drawing Bounding Boxes
+boundingBoxColor = (0, 255, 0)  # BGR Green
+highlightColor = (0, 0, 255)    # BGR Red for highlights
+
+# 🪟 Overlay Settings
+overlayWidth = 1920
+overlayHeight = 1080
+overlayAlpha = 200  # 0 (fully transparent) to 255 (fully opaque)
+
+# ⚡ GPU Backend Preferences
+useCuda = True       # Enable CUDA (for NVIDIA GPUs)
+useDirectML = True   # Enable DirectML (for AMD GPUs or fallback)
diff --git a/Run-Application/launcher.bat b/Run-Application/launcher.bat
@@ -0,0 +1,28 @@
+@echo off
+REM Save the current directory
+pushd %~dp0
+
+REM Check if main_cpu.py exists
+if not exist main_cpu.py (
+    echo Error: main_cpu.py not found in the current directory.
+    popd
+    pause
+    exit /b 1
+)
+
+REM Run the Python script and check for errors
+echo Running main_cpu.py...
+python main_cpu.py
+if %errorlevel% neq 0 (
+    echo Error: main_cpu.py did not run successfully. Error level: %errorlevel%
+    popd
+    pause
+    exit /b %errorlevel%
+)
+
+REM Provide success feedback
+echo main_cpu.py ran successfully.
+
+REM Return to the original directory
+popd
+pause
diff --git a/Run-Application/main.py b/Run-Application/main.py
@@ -0,0 +1,183 @@
+import sys
+import os
+import time
+import numpy as np
+import cupy as cp
+import cv2
+import torch
+import torchvision
+import torchaudio
+import torch_directml
+import onnx
+import onnxsim
+import onnxruntime as ort
+import onnxruntime_directml
+import bettercam
+import config
+import customtkinter as ctk
+import win32api
+import win32con
+import win32gui
+from ultralytics import YOLO
+from colorama import Fore, Style, init
+from overlay import Overlay
+import tensorrt as trt
+
+# ------------------ TensorRT Inference Class ------------------ #
+class TensorRTInference:
+    def __init__(self, engine_path):
+        self.logger = trt.Logger(trt.Logger.WARNING)
+        self.runtime = trt.Runtime(self.logger)
+        with open(engine_path, 'rb') as f:
+            self.engine = self.runtime.deserialize_cuda_engine(f.read())
+        self.context = self.engine.create_execution_context()
+        self.input_shape = self.engine.get_binding_shape(0)
+        self.output_shape = self.engine.get_binding_shape(1)
+        self.input_size = np.prod(self.input_shape) * np.float32().nbytes
+        self.output_size = np.prod(self.output_shape) * np.float32().nbytes
+        self.d_input = cp.cuda.alloc(self.input_size)
+        self.d_output = cp.cuda.alloc(self.output_size)
+        self.bindings = [int(self.d_input.ptr), int(self.d_output.ptr)]
+
+    def infer(self, input_data: np.ndarray) -> np.ndarray:
+        input_data = input_data.astype(np.float32).ravel()
+        output_data = np.empty(self.output_shape, dtype=np.float32)
+        cp.cuda.runtime.memcpy(self.d_input.ptr, input_data.ctypes.data, self.input_size, cp.cuda.runtime.memcpyHostToDevice)
+        self.context.execute_v2(self.bindings)
+        cp.cuda.runtime.memcpy(output_data.ctypes.data, self.d_output.ptr, self.output_size, cp.cuda.runtime.memcpyDeviceToHost)
+        return output_data.reshape(self.output_shape)
+
+# ------------------ BetterCam Enhanced ------------------ #
+class BetterCamEnhanced:
+    def __init__(self, max_buffer_len=config.maxBufferLen, target_fps=config.targetFPS, region=None, monitor_idx=0):
+        self.camera = None
+        self.max_buffer_len = max_buffer_len
+        self.target_fps = target_fps
+        self.region = region
+        self.monitor_idx = monitor_idx
+        self.is_capturing = False
+
+    def start(self):
+        self.camera = bettercam.create(monitor_idx=self.monitor_idx, max_buffer_len=self.max_buffer_len)
+        self.camera.start(target_fps=self.target_fps)
+        self.is_capturing = True
+
+    def grab_frame(self):
+        return self.camera.grab(region=self.region) if self.region else self.camera.grab()
+
+    def stop(self):
+        self.camera.stop()
+        self.is_capturing = False
+
+# ------------------ Load Model ------------------ #
+def load_model():
+    model_type = config.modelType.lower()
+
+    if model_type == 'torch':
+        print(Fore.CYAN + "[INFO] Loading PyTorch model using Ultralytics...")
+        return YOLO(config.torchModelPath), 'torch'
+
+    elif model_type == 'onnx':
+        available_providers = ort.get_available_providers()
+
+        if 'CUDAExecutionProvider' in available_providers:
+            providers = ['CUDAExecutionProvider']
+            print(Fore.GREEN + "[INFO] ONNX model will use CUDAExecutionProvider (NVIDIA GPU).")
+        elif 'DmlExecutionProvider' in available_providers or 'DirectMLExecutionProvider' in available_providers:
+            providers = ['DmlExecutionProvider'] if 'DmlExecutionProvider' in available_providers else ['DirectMLExecutionProvider']
+            print(Fore.YELLOW + "[INFO] ONNX model will use DirectMLExecutionProvider (AMD GPU).")
+        else:
+            providers = ['CPUExecutionProvider']
+            print(Fore.RED + "[INFO] ONNX model will use CPUExecutionProvider (CPU only).")
+
+        session = ort.InferenceSession(config.onnxModelPath, providers=providers)
+        return session, 'onnx'
+
+    elif model_type == 'engine':
+        print(Fore.MAGENTA + "[INFO] Loading TensorRT engine with CuPy...")
+        return TensorRTInference(config.tensorrtModelPath), 'engine'
+
+    else:
+        raise ValueError(Fore.RED + "Unsupported modelType in config.py. Use 'torch', 'onnx', or 'engine'.")
+
+# ------------------ Object Detection ------------------ #
+def detect_objects(model, model_type, frame):
+    if model_type == 'torch':
+        return model.predict(source=frame, imgsz=(config.screenWidth, config.screenHeight), conf=config.confidenceThreshold, verbose=False)
+    elif model_type == 'onnx':
+        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        resized = cv2.resize(frame_rgb, (config.screenWidth, config.screenHeight)).astype(np.float32)
+        tensor = resized.transpose(2, 0, 1)[np.newaxis] / 255.0
+        return model.run(None, {model.get_inputs()[0].name: tensor})
+    elif model_type == 'engine':
+        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        resized = cv2.resize(frame_rgb, (config.screenWidth, config.screenHeight)).astype(np.float32)
+        tensor = resized.transpose(2, 0, 1)[np.newaxis] / 255.0
+        return model.infer(tensor)
+    return None
+
+# ------------------ Draw Bounding Boxes ------------------ #
+def draw_boxes(frame, results, model_type):
+    if results is None:
+        return frame
+    if model_type == 'torch':
+        for r in results:
+            for box in r.boxes.xyxy.cpu().numpy().astype(int):
+                x1, y1, x2, y2 = box[:4]
+                cv2.rectangle(frame, (x1, y1), (x2, y2), config.boundingBoxColor, 2)
+    elif model_type in ['onnx', 'engine']:
+        for det in results[0]:
+            if det[4] > config.confidenceThreshold:
+                x1, y1, x2, y2 = map(int, det[:4])
+                cv2.rectangle(frame, (x1, y1), (x2, y2), config.boundingBoxColor, 2)
+    return frame
+
+def extract_boxes(results, model_type):
+    boxes = []
+    if model_type == 'torch':
+        for r in results:
+            for box in r.boxes.xyxy.cpu().numpy().astype(int):
+                x1, y1, x2, y2 = box[:4]
+                boxes.append([x1, y1, x2, y2])
+    elif model_type in ['onnx', 'engine']:
+        for det in results[0]:
+            if det[4] > config.confidenceThreshold:
+                x1, y1, x2, y2 = map(int, det[:4])
+                boxes.append([x1, y1, x2, y2])
+    return boxes
+
+# ------------------ Main ------------------ #
+def main():
+    init(autoreset=True)
+    input("Start your game and press Enter to continue...")
+
+    camera = BetterCamEnhanced(target_fps=config.targetFPS, monitor_idx=config.monitorIdx)
+    camera.start()
+
+    model, model_type = load_model()
+    overlay = Overlay(width=config.overlayWidth, height=config.overlayHeight, alpha=config.overlayAlpha)
+    overlay.toggle()
+
+    try:
+        while True:
+            frame = camera.grab_frame()
+            if frame is None:
+                continue
+
+            results = detect_objects(model, model_type, frame)
+            frame = draw_boxes(frame, results, model_type)
+            boxes = extract_boxes(results, model_type)
+            overlay.update(boxes)
+
+            cv2.imshow("YOLO Detection", frame)
+            if cv2.waitKey(1) & 0xFF == ord('q'):
+                break
+    except KeyboardInterrupt:
+        pass
+    finally:
+        camera.stop()
+        overlay.toggle()
+        cv2.destroyAllWindows()
+
+if __name__ == "__main__":
+    main()
diff --git a/Run-Application/overlay.py b/Run-Application/overlay.py
@@ -0,0 +1,68 @@
+import win32gui
+import win32con
+import win32api
+
+class Overlay:
+    def __init__(self, width=1920, height=1080, alpha=200):
+        self.hInstance = win32api.GetModuleHandle()
+        self.className = "OverlayWindowClass"
+        self.hWnd = None
+        self.width = width
+        self.height = height
+        self.alpha = alpha
+        self.visible = False
+
+        self._register_class()
+        self._create_window()
+
+    def _register_class(self):
+        wndClass = win32gui.WNDCLASS()
+        wndClass.lpfnWndProc = self._wnd_proc
+        wndClass.hInstance = self.hInstance
+        wndClass.lpszClassName = self.className
+        wndClass.hCursor = win32gui.LoadCursor(None, win32con.IDC_ARROW)
+        wndClass.hbrBackground = win32con.COLOR_WINDOW
+        win32gui.RegisterClass(wndClass)
+
+    def _create_window(self):
+        screen_w = win32api.GetSystemMetrics(0)
+        screen_h = win32api.GetSystemMetrics(1)
+
+        self.hWnd = win32gui.CreateWindowEx(
+            win32con.WS_EX_LAYERED | win32con.WS_EX_TOPMOST | win32con.WS_EX_TOOLWINDOW | win32con.WS_EX_TRANSPARENT,
+            self.className,
+            None,
+            win32con.WS_POPUP,
+            0,
+            0,
+            screen_w,
+            screen_h,
+            None,
+            None,
+            self.hInstance,
+            None
+        )
+
+        win32gui.SetLayeredWindowAttributes(self.hWnd, 0x000000, self.alpha, win32con.LWA_ALPHA)
+        win32gui.ShowWindow(self.hWnd, win32con.SW_SHOW)
+
+    def toggle(self):
+        self.visible = not self.visible
+        win32gui.ShowWindow(self.hWnd, win32con.SW_SHOW if self.visible else win32con.SW_HIDE)
+
+    def update(self, bounding_boxes):
+        hdc = win32gui.GetDC(self.hWnd)
+        pen = win32gui.CreatePen(win32con.PS_SOLID, 2, win32api.RGB(0, 255, 0))
+        win32gui.SelectObject(hdc, pen)
+
+        for box in bounding_boxes:
+            x1, y1, x2, y2 = box
+            win32gui.Rectangle(hdc, x1, y1, x2, y2)
+
+        win32gui.DeleteObject(pen)
+        win32gui.ReleaseDC(self.hWnd, hdc)
+
+    def _wnd_proc(self, hWnd, msg, wParam, lParam):
+        if msg == win32con.WM_DESTROY:
+            win32gui.PostQuitMessage(0)
+        return win32gui.DefWindowProc(hWnd, msg, wParam, lParam)
