|
|
|
"""Gen AI Project1 |
|
|
|
Automatically generated by Colab. |
|
|
|
Original file is located at |
|
https://colab.research.google.com/drive/1Q27-bhi-hIw4U_QKiDXy3bwfLjPOI02o |
|
""" |
|
|
|
|
|
|
|
|
|
|
|
!pip install datasets |
|
|
|
!pip install ultralytics |
|
|
|
from datasets import load_dataset |
|
from moviepy.editor import ImageSequenceClip |
|
from ultralytics import YOLO |
|
import os |
|
import cv2 |
|
from PIL import Image |
|
|
|
|
|
dataset = load_dataset("VarunB31990/Video-Editing-Dataset") |
|
|
|
|
|
model = YOLO("yolov8n.pt") |
|
|
|
|
|
image_dir = "images/" |
|
os.makedirs(image_dir, exist_ok=True) |
|
processed_dir = "processed_frames/" |
|
os.makedirs(processed_dir, exist_ok=True) |
|
|
|
processed_paths = [] |
|
for i, item in enumerate(dataset["train"]): |
|
if "original_image" in item: |
|
image = item["original_image"] |
|
image_path = os.path.join(image_dir, f"frame_{i}.jpg") |
|
image.save(image_path) |
|
|
|
results = model(image_path) |
|
for result in results: |
|
im_array = result.plot() |
|
im = Image.fromarray(im_array) |
|
detected_path = os.path.join(processed_dir, f"detected_{i}.jpg") |
|
im.save(detected_path) |
|
processed_paths.append(detected_path) |
|
|
|
if len(processed_paths)>1: |
|
clip = ImageSequenceClip(processed_paths, fps=10) |
|
clip.write_videofile("yolo_detection_video.mp4", codec="libx264", fps=10) |
|
print("🎥 Video created: yolo_detection_video.mp4") |
|
else: |
|
print("⚠️ Not enough images to create a video.") |
|
|
|
from IPython.display import display, Video |
|
display(Video("yolo_detection_video.mp4", embed=True)) |