File size: 2,877 Bytes
6bc98b5 c5c2233 6bc98b5 d56b3fc 6bc98b5 3ae25b7 6bc98b5 3ae25b7 6bc98b5 3ae25b7 6bc98b5 bcbdf9f 6bc98b5 6d5ddb4 6bc98b5 729b4bd 6bc98b5 c5c2233 6bc98b5 c5c2233 6bc98b5 c5c2233 2f64521 6bc98b5 02401bf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
#!/usr/bin/env python3
"""
square_crop.py
Crop each COCO bounding box to the smallest square that contains it.
"""
import argparse
import json
from pathlib import Path
import os
from PIL import Image, ImageOps
from tqdm import tqdm
def load_coco(json_path):
with open(json_path, "r") as f:
coco = json.load(f)
id2fname = {img["id"]: img["filename"] for img in coco["images"]}
return id2fname, coco["annotations"]
def square_from_bbox(x1, y1, x2, y2, img_w, img_h):
"""
Compute (left, top, side) of the smallest square fully containing the bbox.
The square is centred on the bbox; if it overflows the image, it is shifted
(but not resized) so it lies inside the image. Returns the final crop box
(left, top, side).
"""
side = max(x2-x1, y2-y1)
cx, cy = (x1+x2) / 2.0, (y1+y2) / 2.0
left = int(round(cx - side / 2.0))
top = int(round(cy - side / 2.0))
# Shift the square so it fits inside the image
left = max(0, min(left, img_w - side))
top = max(0, min(top, img_h - side))
return left, top, int(side)
def crop_annotation(img_path, ann, out_dir, pad_color=0):
with Image.open(img_path) as img:
img_w, img_h = img.size
x1, y1, x2, y2 = ann["bbox"]
left, top, side = square_from_bbox(x1, y1, x2, y2, img_w, img_h)
# Perform crop (may be smaller than 'side' at edges)
crop = img.crop((left, top, left + side, top + side))
print(f"final image {crop.size} from original image {img.size}")
# If we lost pixels at the edge, pad back to full square
if crop.size != (side, side):
delta_w = side - crop.size[0]
delta_h = side - crop.size[1]
padding = (0, 0, delta_w, delta_h) # (left, top, right, bottom)
crop = ImageOps.expand(crop, padding, fill=pad_color)
# resize smaller so it takes less time to load
crop = crop.resize((640,640))
# Build output filename: <stem>_ann<id>.ext
stem = Path(img_path).stem
suffix = Path(img_path).suffix
out_name = f"{stem}_ann{ann['id']}{suffix}"
crop.save(os.path.join(out_dir, out_name))
def run_square_crop(input_dir, coco_json_path, cropped_dir):
id2fname, annotations = load_coco(coco_json_path)
# Group annotations by image for efficient loading
im2anns = {}
for ann in annotations:
im2anns.setdefault(ann["image_id"], []).append(ann)
for img_id, anns in tqdm(im2anns.items(), desc="Processing images"):
img_path = os.path.join(input_dir, id2fname[img_id])
for ann in anns:
# try:
crop_annotation(img_path, ann, cropped_dir)
# except Exception as e:
# print(f"Error on {img_path}.")
return [os.path.join(cropped_dir, fname) for fname in sorted(os.listdir(cropped_dir))] |