Write lots of code

.gitignore

@@ -1,3 +1,6 @@
+*.mp4
+.idea/
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

downloader.py

@@ -0,0 +1,55 @@
+import os
+import sys
+import json
+
+import requests
+from pytubefix import YouTube
+from pytubefix.cli import on_progress
+
+TBAip = "https://www.thebluealliance.com/api/v3"
+# Free TBA Key!
+headers = {"User-Agent": "Mozilla/5.0", "X-TBA-Auth-Key": "fzQY0pv6qwfwuII5Xx2bmP57BBSuE0maxKailYlrI0e1EdfKCq6F3Th9FFDqpW7f"}
+RES = '1080p'
+
+if len(sys.argv) < 2:
+  sys.exit("Usage: downloader.py <event code>")
+
+evcode = sys.argv[1]
+
+event_matches = requests.get(TBAip + "/event/" + evcode + "/matches", headers=headers).json()
+
+video_urls = [None for _ in range(len(event_matches))]
+
+for match in event_matches:
+  if match["comp_level"] != "qm": continue
+  match_urls = []
+  for video_url in match["videos"]:
+    if video_url["type"] != "youtube": continue
+    match_urls.append(video_url["key"])
+  if len(match_urls) != 0:
+    video_urls[match["match_number"]] = match_urls
+
+if not os.path.exists(evcode):
+    os.makedirs(evcode)
+
+for i, match in enumerate(video_urls):
+  if match is None: continue
+  # print(match)
+  for url in match:
+    if os.path.exists(os.path.join(evcode, str(i)+".mp4")):
+      continue
+    yt = YouTube("https://youtube.com/watch?v="+url,  on_progress_callback=on_progress)
+    if yt.author != "FIRSTRoboticsCompetition": continue
+
+
+    for idx, stream in enumerate(yt.streams):
+      if stream.resolution == RES:
+        break
+
+    print(f"Downloading match {i}, {stream.resolution}")
+
+    # print(yt.streams[idx])
+    yt.streams[idx].download(output_path = evcode, filename = str(i)+".mp4")
+
+    # ys = yt.streams.get_highest_resolution()
+    # print(yt.title)

manualPath.py

