add segment

.gitignore

@@ -160,3 +160,6 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+# checkpoints
+*.pt

README.md

@@ -139,3 +139,23 @@ ai-video convert input_video.mp4 output_video.mp4 --video_codec libx265 --audio_
 This command will convert `input_video.mp4` to `output_video.mp4`, using the libx265 video codec and mp3 audio codec. It will also resize and crop the video to 1920x1080 pixels. If the input video's aspect ratio doesn't match 1920x1080, the video will be scaled to fit within these dimensions and centered.
 
 If you don't specify an output file, the tool will automatically generate one with the suffix "_converted" added to the input filename.
+
+### 6. AI Segment Video
+
+Segments people in a video using AI models (YOLO and SAM2) and replaces the background with green.
+
+```
+ai-segment <input_file> <output_file>
+```
+
+- `<input_file>`: Path to the input video file
+- `<output_file>`: Path for the output segmented video file
+
+Example:
+```
+ai-segment input_video.mp4 segmented_output.mp4
+```
+
+This command will process `input_video.mp4` using YOLO for person detection and SAM2 for segmentation. It will create a new video `segmented_output.mp4` where all detected people are segmented, and the background is replaced with green.
+
+Note: This command requires additional AI models (YOLO and SAM2) which will be downloaded automatically on first use. The process may take some time depending on the length of the video and your hardware capabilities.

ai_video_cli/segment.py

@@ -0,0 +1,62 @@
+import cv2
+import numpy as np
+from ultralytics import YOLO, SAM
+import argparse
+
+def process_video(video_path, output_path):
+    # Load models - Ultralytics will handle caching automatically
+    yolo_model = YOLO("yolo11n.pt")
+    sam2_model = SAM("sam2_b.pt")
+
+    cap = cv2.VideoCapture(video_path)
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps = cap.get(cv2.CAP_PROP_FPS)
+
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        # Detect people using YOLO model
+        yolo_results = yolo_model(frame)
+
+        # Filter for person class (adjust the class index if needed)
+        person_boxes = yolo_results[0].boxes[yolo_results[0].boxes.cls == 0].xyxy.cpu().numpy()
+
+        # Use SAM 2 for segmentation
+        sam_results = sam2_model(frame, bboxes=person_boxes)
+
+        # Combine all person masks
+        combined_mask = np.zeros(frame.shape[:2], dtype=bool)
+        for mask in sam_results[0].masks.data:
+            combined_mask |= mask.cpu().numpy()
+
+        # Apply the mask to the original frame
+        segmented_frame = frame.copy()
+        segmented_frame[~combined_mask] = [
+            0,
+            255,
+            0,
+        ]  # Green background, you can change this
+
+        out.write(segmented_frame)
+
+    cap.release()
+    out.release()
+
+
+def main():
+
+    parser = argparse.ArgumentParser(description="Process video with YOLO and SAM2")
+    parser.add_argument("input_video", help="Path to the input video file")
+    parser.add_argument("output_video", help="Path to the output video file")
+    args = parser.parse_args()
+
+    process_video(args.input_video, args.output_video)
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -7,13 +7,16 @@ readme = "README.md"
 packages = [{ include = "ai_video_cli" }]
 
 [tool.poetry.dependencies]
-python = ">=3.9,<3.12"
+python = ">=3.10,<3.12"
 moviepy = "^1.0.3"
 numpy = "^1.26.0"
 pillow = "^9"
+opencv-python = "^4.10.0.84"
+ultralytics = "^8.3.13"
 
 [tool.poetry.scripts]
 ai-video = "ai_video_cli.main:main"
+ai-segment = "ai_video_cli.segment:main"
 
 [build-system]
 requires = ["poetry-core"]