Env_mixer / bgrem.py
Inmental's picture
Upload folder using huggingface_hub
4c62147 verified
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()