@@ -0,0 +1,250 @@
+import math
+from time import time
+
+import cv2
+import sys
+
+import numpy as np
+from cv2.typing import Point, Scalar
+
+from src.videoCrop import contains_start, contains_end
+import src.imageTools as imageTools
+
+redColor = (120, 70, 238)
+colorRange = 60
+
+if len(sys.argv) < 3:
+  print("Usage: playVideo.py <event code> <match> <undistortion config>")
+  sys.exit(1)
+
+class distorton_config:
+  def __init__(self):
+    self.k_value = 0
+    self.zoom_value = 1
+    self.dots = [
+      [0.,0.],
+      [0.,0.],
+      [0.,0.],
+      [0.,0.]
+    ]
+
+dconf = distorton_config()
+conf_split = sys.argv[3].split(',')
+dconf.k_value = float(conf_split[0])
+dconf.zoom_value = float(conf_split[1])
+dconf.dots[0][0] = float(conf_split[2])
+dconf.dots[0][1] = float(conf_split[3])
+dconf.dots[1][0] = float(conf_split[4])
+dconf.dots[1][1] = float(conf_split[5])
+dconf.dots[2][0] = float(conf_split[6])
+dconf.dots[2][1] = float(conf_split[7])
+dconf.dots[3][0] = float(conf_split[8])
+dconf.dots[3][1] = float(conf_split[9])
+
+filename = sys.argv[1] + "/" + sys.argv[2] + ".mp4"
+
+blueColor = (280, 30, 125)
+blueColorUnselected = (int(blueColor[0]/3), int(blueColor[1]/3), int(blueColor[2]/3))
+redColor = (125, 30, 280)
+redColorUnselected = (int(redColor[0]/3), int(redColor[1]/3), int(redColor[2]/3))
+
+cap = cv2.VideoCapture(filename)
+totalFrames = cap.get(cv2.CAP_PROP_FRAME_COUNT)
+FPS = cap.get(cv2.CAP_PROP_FPS)
+width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
+
+
+selAlliance = "blue"
+selAllianceIndex = 1
+allianceIndex = [
+  "blue-1",
+  "blue-2",
+  "blue-3",
+  "red-1",
+  "red-2",
+  "red-3"
+]
+
+boxes = {
+  "blue-1": np.zeros((int(totalFrames), 4)),
+  "blue-2": np.zeros((int(totalFrames), 4)),
+  "blue-3": np.zeros((int(totalFrames), 4)),
+  "red-1": np.zeros((int(totalFrames), 4)),
+  "red-2": np.zeros((int(totalFrames), 4)),
+  "red-3": np.zeros((int(totalFrames), 4))
+}
+
+def undistort(frame):
+  display_frame = imageTools.radialUndistort(frame, dconf.k_value, dconf.zoom_value)
+  return imageTools.perspectiveUndistort(display_frame, dconf.dots)
+
+def dispColorCircle(frame, text, location, color):
+  cv2.circle(frame, location, 20, color, -1)
+  cv2.putText(frame, text, (location[0]-7, location[1]+8), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (255, 255, 255), 2)
+
+def distance(p1, p2):
+  return math.sqrt((p1[0]-p2[0])**2 + (p1[1]-p2[1])**2)
+
+clicked = False
+def click_event(event, x, y, flags, param):
+  global selAllianceIndex
+  global selAlliance
+
+  global clicked
+  global clickStart
+  global mousePos
+
+  if not isPaused: return
+  if event == cv2.EVENT_LBUTTONDOWN:
+    clicked = True
+    for i in range(1,4):
+      if distance((x, y), (int(width/2)+50*i,50)) < 20:
+        print("Click event detected")
+        selAllianceIndex = i
+        selAlliance = "red"
+        clicked = False
+        rerender()
+      elif distance((x, y), (int(width/2)+50*i-200,50)) < 20:
+        selAllianceIndex = i
+        selAlliance = "blue"
+        clicked = False
+        rerender()
+    if clicked:
+      clickStart = (x, y)
+  if clicked and event == cv2.EVENT_LBUTTONUP:
+    clicked = False
+    boxes[selAlliance+"-"+str(selAllianceIndex)][int(curFrame)] = np.array([
+      clickStart[0], clickStart[1],
+      x-clickStart[0], y-clickStart[1]
+    ])
+    restartTracking()
+    rerender()
+  if clicked and event == cv2.EVENT_MOUSEMOVE:
+    mousePos = (x,y)
+    rerender()
+
+
+cv2.imshow('frame', np.zeros((1,1)))
+cv2.setMouseCallback('frame', click_event)
+
+isPaused = True
+def rerender():
+  global frame
+
+  global clicked
+  global clickStart
+  global mousePos
+
+  framecopy = frame.copy()
+  cv2.rectangle(framecopy, (0,0),(500,100), (0,0,0), -1)
+  cv2.putText(framecopy, "Paused" if isPaused else "Unpaused", (0,20), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (255, 255, 255), 2)
+  cv2.putText(framecopy, "Frame: "+str(curFrame)+"/"+str(totalFrames), (0,40), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (255, 255, 255), 2)
+  cv2.putText(framecopy, "FPS: "+str(round(1.0 / (time() - start_time)))+"/"+str(round(FPS)), (0,60), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (255, 255, 255), 2)
+
+  for i in range(1,4):
+    dispColorCircle(framecopy, str(i), (int(width/2)+50*i,50), redColor if selAlliance == "red" and selAllianceIndex == i else redColorUnselected)
+    dispColorCircle(framecopy, str(i), (int(width/2)+50*i-200,50), blueColor if selAlliance == "blue" and selAllianceIndex == i else blueColorUnselected)
+
+  for alliance in allianceIndex:
+    if alliance.startswith("red"):
+      color = redColor
+    else:
+      color = blueColor
+
+    box = boxes[alliance][int(curFrame)]
+    if box is None: continue
+    if np.sum(box) == 0: continue
+
+    cv2.rectangle(framecopy, (int(box[0]), int(box[1])), (int(box[0]+box[2]), int(box[1]+box[3])), color, 3)
+    dispColorCircle(framecopy, alliance.split("-")[1], (int(box[0]+box[2]/2), int(box[1]+box[3]/2)), color)
+
+
+    # cv2.rectangle(framecopy, ())
+
+  if clicked:
+    if selAlliance.startswith("red"):
+      color = redColor
+    else:
+      color = blueColor
+    cv2.rectangle(framecopy, (clickStart[0], clickStart[1]), (mousePos[0], mousePos[1]), color, 3)
+
+
+  cv2.imshow('frame', framecopy)
+
+
+multiTracker = None
+trackingEnabled = [False, False, False, False, False, False]
+
+def restartTracking():
+  global multiTracker
+  global trackingEnabled
+  multiTracker = cv2.legacy.MultiTracker_create()
+  # Initialize MultiTracker
+  for i, alliance in enumerate(allianceIndex):
+    box = boxes[alliance][int(curFrame)]
+    if np.sum(box) == 0:
+      trackingEnabled[i] = False
+      continue
+    multiTracker.add(cv2.legacy.TrackerCSRT_create(), frame, box)
+    trackingEnabled[i] = True
+
+def track(frame):
+  global multiTracker
+  global trackingEnabled
+
+  if multiTracker is None:
+    restartTracking()
+
+
+  success, trackedBoxes = multiTracker.update(frame)
+  if not success: return
+  for i, box in enumerate(trackedBoxes):
+    boxes[allianceIndex[i]][int(curFrame)] = np.array(box)
+
+
+isInMatch = False
+while cap.isOpened():
+  global frame
+  global start_time
+  start_time = time()
+  ret, frame = cap.read()
+  if not ret:
+    break
+
+  if not isInMatch and contains_start(frame):
+    isInMatch = True
+  elif isInMatch and contains_end(frame):
+    break
+
+  # print(f"{i}, {isInMatch}")
+  curFrame = cap.get(cv2.CAP_PROP_POS_FRAMES)
+  if not isInMatch: continue
+  if curFrame % 5 != 0: continue
+
+  if not isPaused:
+    if cv2.waitKey(1) & 0xff == 32:
+      isPaused = not isPaused
+
+  frame = undistort(frame)
+  track(frame)
+  rerender()
+
+  while isPaused:
+    k = cv2.waitKey(0) & 0xff
+    if k == 81: # Left arrow
+      cap.set(cv2.CAP_PROP_POS_FRAMES, cap.get(cv2.CAP_PROP_POS_FRAMES) - 1 - 15)
+      curFrame = cap.get(cv2.CAP_PROP_POS_FRAMES)
+      restartTracking()
+      break
+    # if k == 83: # right arrow
+    #   cap.set(cv2.CAP_PROP_POS_FRAMES, cap.get(cv2.CAP_PROP_POS_FRAMES) - 1 + 15)
+    #   break
+    if k == 27:
+      cap.release()
+      break
+    if k == 32:
+      isPaused = not isPaused
+      break
+
+
+cap.release()

