Spaces:
Running
on
Zero
Running
on
Zero
| import os | |
| import rembg | |
| import numpy as np | |
| import cv2 | |
| from PIL import Image | |
| def recenter(image, h_begin=100, w_begin=220, res=256): | |
| h_image, w_image = image.shape[:2] | |
| new_image = np.zeros((res, res, 4), dtype=np.uint8) | |
| h_begin_new = -min(0, h_begin) | |
| w_begin_new = -min(0, w_begin) | |
| if h_begin > 0 and w_begin > 0: | |
| new_image = image[h_begin:h_begin+res, w_begin:w_begin+res] | |
| else: | |
| new_image[h_begin_new:h_begin_new+h_image, w_begin_new:w_image] = image | |
| return new_image | |
| def recover(image, original_size=(720, 480), h_begin=100, w_begin=220, res=256): | |
| target_w, target_h = original_size | |
| recovered_image = np.zeros((target_h, target_w, 4), dtype=np.uint8) | |
| h_begin_new = -min(0, h_begin) | |
| w_begin_new = -min(0, w_begin) | |
| if h_begin > 0 and w_begin > 0: | |
| recovered_image[h_begin:h_begin+res, w_begin:w_begin+res] = image | |
| else: | |
| recovered_image = image[h_begin_new:h_begin_new+target_h, w_begin_new:w_begin_new+target_w] | |
| return recovered_image.astype(np.uint8) | |
| def resize_and_center_crop(image, target_h=480, target_w=720): | |
| w, h = image.size | |
| image_ratio = w / h | |
| if target_w / target_h > image_ratio: | |
| new_w = target_w | |
| new_h = int(h * (target_w / w)) | |
| else: | |
| new_h = target_h | |
| new_w = int(w * (target_h / h)) | |
| image = image.resize((new_w, new_h), Image.LANCZOS) | |
| left = max(0, (new_w - target_w) // 2) | |
| top = max(0, (new_h - target_h) // 2) | |
| right = left + target_w | |
| bottom = top + target_h | |
| image = image.crop((left, top, right, bottom)) | |
| return image | |
| if __name__ == "__main__": | |
| base_dir = 'data_test' | |
| task_name = 'plane' | |
| raw_path = os.listdir(f'{base_dir}/raw_data') | |
| bg_remover = rembg.new_session() | |
| for image_path in raw_path: | |
| if not f'{task_name}_original' in image_path: | |
| continue | |
| input_image = Image.open(f'{base_dir}/raw_data/{image_path}') | |
| image = resize_and_center_crop(input_image) | |
| image.save(f'{base_dir}/raw_data/{image_path[:-4]}_resized.png') | |
| image.save(f'{base_dir}/{image_path.split("_")[0]}.png') | |
| image = np.array(image) | |
| carved_image = rembg.remove(image, session=bg_remover) # [H, W, 4] | |
| Image.fromarray(carved_image).save(f'{base_dir}/raw_data/{image_path[:-4]}_carved.png') | |
| ### Test | |
| # mask = carved_image[..., -1] > 0 | |
| # image = recenter(carved_image) | |
| # image = cv2.resize(image, (256, 256), interpolation=cv2.INTER_AREA) | |
| # Image.fromarray(image).save(f'{base_dir}/raw_data/{image_path[:-4]}_recentered.png') | |
| # image = cv2.resize(image, (280, 280), interpolation=cv2.INTER_AREA) | |
| # image = recover(image, (720, 480)) | |
| # Image.fromarray(image).save(f'{base_dir}/raw_data/{image_path[:-4]}_recovered.png') | |