228 lines
7.6 KiB
Python
228 lines
7.6 KiB
Python
import glob
|
|
import os
|
|
import random
|
|
import cv2
|
|
|
|
# --- CONFIGURATION ---
|
|
OUTPUT_DIRECTORY = "dataset"
|
|
FRAMES_TO_EXTRACT = 20 # Number of random frames to pull
|
|
SCALE_FACTOR = 0.5 # Resize video to 1/2 size for easier viewing
|
|
|
|
# Define the EXACT size of the cropped images you want to save
|
|
BOX_WIDTH = 64
|
|
BOX_HEIGHT = 64
|
|
|
|
# Defined classification options
|
|
CLASSIFICATION_OPTIONS = ["guardian", "interceptor", "sentinel", "rocket", "bomber", "background", "asteroid"]
|
|
# ---------------------
|
|
|
|
|
|
def get_dataset_counts():
|
|
"""Counts the number of existing images in each classification folder."""
|
|
counts = {}
|
|
for option in CLASSIFICATION_OPTIONS:
|
|
folder_path = os.path.join(OUTPUT_DIRECTORY, option)
|
|
if os.path.exists(folder_path):
|
|
# Count only common image files
|
|
files = [
|
|
f
|
|
for f in os.listdir(folder_path)
|
|
if f.lower().endswith((".jpg", ".jpeg", ".png"))
|
|
]
|
|
counts[option] = len(files)
|
|
else:
|
|
counts[option] = 0
|
|
return counts
|
|
|
|
|
|
def mouse_click_callback(event, x, y, flags, param):
|
|
"""Handles mouse clicks.
|
|
|
|
Centers a fixed-size box on the click, crops it, and saves it.
|
|
"""
|
|
if event == cv2.EVENT_LBUTTONDOWN:
|
|
frame = param["frame"]
|
|
display_frame = param["display_frame"]
|
|
label_dir = param["label_dir"]
|
|
frame_idx = param["frame_idx"]
|
|
window_name = param["window_name"]
|
|
|
|
# Calculate top-left corner based on centering the box on the click
|
|
x1 = x - BOX_WIDTH // 2
|
|
y1 = y - BOX_HEIGHT // 2
|
|
x2 = x1 + BOX_WIDTH
|
|
y2 = y1 + BOX_HEIGHT
|
|
|
|
# --- Edge Case Handling ---
|
|
if x1 < 0:
|
|
x1 = 0
|
|
x2 = BOX_WIDTH
|
|
if y1 < 0:
|
|
y1 = 0
|
|
y2 = BOX_HEIGHT
|
|
if x2 > frame.shape[1]:
|
|
x2 = frame.shape[1]
|
|
x1 = x2 - BOX_WIDTH
|
|
if y2 > frame.shape[0]:
|
|
y2 = frame.shape[0]
|
|
y1 = y2 - BOX_HEIGHT
|
|
|
|
# Extract the crop from the scaled frame
|
|
crop = frame[y1:y2, x1:x2]
|
|
|
|
# Save the crop
|
|
count = param["counter"]
|
|
filename = f"frame_{frame_idx}_crop_{count}.jpg"
|
|
save_path = os.path.join(label_dir, filename)
|
|
cv2.imwrite(save_path, crop)
|
|
|
|
print(f" Saved crop {count} to {save_path}")
|
|
param["counter"] += 1
|
|
|
|
# Draw visual feedback (a green box) on the screen where you clicked
|
|
cv2.rectangle(display_frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
|
cv2.imshow(window_name, display_frame)
|
|
cv2.moveWindow(window_name, 0, 0)
|
|
|
|
|
|
def label_video_session():
|
|
print("=========================================")
|
|
print(" AI DATASET SESSION LABELER ")
|
|
print("=========================================\n")
|
|
|
|
# 1. Automatically find and select a random MP4 video
|
|
mp4_files = glob.glob("*.mp4")
|
|
if not mp4_files:
|
|
print(
|
|
"Error: No .mp4 files found in the current directory.\n"
|
|
"Please place this script in the same folder as your video files."
|
|
)
|
|
return
|
|
|
|
selected_video = random.choice(mp4_files)
|
|
print(f"-> Selected Video File: '{selected_video}'")
|
|
|
|
# 2. Fetch and display current dataset balance dashboard
|
|
counts = get_dataset_counts()
|
|
print("\n-----------------------------------------")
|
|
print(" CURRENT DATASET BALANCE STATUS ")
|
|
print("-----------------------------------------")
|
|
for option, count in counts.items():
|
|
print(f" • {option.capitalize():<12} : {count} items")
|
|
print("-----------------------------------------")
|
|
print(
|
|
"💡 Tip: Try to choose options with lower counts to keep data balanced!"
|
|
)
|
|
|
|
# 3. Force selection of a valid label
|
|
print("\nAvailable Classification Targets:")
|
|
for idx, option in enumerate(CLASSIFICATION_OPTIONS, start=1):
|
|
print(f" [{idx}] {option} (Current: {counts[option]})")
|
|
|
|
while True:
|
|
try:
|
|
choice = input(
|
|
f"\nEnter the number (1-{len(CLASSIFICATION_OPTIONS)}) of what you are labeling: "
|
|
).strip()
|
|
choice_idx = int(choice) - 1
|
|
if 0 <= choice_idx < len(CLASSIFICATION_OPTIONS):
|
|
session_label = CLASSIFICATION_OPTIONS[choice_idx]
|
|
break
|
|
else:
|
|
print("Invalid selection. Please choose a number from the list.")
|
|
except ValueError:
|
|
print("Invalid input. Please enter a number.")
|
|
|
|
# Open video and check validity
|
|
cap = cv2.VideoCapture(selected_video)
|
|
if not cap.isOpened():
|
|
print(f"Error: Could not open video file {selected_video}")
|
|
return
|
|
|
|
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
|
if total_frames == 0:
|
|
print("Error: Video has 0 frames.")
|
|
return
|
|
|
|
# Select random unique frames
|
|
num_frames = min(FRAMES_TO_EXTRACT, total_frames)
|
|
frame_indices = random.sample(range(total_frames), num_frames)
|
|
frame_indices.sort()
|
|
|
|
# Set up directories
|
|
label_dir = os.path.join(OUTPUT_DIRECTORY, session_label)
|
|
os.makedirs(label_dir, exist_ok=True)
|
|
|
|
# Comprehensive User Instructions
|
|
print("\n" + "=" * 50)
|
|
print(" HOW TO LABEL")
|
|
print("=" * 50)
|
|
print(f"• TARGET OBJECT: {session_label.upper()}")
|
|
print(f"• SESSION SCOPE: {num_frames} random frames from '{selected_video}'")
|
|
print("• CROPPING SIZE: 64x64 pixels (centered automatically around your click)")
|
|
print("-" * 50)
|
|
print("► CONTROLS:")
|
|
print(" [Left-Click] - Click directly on the target object to save a snapshot.")
|
|
print(" (You can click multiple objects if they appear on screen)")
|
|
print(" [SPACE or ENTER] - Advance to the next random video frame.")
|
|
print(" [Q Key] - Quit and save all progress up to this point.")
|
|
print("=" * 50)
|
|
input("\nPress ENTER when you are ready to begin...")
|
|
|
|
for i, frame_idx in enumerate(frame_indices):
|
|
cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
continue
|
|
|
|
# Resize the frame by 1/2 for laptop viewability
|
|
frame_resized = cv2.resize(
|
|
frame, (0, 0), fx=SCALE_FACTOR, fy=SCALE_FACTOR
|
|
)
|
|
|
|
# Create a copy for drawing the green boxes dynamically
|
|
frame_display = frame_resized.copy()
|
|
|
|
window_name = (
|
|
f"Labeling: {session_label} | Frame {i+1}/{num_frames}"
|
|
)
|
|
cv2.namedWindow(window_name)
|
|
|
|
# Dictionary to pass state/variables into the mouse callback function
|
|
session_state = {
|
|
"frame": frame_resized,
|
|
"display_frame": frame_display,
|
|
"label_dir": label_dir,
|
|
"frame_idx": frame_idx,
|
|
"window_name": window_name,
|
|
"counter": 0,
|
|
}
|
|
|
|
# Bind the mouse click event to the window
|
|
cv2.setMouseCallback(window_name, mouse_click_callback, session_state)
|
|
|
|
# Keep window open until user hits Space/Enter or 'q'
|
|
cv2.imshow(window_name, frame_display)
|
|
while True:
|
|
cv2.moveWindow(window_name, 0, 0)
|
|
key = cv2.waitKey(1) & 0xFF
|
|
if key == 13 or key == 32:
|
|
break
|
|
elif key == ord("q") or key == ord("Q"):
|
|
print("\nQuitting session early... Saving progress.")
|
|
cap.release()
|
|
cv2.destroyAllWindows()
|
|
return
|
|
|
|
cv2.destroyWindow(window_name)
|
|
|
|
cap.release()
|
|
cv2.destroyAllWindows()
|
|
print(
|
|
f"\nSession finished! All crops saved successfully in '{label_dir}'."
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
label_video_session()
|