import cv2
import numpy as np
from rembg import remove

def remove_background(frame):
    result = remove(frame)
    return result

def main():
    cap = cv2.VideoCapture(0)  # Capture video from webcam

    if not cap.isOpened():
        print("Error: Could not open webcam.")
        return

    while True:
        ret, frame = cap.read()
        if not ret:
            print("Error: Could not read frame.")
            break

        # Resize frame to improve performance
        frame = cv2.resize(frame, (640, 480))

        # Remove background
        result = remove_background(frame)

        # Convert to a format suitable for OpenCV
        result = cv2.cvtColor(np.array(result), cv2.COLOR_RGB2BGR)

        # Display the resulting frame
        cv2.imshow('Background Removal', result)

        # Press 'q' to exit the loop
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()