Last active
April 7, 2025 07:01
-
-
Save motebaya/5254430cf32ac00bae4e87e26af99680 to your computer and use it in GitHub Desktop.
crop images from all results face detection from: https://github.com/zymk9/yolov5_anime
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/usr/bin/python | |
# crop image area with cv2 | |
# @github.com/motebaya - 7/10/2024 | |
import cv2 | |
import os | |
import logging | |
logging.basicConfig(level=logging.INFO) | |
from argparse import ArgumentParser, RawTextHelpFormatter | |
from colorama.ansi import Fore | |
import numpy as np | |
ext_list = [ | |
'.bmp', '.dib', '.jpeg', '.jpg', '.jpe', | |
'.jp2', '.png', '.webp', '.pbm', '.pgm', | |
'.ppm', '.sr', '.ras', '.tiff', '.tif', | |
'.exr', '.hdr', '.pic' | |
] | |
def crop_img( | |
img: str, x: int | float, y: int | float, w: int | float, h: int | float, margin: float = 0.5 | |
) -> None: | |
im = cv2.imdecode(np.fromfile(img, np.uint8), cv2.IMREAD_COLOR) | |
width, height = im.shape[1], im.shape[0] | |
x1 = int((x - w / 2) * width) | |
y1 = int((y - h / 2) * height) | |
x2 = int((x + w / 2) * width) | |
y2 = int((y + h / 2) * height) | |
margin_w, margin_h = int(margin * (x2 - x1)), int(margin * (y2 - y1)) # margin extend default 50%, from yolo it's too small! | |
croped_outp = os.path.join(os.path.dirname(img), "cropped") | |
if not os.path.exists(croped_outp): | |
os.makedirs(croped_outp, exist_ok=True) | |
outp = f"{croped_outp}/crop_{os.path.basename(img)}" | |
fname, ext = os.path.splitext(outp) | |
if os.path.exists(outp): | |
outp = f"{fname}_1{ext}" | |
# cv2 cannot read unicode filename | |
is_ok, buffer = cv2.imencode( | |
ext, | |
im[max(y1 - margin_h, 0):min(y2 + margin_h, height), max(x1 - margin_w, 0):min(x2 + margin_w, width)] | |
) | |
if is_ok: | |
buffer.tofile(outp) | |
logging.info(f"Image cropped and saved to {outp}") | |
return | |
logging.error(f"{Fore.RED}Error cropping image: {img}") | |
return | |
def bulk_crop( | |
orig_target: str, | |
ratio_target: str, | |
margin: float = 0.5 | |
) -> None: | |
"""Bulk crop images based on ratio files. | |
:param orig_target: Path to the directory containing the original images. | |
:type orig_target: str | |
:param ratio_target: Path to the directory containing the ratio files. | |
Each ratio file should have cropping coordinates for the corresponding image. | |
:type ratio_target: str | |
:example | |
python detect.py --source inference/images --weights weights/yolov5x_anime.pt --output inference/output --save-txt | |
""" | |
if os.path.isdir(orig_target) and os.path.isdir(ratio_target): | |
img_list = os.listdir(orig_target) | |
for i, img in enumerate(img_list, 1): | |
img_name, ext = os.path.splitext(img) | |
if ext.lower() not in ext_list: | |
logging.warning(f"{Fore.YELLOW}Unsupported file type: {img_name}{ext}{Fore.RESET}") | |
continue | |
ratio = os.path.join(ratio_target, f"{img_name}.txt") | |
if not os.path.exists(ratio): | |
logging.warning(f"{Fore.YELLOW}Ratio file not found: {ratio}{Fore.RESET}") | |
continue | |
ratio_list = open(ratio, 'r').read().strip().splitlines() | |
for ix, rt in enumerate(ratio_list, 1): | |
x, y, w, h = rt.strip().split()[1:] | |
logging.info( | |
f" * Cropping:{ix} of {len(ratio_list)} RATIO - {i} of {len(img_list)} IMAGES - {img_name} - {x}, {y}, {w}, {h}" | |
) | |
try: | |
crop_img( | |
os.path.join(orig_target, img), | |
float(x), float(y), float(w), float(h), | |
margin=margin | |
) | |
except Exception as e: | |
logging.error(f"{Fore.RED}Error cropping {img_name}: {e}{Fore.RESET}") | |
continue | |
else: | |
logging.error(f"{Fore.RED}Invalid directories provided.{Fore.RESET}") | |
return | |
if __name__=="__main__": | |
parser = ArgumentParser( | |
description="Crop images based on coordinates.", | |
formatter_class=RawTextHelpFormatter | |
) | |
parser.add_argument( | |
'-i', '--input', type=str, help="Path to the input images.", metavar="" | |
) | |
parser.add_argument( | |
"-r", "--ratio", type=str, help="Path to the ratio files.", metavar="" | |
) | |
parser.add_argument( | |
"-m", "--margin", type=float, default=0.5, help="Optional Extend Margin for cropping.", metavar="" | |
) | |
args = parser.parse_args() | |
if not args.input or not args.ratio: | |
parser.print_help() | |
exit(1) | |
bulk_crop( | |
args.input, | |
args.ratio, | |
args.margin or 0.5 | |
) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment