Added YOLO model

.gitattributes

@@ -0,0 +1,2 @@
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,21 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+.Python
+env/
+build/
+develop-eggs/
+dist/
+eggs/
+*.egg-info/
+
+# Data
+data/
+
+# Models
+runs/
+yolov8n.pt
+
+# Visual Studio
+*.vs

models/best.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e5511f2d7d2c6d85e3011c133c3284025c3bfb4ab2e39c9ef70af1a0d0e8e7ea
+size 12274092

models/best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fbb4fcbd99e83e1fae188d27a443754981e271e1a0a8b3e4903b40f94014075b
+size 6244707

src/detectobjects.py

@@ -0,0 +1,79 @@
+from ultralytics import YOLO
+from ultralytics.utils.plotting import Annotator, colors
+import cv2
+from typing import List
+import numpy as np
+
+class ObjectDetector():
+    def __init__(self, pretrained_model: str = 'yolov8n.pt', debug: bool = False):
+        self.model = YOLO(pretrained_model)
+        self.debug = debug
+        
+        self.color_map = {
+            0: (0, 255, 0),     # player: green
+            1: (255, 128, 0),   # Storm Timer: light blue
+            2: (0, 0, 255),     # Killfeed: red
+            3: (255, 0, 0),     # Player Count: blue
+            4: (0, 255, 255),   # Minimap: yellow
+            5: (128, 0, 255),   # Storm Shrink Warning: dark red
+            6: (255, 0, 255),   # Eliminations: magenta
+            7: (0, 128, 255),   # Health: orange
+            8: (255, 255, 0),   # Shield: cyan
+            9: (128, 255, 0),   # Inventory: light green
+            10: (0, 165, 255),  # Buildings: orange-yellow
+            11: (139, 69, 19),  # Wood Material: brown
+            12: (128, 128, 128),# Brick Material: gray
+            13: (192, 192, 192),# Metal Material: light gray
+            14: (255, 191, 0),  # Compass: deep sky blue
+            15: (255, 0, 128),  # Equipped Item: purple
+            16: (0, 255, 191),  # Waypoint: yellow-green
+            17: (128, 128, 0),  # Sprint Meter: teal
+            18: (0, 140, 255),  # Safe Zone: orange-red
+            19: (0, 215, 255),  # playerIcon: gold
+            20: (34, 139, 34),  # Tree: forest green
+            21: (75, 75, 75),   # Stone: dark gray
+            22: (0, 69, 255),   # Building: orange-red
+            23: (122, 61, 0),   # Wood Building: dark brown
+            24: (108, 108, 108),# Stone Building: medium gray
+            25: (211, 211, 211), # Metal Building: silver
+            26: (0, 43, 27), # Wall: dark green
+            27: (22, 22, 22), # Ramp: dark gray
+            28: (17, 211, 0), # Pyramid: bright green
+            29: (121, 132, 9) # Floor: olive green
+}
+        
+    def train_model(self, yaml_filepath):
+        self.model.train(data=yaml_filepath, epochs=100, imgsz=640, batch=16, patience=50)
+        
+    def detect_object(self, frames: List[np.ndarray]):
+        for frame in frames:
+            results = self.model.track(frame, stream=True)
+            
+            for result in results:
+                class_names = result.names
+                annotated_frame = frame.copy()
+                
+                for box in result.boxes:
+                    if box.conf[0] > 0.4:
+                        [x1, y1, x2, y2] = box.xyxy[0] # coords
+                        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
+                        
+                        cls = int(box.cls[0]) # class
+                        color = self.color_map.get(cls, (0,255,0))
+                        
+                        cv2.rectangle(annotated_frame, (x1,y1), (x2,y2), color=color, thickness=2)
+
+                        text = f'{class_names[cls]} {box.conf[0]:.2f}'
+                        (text_width, text_height), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 1, 2)
+                        cv2.rectangle(annotated_frame, (x1, y1-text_height-5), (x1+text_width, y1), color, -1)
+
+                        cv2.putText(annotated_frame, text, (x1, y1-5), 
+                                  cv2.FONT_HERSHEY_SIMPLEX, 1, (255,255,255), thickness=2)
+                while self.debug:
+                    cv2.imshow('frame', annotated_frame)
+                    if cv2.waitKey(1) & 0xFF == ord('q'):
+                        break
+        cv2.destroyAllWindows()
+        
+    def export_model(self, format: str = 'onnx'):
+        self.model.export(format=format)

src/vidprocessing.py

@@ -0,0 +1,70 @@
+import cv2
+from typing import List
+from os.path import dirname, abspath
+from pathlib import Path
+import numpy as np
+from detectobjects import ObjectDetector
+import yaml
+
+base_dir = Path(dirname(dirname(abspath(__file__))))
+
+def get_video(path: str) -> cv2.VideoCapture:
+    video = cv2.VideoCapture(path)
+    if not video.isOpened():
+        raise ValueError(f'Could not open video file: {path}')
+    
+    fps = video.get(cv2.CAP_PROP_FPS)
+    frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
+    duration = frame_count/fps
+    
+    print(f'FPS: {fps}\nFrame Count: {frame_count}\nDuration: {duration}')
+    
+    return video
+
+def get_frames(video: cv2.VideoCapture, frame_start: int, frame_end: int) -> List[np.ndarray]:
+    
+    frames = []
+    for i in range(frame_start, frame_end+1):
+        video.set(cv2.CAP_PROP_POS_FRAMES, i)
+        ret, frame = video.read()
+        if not ret:
+            raise ValueError(f'Could not read frame {i}')
+        
+        frames.append(frame)
+        
+    return frames
+
+
+def create_images_of_video(video: cv2.VideoCapture, interval: int = 100):
+    frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
+    
+    # read frame_count / interval of frames
+    for i in range(0, frame_count, interval):
+        video.set(cv2.CAP_PROP_POS_FRAMES, i)
+        success, frame = video.read()
+        if not success:
+            raise ValueError(f'Could not read frame {i}')
+        
+        # write photo to file
+        cv2.imwrite(base_dir / "data" / 'model_data' / 'temp_vid_folder' /f'{i}.png', frame)
+                    
+
+if __name__ == "__main__":
+    vid = get_video(base_dir / "data" / "video_data" / "fortnite_remo_three.mp4")
+    frames = get_frames(vid, 21100, 21110)
+    
+    ### PRETRAINED MODEL DETECTION CODE
+    yolo = ObjectDetector(pretrained_model=(base_dir / 'best.pt'), debug=True)
+    yolo.detect_object(frames)
+
+
+    ### TRAINING CODE FOR YOLO MODEL
+    # yolo = ObjectDetector()
+    # yolo.train_model(base_dir / 'data' / 'model_data' / 'fortnite_train.yaml')
+
+    ### VIDEO FRAME CUTTING CODE
+    # create_images_of_video(vid)
+
+    ### MODEL EXPORT + VIDEO RELEASE
+    yolo.export_model()
+    vid.release()