diff --git a/Data/Collect/CollectImage.py b/Data/Collect/CollectImage.py index 2fd3cd2..e0336f5 100644 --- a/Data/Collect/CollectImage.py +++ b/Data/Collect/CollectImage.py @@ -1,5 +1,5 @@ # BUROBOT -import time, os, cv2, keyboard, sys +import time, os, cv2, keyboard, sys, pyautogui # Import the necessary modules and packages sys.path.append(os.path.join(os.path.abspath(__file__).split("Burobot")[0], "Burobot")) @@ -7,10 +7,10 @@ # Helper function to check input errors -def _check_errs(save_to_path: str, resolution: tuple, delay: float = 0.1, i: int = 0): +def _check_errs(saveToPath: str, resolution: tuple, delay: float = 0.1, i: int = 0): # Check if the save path exists - if not os.path.exists(save_to_path): - raise FileNotFoundError("Can't find path 🤷\nsave_to_path:" + str(save_to_path)) + if not os.path.exists(saveToPath): + raise FileNotFoundError("Can't find path 🤷\nsaveToPath:" + str(saveToPath)) # Check if the resolution has correct dimensions if len(resolution) != 2: raise ValueError("resolution must be like (w>0, h>0) 🖼️") @@ -23,9 +23,9 @@ def _check_errs(save_to_path: str, resolution: tuple, delay: float = 0.1, i: int # Function to collect images from the webcam -def collect_from_webcam( - save_to_path: str, - output_resolution: tuple, +def collectFromWebcam( + saveToPath: str, + outputResolution: tuple, webcam: int = 0, delay: float = 0.1, i: int = 0, @@ -33,16 +33,16 @@ def collect_from_webcam( # Clear the outputs BurobotOutput.clearAndMemoryTo() # Check input errors - _check_errs(save_to_path, output_resolution, delay, i) + _check_errs(saveToPath, outputResolution, delay, i) # Create a 'cv2.VideoCapture' object to capture video from the webcam cap = cv2.VideoCapture(webcam) collect = False while True: # Read a frame from the webcam - _, frame_orj = cap.read() + _, frameOriginal = cap.read() # Resize the frame to the desired output resolution - frame_orj = cv2.resize(frame_orj, output_resolution) - frame = frame_orj.copy() + frameOriginal = cv2.resize(frameOriginal, outputResolution) + frame = frameOriginal.copy() # Add instructions to the frame frame = cv2.putText( frame, @@ -60,12 +60,12 @@ def collect_from_webcam( cv2.destroyAllWindows() break if collect: - img_name = os.path.join(save_to_path, f"{str(i)}.jpg") + imgName = os.path.join(saveToPath, f"{str(i)}.jpg") try: - cv2.imwrite(img_name, frame_orj) + cv2.imwrite(imgName, frameOriginal) except: - os.remove(img_name) - cv2.imwrite(img_name, frame_orj) + os.remove(imgName) + cv2.imwrite(imgName, frameOriginal) i += 1 time.sleep(delay) frame = cv2.putText( @@ -85,11 +85,11 @@ def collect_from_webcam( # Function to collect images from the screen -def collect_from_screen( - save_to_path: str, output_resolution: tuple, delay: float = 0.1, i: int = 0 +def collectFromScreen( + saveToPath: str, outputResolution: tuple, delay: float = 0.1, i: int = 0 ): BurobotOutput.clearAndMemoryTo() - _check_errs(save_to_path, output_resolution, delay, i) + _check_errs(saveToPath, outputResolution, delay, i) from PIL import ImageDraw, ImageGrab import pyautogui @@ -107,7 +107,7 @@ def collect_from_screen( draw = ImageDraw.Draw(screen) - draw.rectangle([x, y], outline="red") #type: ignore + draw.rectangle([x, y], outline="red") # type: ignore # Show the rectangle screen.show() @@ -124,21 +124,21 @@ def collect_from_screen( if collect: img = pyautogui.screenshot() - if not x[0] <= y[0] or not x[1] <= y[1]:#type: ignore - temp_x = x#type: ignore - x = y#type: ignore + if not x[0] <= y[0] or not x[1] <= y[1]: # type: ignore + temp_x = x # type: ignore + x = y # type: ignore y = temp_x del temp_x - img = img.crop((x[0], x[1], y[0], y[1]))#type: ignore - img = img.resize(output_resolution)#type: ignore + img = img.crop((x[0], x[1], y[0], y[1])) # type: ignore + img = img.resize(outputResolution) # type: ignore BurobotOutput.clearAndMemoryTo() try: - img.save(os.path.join(save_to_path, str(i) + ".jpg")) + img.save(os.path.join(saveToPath, str(i) + ".jpg")) print(str(i) + ".jpg saved!") except: - os.remove(os.path.join(save_to_path, str(i) + ".jpg")) + os.remove(os.path.join(saveToPath, str(i) + ".jpg")) print(str(i) + ".jpg deleted!") - img.save(os.path.join(save_to_path, str(i) + ".jpg")) + img.save(os.path.join(saveToPath, str(i) + ".jpg")) print(str(i) + ".jpg saved!") i += 1 time.sleep(delay) @@ -146,4 +146,64 @@ def collect_from_screen( return i +def collectFromVideo(videoPath: str, saveToPath: str, i: int = 0, printInfo:bool = True): + """ + Collects frames from a video by pressing 'c' key and saves them to the specified path. + + :videoPath (str): The path of the input video file. + :saveToPath (str): The directory where the frames will be saved. + :i (int): An optional parameter to set the starting index of the saved frames. + """ + # Open the video file + cap = cv2.VideoCapture(videoPath) + screenWidth, screenHeight = pyautogui.size() + + if not cap.isOpened(): + print("Could not open the video file.") + return + + + + save = False + while True: + try: + # Read a frame + ret, frame = cap.read() + frame = cv2.resize(frame, (int(screenWidth*0.3), int(screenHeight*0.3))) + if not ret: + print("Video ended or could not read.") + break + + # Save the frame when 'c' is pressed + key = cv2.waitKey(30) & 0xFF + if key == ord('v'): + save = not save + if key == ord("c") or save: + framePath = f"{saveToPath}/frame_{i}.png" + cv2.imwrite(framePath, frame) + print(f"Frame saved: {framePath}") + i += 1 + # Exit when 'q' is pressed + elif key == ord("q"): + break + + # Display information messages on the screen + if printInfo: + infoMessage = "Press 'c' to save. Press 'v' to toggle continuous saving. Press 'q' to exit." + frame = cv2.putText( + frame, infoMessage, (int(screenWidth*0.3*0.2), int(screenHeight*0.3*0.8)), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2 + ) + frame = cv2.putText( + frame, "Continuous saving: "+str("YES" if save else "NO"), (int(screenWidth*0.3*0.05), int(screenHeight*0.3*0.05)), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2 + ) + except: + frame = [[0,0,0]] + pass + cv2.imshow("Video", frame) + + # Release the video capture and close the window + cap.release() + cv2.destroyAllWindows() + + # BUROBOT