riversnow/so101-segmentation-model

SO101 segmentation model

This is a model for segementation of images of the so101 robot arm, it was fine tuned over yolo11s

Sample code

Here's some sample code to use it

import cv2
import numpy as np
from ultralytics import YOLO

model = YOLO("weights/best.pt")
cap = cv2.VideoCapture("test_video.mp4")

w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = cap.get(cv2.CAP_PROP_FPS)

fourcc = cv2.VideoWriter_fourcc(*"avc1")
out = cv2.VideoWriter("comparison_output.mp4", fourcc, fps, (w * 2, h))

print("Generating side-by-side video...")

while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break

    results = model(frame)
    left_side = frame
    black_bg = np.zeros_like(frame)
    right_side = results[0].plot(img=black_bg, boxes=False, labels=True)
    combined_frame = np.hstack((left_side, right_side))

    out.write(combined_frame)

cap.release()
out.release()
print("Done! Check comparison_output.mp4")

Disclaimer : I vibe coded most of the code here, since it was one-time use code and I don't expect to publish it anywhere, I used https://github.com/johnsutor/so101-nexus to generate the synthetic images