Spaces:
Runtime error
Runtime error
| import cv2 | |
| # import cvzone | |
| # from cvzone.SelfiSegmentationModule import SelfiSegmentation | |
| from PIL import Image | |
| import rembg | |
| # Load the input image | |
| import os | |
| import subprocess | |
| from .tele import SendToTele | |
| class RemoveBackground(): | |
| #images will be array | |
| def __init__(self,images,aspect_ratio,rgb,output_di,bw=False): | |
| self.images = images | |
| self.aspect_ratio = aspect_ratio | |
| self.rgb = rgb | |
| self.output_di = output_di | |
| self.bw = bw | |
| self.model_name = "silueta" | |
| self.session = rembg.new_session(model_name=self.model_name) | |
| def deleteimage(self,file_path): | |
| # Check if the file exists before deleting | |
| if os.path.exists(file_path): | |
| # Delete the file | |
| os.remove(file_path) | |
| print("File deleted successfully.") | |
| else: | |
| print("File does not exist.") | |
| def remove(self): | |
| aspect_ratio = self.aspect_ratio | |
| for i in self.images: | |
| print(i) | |
| imgread = cv2.imread(i[0]) | |
| self.deleteimage(i[0]) | |
| if self.bw: | |
| rgb = (128,128,128,0) | |
| else: | |
| rgb = (255,0,0,0) | |
| img = rembg.remove(imgread,bgcolor=rgb, alpha_matting=True, alpha_matting_foreground_threshold=270,alpha_matting_background_threshold=20, alpha_matting_erode_size=11) | |
| # Convert the image to grayscale | |
| gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) | |
| # Load the face detector | |
| face_cascade = cv2.CascadeClassifier('haarcascade_frontalface_default.xml') | |
| if face_cascade.empty(): | |
| raise IOError('Unable to load the face cascade classifier xml file') | |
| # Detect faces in the grayscale image | |
| faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5) | |
| if len(faces) > 0: | |
| # Calculate the center of the bounding box | |
| x, y, w, h = faces[0] | |
| center_x = x + w // 2 | |
| center_y = y + h // 2 | |
| # Calculate the width and height of the cropped image | |
| aspect_ratio = self.aspect_ratio | |
| if w / h > aspect_ratio: | |
| crop_width = int(aspect_ratio * h) | |
| crop_height = h | |
| else: | |
| crop_width = w | |
| crop_height = int(w / aspect_ratio) | |
| # Calculate the coordinates of the top-left corner of the cropped image | |
| margin_x = crop_width // 2 | |
| margin_y = crop_height // 2 | |
| crop_x = max(0, center_x - crop_width // 2 - margin_x) | |
| crop_y = max(0, center_y - crop_height // 2 - margin_y) | |
| # Crop the image | |
| if self.bw: | |
| crop_img = gray[crop_y:crop_y+crop_height+margin_y*2, crop_x:crop_x+crop_width+margin_x*2] | |
| else: | |
| crop_img = img[crop_y:crop_y+crop_height+margin_y*2, crop_x:crop_x+crop_width+margin_x*2] | |
| else: | |
| print("Face not detect") | |
| basename =os.path.basename(i[0]) | |
| cv2.imwrite(self.output_di+basename,crop_img) | |
| SendToTele(document_path=self.output_di+basename) | |
| # images = {('img.jpg',5)} | |
| # paper_size = (1200,1800) | |
| # gap = 40; # gap between images; | |
| # img_max_width = 1.1 * 300 | |
| # img_max_height = 1.3 * 300 | |
| # aspect_ratio = img_max_width / img_max_height | |
| # bg_rgb = (255,0,0) | |
| # output_di = 'img/' | |
| # a = RemoveBackground(images,aspect_ratio,bg_rgb,output_di,True) | |
| # a.remove() | |