Skip to content

Commit

Permalink
Update main.py
Browse files Browse the repository at this point in the history 
Signed-off-by: Bubbles The Dev <152947339+KernFerm@users.noreply.github.com>
  • Loading branch information
@KernFerm
KernFerm authored Sep 4, 2024
1 parent 55aa020 commit bc81bc3
Showing 1 changed file with 27 additions and 19 deletions.
46 changes: 27 additions & 19 deletions main.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import time
import mss # MSS for fast screen capturing
import os # For handling OS-specific paths
from colorama import Fore, Style, init # For colored text output

def load_model(model_path='ultralytics/yolov5s'):
"""
Expand All @@ -16,22 +17,22 @@ def load_model(model_path='ultralytics/yolov5s'):
try:
start_time = time.time() # Start timing
if model_path.endswith('.pt'):
print("Loading PyTorch model...")
print(Fore.WHITE + "Loading PyTorch model...")
model = torch.hub.load('ultralytics/yolov5', 'custom', path=model_path, force_reload=True) # Custom model
model_type = 'torch' # YOLOv5 model
elif model_path.endswith('.onnx'): # ONNX model
print("Loading ONNX model...")
print(Fore.WHITE + "Loading ONNX model...")
model = ort.InferenceSession(model_path) # ONNX model
model_type = 'onnx'
else:
print("Loading YOLOv5 model from repository...")
print(Fore.WHITE + "Loading YOLOv5 model from repository...")
model = torch.hub.load('ultralytics/yolov5', model_path, force_reload=True) # YOLOv5 model from the repository
model_type = 'torch' # YOLOv5 model
end_time = time.time() # End timing
print(f"Model loaded in {end_time - start_time:.2f} seconds") # Print the loading time
print(Fore.WHITE + f"Model loaded in {end_time - start_time:.2f} seconds") # Print the loading time
return model, model_type # Return the loaded model and model type
except Exception as e: # Catch any exceptions during model loading
print(f"Error loading model: {e}") # Print the error message
print(Fore.WHITE + f"Error loading model: {e}") # Print the error message
sys.exit(1)

def capture_screen():
Expand Down Expand Up @@ -68,7 +69,7 @@ def detect_objects(model, model_type, frame, device):
results = outputs[0]
return results
except Exception as e:
print(f"Error during detection: {e}")
print(Fore.WHITE + f"Error during detection: {e}")
return None

def draw_bounding_boxes(frame, results, color, model_type):
Expand All @@ -90,7 +91,7 @@ def draw_bounding_boxes(frame, results, color, model_type):
xmin, ymin, xmax, ymax, conf = results[i][:5]
cv2.rectangle(frame, (int(xmin), int(ymin)), (int(xmax), int(ymax)), color, 2)
else:
print(f"Skipping result {i} due to insufficient values: {results[i]}")
print(Fore.WHITE + f"Skipping result {i} due to insufficient values: {results[i]}")
return frame

def get_color_from_input():
Expand Down Expand Up @@ -125,52 +126,59 @@ def get_color_from_input():
print("Invalid color name. Please try again.")

def main():
# Initialize colorama
init(autoreset=True)

# Ensure the game is running and visible on the screen before starting the script
input("Make sure the game is running and visible on the screen. Press Enter to continue...")
input(Fore.MAGENTA + "Make sure the game is running and visible on the screen. Press Enter to continue...")

print(Fore.RED + "Any issues Join our Discord Server: https://discord.fnbubbles420.org/invite head to gaming-vision-aid channel for help. ping @game-vision-aid-devs for help.")
print(Fore.GREEN + "The program was created & developed by Bubbles The Dev for FNBUBBLES420 ORG. for those with visual impairments and color blindness")
print(Fore.BLUE + "This is a Accessible Gaming Vision Aid Tool for the visually impaired and color blind.")

# Check for NVIDIA GPU availability with CUDA support
if torch.cuda.is_available():
device = 'cuda'
print("CUDA-enabled NVIDIA GPU found. Using GPU acceleration.")
print(Fore.CYAN + "CUDA-enabled NVIDIA GPU found. Using GPU acceleration.")
else:
device = 'cpu'
print("No CUDA-enabled NVIDIA GPU found. Using CPU.")
print(Fore.YELLOW + "No CUDA-enabled NVIDIA GPU found. Using CPU.")

# Load YOLOv5 or ONNX model
# Use relative paths to ensure cross-platform compatibility.
# Assume the model file is in the same directory as the script.
script_dir = os.path.dirname(os.path.abspath(__file__))
model_path = os.path.join(script_dir, 'models\FN_v5.pt') # Change to your custom model path if needed
model_path = os.path.join(script_dir, 'models', 'FN_v5.pt') # Change to your custom model path if needed

start_time = time.time() # Start timing for model loading
model, model_type = load_model(model_path)
if model_type == 'torch':
model = model.to(device)
end_time = time.time() # End timing for model loading
print(f"Model loading time: {end_time - start_time:.2f} seconds")
print(Fore.WHITE + f"Model loading time: {end_time - start_time:.2f} seconds")

# Get a color for the overlay from user input
print("Getting overlay color from user input...")
print(Fore.BLUE +"Getting overlay color from user input...")
start_time = time.time() # Start timing for color input
overlay_color = get_color_from_input()
end_time = time.time() # End timing for color input
print(f"Color input time: {end_time - start_time:.2f} seconds")
print("Overlay color selected.")
print(Fore.YELLOW + f"Color input time: {end_time - start_time:.2f} seconds")
print(Fore.WHITE + "Overlay color selected.")

# Main loop
try:
while True:
# Capture the screen
frame = capture_screen()
print("Screen captured.")
print(Fore.RED + "Screen captured.")

# Detect objects in the frame
results = detect_objects(model, model_type, frame, device)
print("Objects detected.")
print(Fore.GREEN + "Objects detected.")

# Draw bounding boxes around detected objects with the chosen color
frame_with_boxes = draw_bounding_boxes(frame, results, overlay_color, model_type)
print("Bounding boxes drawn.")
print(Fore.BLUE + "Bounding boxes drawn.")

# Display the frame with the overlay
cv2.imshow("Overlay", frame_with_boxes)
Expand All @@ -180,7 +188,7 @@ def main():
break

except KeyboardInterrupt:
print("Exiting program...")
print(Fore.YELLOW + "Exiting program...")

finally:
# Clean up
Expand Down

0 comments on commit bc81bc3

Please sign in to comment.