src/imageTools.py

@@ -0,0 +1,38 @@
+import numpy as np
+import cv2
+
+
+def perspectiveUndistort(image, dots):
+  height, width = image.shape[:2]
+
+  src_pts = np.array(dots, dtype=np.float32)
+  dst_pts = np.array([[0, 0], [width - 1, 0], [width - 1, height - 1], [0, height - 1]], dtype=np.float32)
+  M = cv2.getPerspectiveTransform(src_pts, dst_pts)
+  return cv2.warpPerspective(image, M, (width, height))
+
+def radialUndistort(original_image, k_value=0, zoom_value=1):
+  height, width = original_image.shape[:2]
+
+
+
+  # Radial distortion correction
+  center_x, center_y = width / 2, height / 2
+  x, y = np.meshgrid(np.arange(width), np.arange(height))
+  x = x.astype(np.float32) - center_x
+  y = y.astype(np.float32) - center_y
+  r = np.sqrt(x ** 2 + y ** 2)
+  r_max = np.max(r)
+  scale = float(zoom_value)
+  x_distorted = x * (1 + k_value * (r / r_max)) * scale
+  y_distorted = y * (1 + k_value * (r / r_max)) * scale
+  map_x = (x_distorted + center_x).astype(np.float32)
+  map_y = (y_distorted + center_y).astype(np.float32)
+  return cv2.remap(original_image, map_x, map_y, cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT)
+
+def getFrameFromVideo(vidpath, framenum):
+  cap = cv2.VideoCapture(vidpath)
+  cap.set(1,framenum)
+  ret, frame = cap.read()
+  if ret:
+    return frame
+  return None

