파이썬 rembg 활용해서 이미지 배경제거하기
∣
2023년 9월 19일
rembg는 이미지의 배경을 제거하는 패키지입니다. 이 도구를 사용하면 이미지의 주요 객체를 분리하여 배경 없는 PNG 이미지로 추출할 수 있습니다. rembg는 특히 사진에서 사람, 동물, 물체 등의 객체를 감지하고 배경을 제거하여 투명 배경의 이미지를 생성하는 데 사용됩니다.
# 이미지를 바이트 배열로 변환
image_bytes = cv2.imencode(".png", processed_image)[1].tobytes()
# 배경 제거
result = remove(image_bytes)
# 결과 이미지를 표시
img = Image.open(io.BytesIO(result))
본 포스트에서는 Qt5와 opencv를 사용하여 실시간 배경제거를 할 계획이다. 아래 버젼에서 실행 됐습니다.
PyQt5==5.15.9
PyQt5-Qt5==5.15.2
PyQt5-sip==12.12.2
rembg==2.0.50
pooch==1.7.0
opencv-contrib-python==4.8.0.76
opencv-python==4.8.0.76
opencv-python-headless==4.8.1.78
import sys
import cv2
import numpy as np
from PyQt5.QtCore import Qt, QTimer, QRect, QSize
from PyQt5.QtGui import QImage, QPixmap, QPainter
from PyQt5.QtWidgets import QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, QSplitter, QPushButton, QLabel, QSizePolicy
from PyQt5.QtGui import QPalette, QColor
from PIL import Image
import io
from rembg import remove
먼저 아래는 WebcamApp의 초기화 코드 입니다.
class WebcamApp(QMainWindow):
def __init__(self):
super().__init__()
self.saved_image_count = 0 # 저장된 이미지 수를 추적하기 위한 변수
self.setWindowTitle("Webcam and Capture")
self.setGeometry(100, 100, 800, 400)
central_widget = QWidget(self)
self.setCentralWidget(central_widget)
layout = QHBoxLayout(central_widget)
splitter = QSplitter(Qt.Horizontal)
layout.addWidget(splitter)
left_widget = QWidget(self)
right_widget = QWidget(self)
splitter.addWidget(left_widget)
splitter.addWidget(right_widget)
splitter.setSizes([5, 5])
left_layout = QVBoxLayout(left_widget)
self.webcam_label = QLabel(self)
self.webcam_label.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
left_layout.addWidget(self.webcam_label)
# 좌측 하단에 캡처 버튼 추가
capture_button = QPushButton("Capture", self)
capture_button.clicked.connect(self.capture_image)
left_layout.addWidget(capture_button)
right_layout = QVBoxLayout(right_widget)
self.capture_label = QLabel(self)
self.capture_label.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
right_layout.addWidget(self.capture_label)
# 우측 하단에 저장 버튼 추가
save_button = QPushButton("Save", self)
save_button.clicked.connect(self.save_image)
right_layout.addWidget(save_button)
self.webcam = cv2.VideoCapture(0)
self.timer = QTimer(self)
self.timer.timeout.connect(self.update_webcam)
self.timer.start(30)
# 배경색 설정 (흰색)
palette = QPalette()
palette.setColor(QPalette.Window, QColor(255, 255, 255))
self.setPalette(palette)
화면 구성을 한 부분이며 총 두개의 디스플레이와 두개의 버튼을 구성하였습니다.
class WebcamApp(QMainWindow):
#... 이전 init 코드
def update_webcam(self):
ret, frame = self.webcam.read()
frame = cv2.flip(frame, 1) # 그대 얼굴에 자신감을 주는 코드
if ret:
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) # BGR to RGB 변환
h, w, ch = frame.shape
bytes_per_line = ch * w
qt_image = QImage(frame.data, w, h, bytes_per_line, QImage.Format_RGB888)
pixmap = QPixmap.fromImage(qt_image)
# 웹캠 레이블의 크기를 현재 창 크기에 맞추기
self.webcam_label.setPixmap(pixmap.scaled(self.webcam_label.size(), Qt.KeepAspectRatio))
self.current_frame = frame # 현재 프레임 저장
실시간으로 웹캠 프레임을 읽어 크기를 조정하고 self.current_frame으로 설정합니다.
class WebcamApp(QMainWindow):
#... 이전 init, update_webcam코드
# capture_image 함수 내에서 호출하여 이미지 처리
def capture_image(self):
if hasattr(self, 'current_frame'):
self.captured_image = cv2.cvtColor(self.current_frame, cv2.COLOR_RGB2BGR) # RGB를 BGR로 변환하여 저장
#####################################
# 중앙으로 1대1 비율 조정 #
#####################################
# h, w, _ = self.captured_image.shape
# min_dim = min(h, w)
# top = (h - min_dim) // 2
# left = (w - min_dim) // 2
# bottom = top + min_dim
# right = left + min_dim
# self.captured_image = self.captured_image[top:bottom, left:right]
# # 512x512 크기로 리사이즈
# self.captured_image = cv2.resize(self.captured_image, (512, 512))
#####################################
# 중앙으로 1대1 비율 조정 #
#####################################
#####################################
# 얼굴 밑으로 출력 삭제 #
#####################################
# # 얼굴 감지를 위한 Haar Cascade 분류기 로드
# face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
# # 그레이스케일 이미지로 변환
# gray_image = cv2.cvtColor(self.captured_image, cv2.COLOR_BGR2GRAY)
# # 얼굴 감지
# faces = face_cascade.detectMultiScale(gray_image, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))
# # 감지된 얼굴 영역을 흰색으로 만듭니다.
# for (x, y, w, h) in faces:
# self.captured_image[y+h+30:, :] = [255, 255, 255]
#####################################
# 얼굴 밑으로 출력 삭제 #
#####################################
# 이미지에서 얼굴 감지 및 흰색으로 만들기
processed_image = self.captured_image
# 이미지를 바이트 배열로 변환
image_bytes = cv2.imencode(".png", processed_image)[1].tobytes()
# 배경 제거
result = remove(image_bytes)
# 결과 이미지를 표시
img = Image.open(io.BytesIO(result))
q_image = QImage(img.tobytes(), img.width, img.height, img.width * 4, QImage.Format_RGBA8888)
pixmap = QPixmap.fromImage(q_image)
self.captured_image = q_image
self.capture_label.setPixmap(pixmap.scaled(self.capture_label.size(), Qt.KeepAspectRatio))
해당 이미지를 바이트 배열로 변환 후 배경을 삭제하고 창 사이즈에 맞춰서 이미지를 출력하는 코드입니다. 각 주석은 사진을 1대1로 변경하는 코드, 얼굴을 감지해서 얼굴 밑으로는 출력하지 않는 코드입니다.
class WebcamApp(QMainWindow):
#... init, update_webcam, capture_image 코드
def save_image(self):
if hasattr(self, 'captured_image'):
self.saved_image_count += 1
file_name = f"captured_image_{self.saved_image_count}.png"
pixmap = QPixmap.fromImage(self.captured_image)
##########################
# 저화질로 변경 #
##########################
# pixmap = pixmap.scaled(QSize(256, 256), Qt.KeepAspectRatio, Qt.SmoothTransformation)
# # 새로운 크기 설정
# new_width = 512
# new_height = 512
# # 빈 QImage 생성
# result_image = QImage(new_width, new_height, QImage.Format_ARGB32)
# result_image.fill(Qt.transparent) # 투명한 배경으로 채우기
# # 이미지를 중앙에 배치할 위치 계산
# x_offset = (new_width - pixmap.width()) // 2
# y_offset = (new_height - pixmap.height()) // 2
# # 이미지를 새 이미지에 중앙에 복사
# painter = QPainter(result_image)
# target_rect = QRect(x_offset, y_offset, pixmap.width(), pixmap.height())
# source_rect = pixmap.rect()
# painter.drawPixmap(target_rect, pixmap, source_rect)
# painter.end()
# # 결과 이미지를 QPixmap으로 변환
# pixmap = QPixmap.fromImage(result_image)
##########################
# 저화질로 변경 #
##########################
pixmap.save(file_name, "PNG")
print(f"Image saved as '{file_name}'")
배경 제거된 이미지를 저장하는 코드 입니다. 각각 사이즈를 리사이즈 가능하며 축소된 이미지로 저장하는 코드가 주석으로 남겨져 있습니다.
class WebcamApp(QMainWindow):
#... init, update_webcam, capture_image, save_image 코드
def closeEvent(self, event):
self.webcam.release()
event.accept()
qt5의 스타일을 설정하고 gui를 실행시키는 명령어 입니다.
if __name__ == "__main__":
app = QApplication(sys.argv)
window = WebcamApp()
# PyQt5 스타일 설정
app.setStyle("Fusion")
window.setStyleSheet(
"QPushButton {"
" background-color: #4CAF50;"
" color: white;"
" border: none;"
" padding: 10px 20px;"
" font-size: 16px;"
" cursor: pointer;"
"}"
"QPushButton:hover {"
" background-color: #45a049;"
"}"
"QLabel {"
" font-size: 18px;"
"}"
)
window.show()
sys.exit(app.exec_())
성능은 매우 괜찮다!… 나름 만족만족하는 중이다.
rembg는 정확한 특정 개체를 지정해 배경 제거가 아닌 주가 돼보이는 물체를 중심으로 배경 제거 하기 때문에 세세한 셋팅으로 배경제거는 힘들다. 하지만! 이런 좋은 코드를 무료 배포해주셔서 감사합니다… 끝!
댓글