kadirnar / segment-anything-video

MetaSeg: Packaged version of the Segment Anything repository
Apache License 2.0
952 stars 67 forks source link

use shi to segment video #66

Closed e-tuanzi closed 1 year ago

e-tuanzi commented 1 year ago

I have a problem when i run this code. I don't konw why ModuleNotFoundError: No module named 'yolov5'. I install yolov5 module by following code.

pip install ultralytics

Can someone give me some advice? Thank you very much!

from metaseg.sahi_predict import SahiAutoSegmentation, sahi_sliced_predict
import cv2

cap = cv2.VideoCapture('./test_data/red_girl.mp4')
fourcc = cv2.VideoWriter_fourcc(*'MP4V')  # 视频编解码器
fps = cap.get(cv2.CAP_PROP_FPS)  # 帧数
width, height = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))  # 宽高
out = cv2.VideoWriter('./output/read_girl_shi.mp4', fourcc, fps, (width, height))  # 写入视频
# Read the first frame
ret, frame = cap.read()
while ret:
    cv2.imwrite("./test_data/temp.jpg", frame)

    image_path = "./test_data/temp.jpg"
    boxes = sahi_sliced_predict(
        image_path=image_path,
        detection_model_type="yolov5", #yolov8, detectron2, mmdetection, torchvision
        detection_model_path="yolov5l6.pt",
        conf_th=0.25,
        image_size=1280,
        slice_height=256,
        slice_width=256,
        overlap_height_ratio=0.2,
        overlap_width_ratio=0.2,
    )
    SahiAutoSegmentation().predict(
        source=image_path,
        model_type="vit_b",
        input_box=boxes,
        multimask_output=False,
        random_color=False,
        show=True,
        save=True,
        output_path="./output/temp.jpg"
    )
    image = cv2.imread("./output/temp.jpg")
    # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    out.write(image)

cap.release()
onuralpszr commented 1 year ago

We updated the readme accordingly https://github.com/kadirnar/segment-anything-video/commit/3a874806f4fbd155ba0d5f2c39648eb1f7adbfe8

So this should be more clear how to use it.

Thank you.