src/videoCrop.py

@@ -0,0 +1,63 @@
+import cv2
+import numpy as np
+
+image_size = (640, 360)
+
+start_img = cv2.imread("./images/start.png")
+end_img = cv2.imread("./images/end.png")
+
+def containsImage(frame, image):
+  if frame is None: return False
+  frame = cv2.resize(frame, image_size, interpolation=cv2.INTER_AREA)
+  frame = frame[323:345, 305:335]
+  res = cv2.matchTemplate(frame, image, cv2.TM_CCOEFF_NORMED)
+  found = False
+  for _ in zip(*(np.where(res >= .98))[::-1]):
+    found = True
+    break
+  return found
+
+def contains_start(frame):
+  return containsImage(frame, start_img)
+
+def contains_end(frame):
+  return containsImage(frame, end_img)
+
+def search_video(video_path, search_img, start_index = 0):
+  cap = cv2.VideoCapture(video_path)
+
+  if not cap.isOpened():
+    print("Cannot open camera")
+    exit()
+
+  i = -1
+  while True:
+    # Capture frame-by-frame
+    ret, frame = cap.read()
+    if not ret:
+      print("Can't receive frame (stream end?). Exiting ...")
+      break
+    i += 1
+
+    if i < start_index: continue
+
+    frame = cv2.resize(frame, image_size, interpolation = cv2.INTER_AREA)
+    crop_img = frame[323:345, 305:335]
+
+    res = cv2.matchTemplate(crop_img, search_img, cv2.TM_CCOEFF_NORMED)
+    found = False
+    for _ in zip(*(np.where(res >= .98))[::-1]):
+      found = True
+      break
+    if found:
+      # cv2.waitKey(0)
+      break
+
+  cap.release()
+  return i
+
+
+def videoCrop(video_path):
+  start = search_video(video_path, start_img)
+  end = search_video(video_path, end_img, start)
+  return start, end

undistort.py

