diff --git a/.gitignore b/.gitignore index d9005f2..af367ff 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,6 @@ +*.mp4 +.idea/ + # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] diff --git a/downloader.py b/downloader.py new file mode 100644 index 0000000..bcebaa9 --- /dev/null +++ b/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 ") + +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) diff --git a/images/end.png b/images/end.png new file mode 100644 index 0000000..aaf6131 Binary files /dev/null and b/images/end.png differ diff --git a/images/start.png b/images/start.png new file mode 100644 index 0000000..0fcab19 Binary files /dev/null and b/images/start.png differ diff --git a/manualPath.py b/manualPath.py new file mode 100644 index 0000000..99777f4 --- /dev/null +++ b/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 ") + 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() \ No newline at end of file diff --git a/src/imageTools.py b/src/imageTools.py new file mode 100644 index 0000000..a234e18 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/videoCrop.py b/src/videoCrop.py new file mode 100644 index 0000000..50a77be --- /dev/null +++ b/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 \ No newline at end of file diff --git a/undistort.py b/undistort.py new file mode 100644 index 0000000..61e12f6 --- /dev/null +++ b/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("", self.on_click) + self.canvas.bind("", self.on_drag) + self.canvas.bind("", 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 ") + 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() diff --git a/vid-startstop.py b/vid-startstop.py new file mode 100644 index 0000000..e69de29