diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..00983fe --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +/__pycache__/ \ No newline at end of file diff --git a/README.md b/README.md index fea4135..79bf12b 100644 --- a/README.md +++ b/README.md @@ -1 +1,24 @@ -First Commit \ No newline at end of file +### Start + +See the `conf.py`: +``` +scheme = 'B' + +schemeA = { + 'image_name': 'd1.jpg', + 'is_lama': False, + 'inpaint_radius': 3, + 'is_gaussianblur': True, + 'gaussian_radius': 9 # Odd number +} + +schemeB = { + 'image_name': 'd1.jpg', + 'gaussian_radius': 51, # Odd number + 'is_use_fill_color': False, + 'fill_color': [227, 234, 244] +} +``` +You can choose scheme whcih you want. +The A is useof openCV inpaint or model(LaMa); And B is useof openCV but no inpaint. +When you finished the conf, run `python main.py`. \ No newline at end of file diff --git a/conf.py b/conf.py new file mode 100644 index 0000000..fbe7e64 --- /dev/null +++ b/conf.py @@ -0,0 +1,16 @@ +scheme = 'B' + +schemeA = { + 'image_name': 'd1.jpg', + 'is_lama': False, + 'inpaint_radius': 3, + 'is_gaussianblur': True, + 'gaussian_radius': 9 # Odd number +} + +schemeB = { + 'image_name': 'd2.jpg', + 'gaussian_radius': 51, # Odd number + 'is_use_fill_color': False, + 'fill_color': [227, 234, 244] +} \ No newline at end of file diff --git a/images/d1.jpg b/images/d1.jpg new file mode 100644 index 0000000..2525389 Binary files /dev/null and b/images/d1.jpg differ diff --git a/images/d2.jpg b/images/d2.jpg new file mode 100644 index 0000000..98546ec Binary files /dev/null and b/images/d2.jpg differ diff --git a/images/d3.jpg b/images/d3.jpg new file mode 100644 index 0000000..cc8b8e1 Binary files /dev/null and b/images/d3.jpg differ diff --git a/main.py b/main.py new file mode 100644 index 0000000..c550238 --- /dev/null +++ b/main.py @@ -0,0 +1,28 @@ +import conf +import subprocess + + +scheme = conf.scheme + + +def run_script(script_name): + try: + subprocess.run(['python', script_name], check=True) + print(f"{script_name} success!") + + except subprocess.CalledProcessError as e: + print(f'Run {script_name} error: {e}') + + except FileNotFoundError: + print(f'Can not found: {script_name}') + + +if __name__ == '__main__': + if scheme == 'A': + run_script("schemeA.py") + + elif scheme == 'B': + run_script("schemeB.py") + + else: + print('请检查配置文件conf.py中的scheme字段') \ No newline at end of file diff --git a/outputs/output_A_d1.jpg b/outputs/output_A_d1.jpg new file mode 100644 index 0000000..53cce33 Binary files /dev/null and b/outputs/output_A_d1.jpg differ diff --git a/outputs/output_A_d2.jpg b/outputs/output_A_d2.jpg new file mode 100644 index 0000000..b7243cf Binary files /dev/null and b/outputs/output_A_d2.jpg differ diff --git a/outputs/output_B_d1.jpg b/outputs/output_B_d1.jpg new file mode 100644 index 0000000..46d3721 Binary files /dev/null and b/outputs/output_B_d1.jpg differ diff --git a/outputs/output_B_d2.jpg b/outputs/output_B_d2.jpg new file mode 100644 index 0000000..39f3095 Binary files /dev/null and b/outputs/output_B_d2.jpg differ diff --git a/outputs/output_lama_A_d1.jpg b/outputs/output_lama_A_d1.jpg new file mode 100644 index 0000000..00b815b Binary files /dev/null and b/outputs/output_lama_A_d1.jpg differ diff --git a/outputs/output_lama_A_d2.jpg b/outputs/output_lama_A_d2.jpg new file mode 100644 index 0000000..e8c8f81 Binary files /dev/null and b/outputs/output_lama_A_d2.jpg differ diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..f106479 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,7 @@ +numpy==2.2.4 +opencv_python==4.11.0.86 +opencv_python_headless==4.11.0.86 +Pillow==9.5.0 +Pillow==11.1.0 +simple_lama_inpainting==0.1.2 +ultralytics==8.3.85 diff --git a/schemeA.py b/schemeA.py new file mode 100644 index 0000000..58d4969 --- /dev/null +++ b/schemeA.py @@ -0,0 +1,93 @@ +import cv2 +from PIL import Image +import numpy as np +from ultralytics import YOLO +from simple_lama_inpainting import SimpleLama +import conf + + +image_name = conf.schemeA['image_name'] +is_lama = conf.schemeA['is_lama'] +inpaint_radius = conf.schemeA['inpaint_radius'] +is_gaussianblur = conf.schemeA['is_gaussianblur'] +gaussian_radius = conf.schemeA['gaussian_radius'] + + +model = YOLO('yolov5s.pt') + + +if is_lama == True: + simple_lama = SimpleLama() + + +def imwrite(image): + output_path = 'outputs/' + 'output_lama_A_' + image_name if is_lama == True else 'outputs/' + 'output_A_' + image_name + cv2.imwrite(output_path, image) + print(f'Person has been removed, and the processed image has been saved in ./{output_path}') + + +def use_lama(box, image, mask): + if is_gaussianblur == True: + mask_blurred = cv2.GaussianBlur(mask, (gaussian_radius, gaussian_radius), 0) + inpainted_image = simple_lama(image, mask_blurred) + else: + inpainted_image = simple_lama(image, mask) + + if isinstance(inpainted_image, Image.Image): + inpainted_image = np.array(inpainted_image) # make PIL.Image to numpy + + inpainted_image_bgr = cv2.cvtColor(inpainted_image, cv2.COLOR_RGB2BGR) # RGB to BGR + imwrite(inpainted_image_bgr) + + +def use_opencv(box, image, mask): + if is_gaussianblur == True: + mask_blurred = cv2.GaussianBlur(mask, (gaussian_radius, gaussian_radius), 0) + + inpainted_image = cv2.inpaint(image, mask_blurred, inpaintRadius=inpaint_radius, flags=cv2.INPAINT_TELEA) + + mask_edges = cv2.Canny(mask_blurred, 100, 200) + inpainted_image = cv2.inpaint(inpainted_image, mask_edges, inpaintRadius=inpaint_radius, flags=cv2.INPAINT_TELEA) + else: + inpainted_image = cv2.inpaint(image, mask, inpaintRadius=inpaint_radius, flags=cv2.INPAINT_TELEA) + + imwrite(inpainted_image) + + +def inpaint_mask(box, image, mask): + if is_lama == True: + use_lama(box, image, mask) + else: + use_opencv(box, image, mask) + + +def get_mask(box): + x1, y1, x2, y2 = map(int, box.xyxy[0]) + mask = np.zeros(image.shape[:2], dtype=np.uint8) + mask[y1:y2, x1:x2] = 255 + return mask + + +def detect_person(image, box): + if model.names[int(box.cls)] == 'person': + mask = get_mask(box) + inpaint_mask(box, image, mask) + + +def detect_results(results, image): + for result in results: + boxes = result.boxes + + for box in boxes: + detect_person(image, box) + + +if __name__ == '__main__': + image = cv2.imread('./images/' + image_name) + if is_lama == True: + image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # SimpleLama need to RGB + results = model(image) + else: + results = model(image) + + detect_results(results, image) \ No newline at end of file diff --git a/schemeB.py b/schemeB.py new file mode 100644 index 0000000..b87177f --- /dev/null +++ b/schemeB.py @@ -0,0 +1,69 @@ +import cv2 +import numpy as np +from ultralytics import YOLO +import conf + + +image_name = conf.schemeB['image_name'] +gaussian_radius = conf.schemeB['gaussian_radius'] +is_use_fill_color = conf.schemeB['is_use_fill_color'] +fill_color_c = conf.schemeB['fill_color'] + +model = YOLO('yolov5s.pt') + + +def imwrite(image): + output_path = 'outputs/output_B_' + image_name + cv2.imwrite(output_path, image) + + print(f'Person has been removed, and the processed image has been saved in ./{output_path}') + + +def get_blended_image(image, blurred_mask): + fill_color = None + if is_use_fill_color: + fill_color = fill_color_c + else: + mean_color = np.mean(image, axis=0) + fill_color = mean_color + + fill_image = np.zeros_like(image) + fill_image[:] = fill_color + + blended_image = (image * (1 - blurred_mask[..., np.newaxis]) + + fill_image * blurred_mask[..., np.newaxis]).astype(np.uint8) + + return blended_image + + +def detect_person(image, box): + if model.names[int(box.cls)] == 'person': + x1, y1, x2, y2 = map(int, box.xyxy[0]) + + mask = np.zeros(image.shape[:2], dtype=np.uint8) + mask[y1:y2, x1:x2] = 255 + + rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # BGR to RGB + + blurred_mask = cv2.GaussianBlur(mask, (gaussian_radius, gaussian_radius), 0) + blurred_mask = blurred_mask / 255.0 + blended_image = get_blended_image(rgb_image, blurred_mask) + + bgr_image = cv2.cvtColor(blended_image, cv2.COLOR_RGB2BGR) # RGB to BGR + + imwrite(bgr_image) + + +def detect_results(results, image): + for result in results: + boxes = result.boxes + + for box in boxes: + detect_person(image, box) + + +if __name__ == '__main__': + image = cv2.imread('./images/' + image_name) + results = model(image) + + detect_results(results, image) \ No newline at end of file diff --git a/yolov5su.pt b/yolov5su.pt new file mode 100644 index 0000000..e6cba26 Binary files /dev/null and b/yolov5su.pt differ