@@ -0,0 +1,160 @@
+import os
+import sys
+
+import cv2
+import numpy as np
+import tkinter as tk
+from tkinter import filedialog, messagebox
+from PIL import Image, ImageTk
+from scipy.optimize import minimize_scalar
+
+from src import videoCrop
+from src.imageTools import radialUndistort, perspectiveUndistort, getFrameFromVideo
+
+class AdvancedDistortionCorrectionTool:
+    def __init__(self, master):
+        self.master = master
+        self.master.title("Advanced Distortion Correction Tool")
+        self.master.geometry("1200x800")
+
+        self.image = None
+        self.original_image = None
+        self.displayed_image = None
+        self.k = 0  # Initialize K value
+        self.dots = []
+        self.dragging = None
+
+        # Create UI elements
+        # self.load_button = tk.Button(self.master, text="Load Image", command=self.load_image)
+        # self.load_button.pack(pady=10)
+
+        self.canvas = tk.Canvas(self.master)
+        self.canvas.pack()
+        # self.canvas.width
+        self.canvas.bind("<Button-1>", self.on_click)
+        self.canvas.bind("<B1-Motion>", self.on_drag)
+        self.canvas.bind("<ButtonRelease-1>", self.on_release)
+
+        self.k_label = tk.Label(self.master, text="K value: 0")
+        self.k_label.pack(pady=5)
+
+        self.k_slider = tk.Scale(self.master, from_=-1.0, to=1.0, resolution=0.01, orient=tk.HORIZONTAL, length=300, command=self.update_k)
+        self.k_slider.pack(pady=10)
+
+        self.zoom_slider = tk.Scale(self.master, from_=0.5, to=5.0, resolution=0.01, orient=tk.HORIZONTAL, label="Zoom", length=300, command=self.update_zoom)
+        self.zoom_slider.pack(pady=10)
+
+        self.perspective_var = tk.BooleanVar()
+        self.perspective_check = tk.Checkbutton(self.master, text="Apply Perspective Correction", variable=self.perspective_var, command=self.update_persp)
+        self.perspective_check.pack(pady=10)
+
+        self.save_button = tk.Button(self.master, text="Save Image", command=self.save_image, state=tk.DISABLED)
+        self.save_button.pack(pady=10)
+
+        self.display_image_scale = 0.5
+
+        self.k_value = 0
+        self.zoom_value = 1
+
+    def update_zoom(self, zoom_value):
+        self.zoom_value = zoom_value
+        self.update_image(self.k_value, zoom_value)
+
+    def update_k(self, k_value):
+        self.k_value = k_value
+        self.update_image(k_value, self.zoom_value)
+
+    def update_persp(self):
+        self.update_image(self.k_value, self.zoom_value)
+
+    def load_image(self, image):
+        print(image)
+        self.original_image = image
+        self.save_button.config(state=tk.NORMAL)
+        height, width = self.original_image.shape[:2]
+        self.dots = [
+            [10, 10],
+            [width-10, 10],
+            [width-10, height-10],
+            [10, height-10]
+        ]
+        self.update_image(self.k_value, self.zoom_value)
+
+    def display_image(self):
+        if self.image is not None:
+            self.displayed_image = self.image.copy()
+            if not self.perspective_var.get():
+                for x, y in self.dots:
+                    cv2.circle(self.displayed_image, (int(x), int(y)), 15, (0, 255, 0), -1)
+            image = cv2.cvtColor(self.displayed_image, cv2.COLOR_BGR2RGB)
+            image = Image.fromarray(image)
+            image.thumbnail((image.width*self.display_image_scale, image.height*self.display_image_scale), Image.LANCZOS)
+            photo = ImageTk.PhotoImage(image)
+            self.canvas.config(width=photo.width(), height=photo.height())
+            self.canvas.create_image(0, 0, anchor=tk.NW, image=photo)
+            self.canvas.image = photo
+
+    def update_image(self, k_value, zoom_value):
+        if self.original_image is None:
+            return
+
+        self.image = radialUndistort(self.original_image, float(k_value), float(zoom_value))
+
+        if self.perspective_var.get():
+            self.image = perspectiveUndistort(self.image, self.dots)
+
+        self.display_image()
+        self.k_label.config(text=f"K value: {self.k:.2f}")
+
+    def count_lines(self, image):
+        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+        edges = cv2.Canny(gray, 50, 150, apertureSize=3)
+        lines = cv2.HoughLines(edges, 1, np.pi/180, 200)
+        return len(lines) if lines is not None else 0
+
+    def on_click(self, event):
+        x, y = event.x / self.display_image_scale, event.y / self.display_image_scale
+        for i, (dot_x, dot_y) in enumerate(self.dots):
+            if abs(x - dot_x) < 50 and abs(y - dot_y) < 50:
+                self.dragging = i
+                break
+
+
+    def on_drag(self, event):
+        if self.dragging is not None:
+            self.dots[self.dragging] = [event.x / self.display_image_scale, event.y / self.display_image_scale]
+            self.update_image(self.k_value, self.zoom_value)
+
+    def on_release(self, event):
+        self.dragging = None
+
+    def save_image(self):
+        print(",".join([
+            str(self.k_value),
+            str(self.zoom_value),
+
+            str(self.dots[0][0]),
+            str(self.dots[0][1]),
+            str(self.dots[1][0]),
+            str(self.dots[1][1]),
+            str(self.dots[2][0]),
+            str(self.dots[2][1]),
+            str(self.dots[3][0]),
+            str( self.dots[3][1])
+        ]))
+
+if __name__ == "__main__":
+    if len(sys.argv) < 1:
+        print("Usage: undistort.py <event code>")
+    filename = sys.argv[1] + "/" + (os.listdir(sys.argv[1])[0])
+
+    root = tk.Tk()
+    app = AdvancedDistortionCorrectionTool(root)
+
+    start, stop = videoCrop.videoCrop(filename)
+    if start != stop:
+        print("Found video start and end")
+
+    app.load_image(getFrameFromVideo(filename, start))
+
+    root.mainloop()