I am trying to take extracted frames from a video and create masks for those frames using YOLO segmentation models. When I attempt to run the code, it gives me an attribute error.
from ultralytics import YOLO
import cv2
import os
model_path = 'C:\\Users\\Analysis\\Video_Analysis\\best-maskmodel.pt'
folder_path = 'C:\\Users\\Analysis\\Video_Analysis\\Extracted_frames_idea'
model = YOLO(model_path)
output_folder_path = 'C:\\Users\\admin\\Desktop\\Masks' # Update with the desired output folder path
for filename in os.listdir(folder_path):
if filename.endswith('.jpg') or filename.endswith('.jpeg') or filename.endswith('.png'):
image_path = os.path.join(folder_path, filename)
img = cv2.imread(image_path)
H, W, _ = img.shape
results = model(img)
for result in results:
if result.masks is not None: # Check if masks are available
for j, mask in enumerate(result.masks.data):
mask = mask.numpy() * 255
mask = cv2.resize(mask, (W, H))
output_filename = 'mask_' + filename
output_path = os.path.join(output_folder_path, output_filename)
cv2.imwrite(output_path, mask)
I tried adding the v8SegmentationLoss
to the loss.py
and had no luck.