Upload square_crop.py

square_crop.py

@@ -0,0 +1,109 @@
+#!/usr/bin/env python3
+"""
+square_crop.py
+
+Crop each COCO bounding box to the smallest square that contains it.
+
+Usage
+-----
+    python square_crop.py --images_dir <IMG_DIR> \
+                          --coco_json  <COCO_JSON> \
+                          --output_dir <OUT_DIR>
+
+Dependencies
+------------
+    pip install pillow pycocotools tqdm
+"""
+import argparse
+import json
+from pathlib import Path
+
+from PIL import Image, ImageOps
+from tqdm import tqdm
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Crop COCO bounding boxes to squares.")
+    parser.add_argument("--images_dir", required=True, type=Path,
+                        help="Directory containing the source images.")
+    parser.add_argument("--coco_json", required=True, type=Path,
+                        help="Path to the COCO annotation file (JSON).")
+    parser.add_argument("--output_dir", required=True, type=Path,
+                        help="Directory to write the cropped images.")
+    parser.add_argument("--pad_color", default=0, type=int,
+                        help="Gray-scale padding color 0–255 (default 0 = black).")
+    return parser.parse_args()
+
+
+def load_coco(json_path):
+    with open(json_path, "r") as f:
+        coco = json.load(f)
+    id2fname = {img["id"]: img["file_name"] for img in coco["images"]}
+    return id2fname, coco["annotations"]
+
+
+def square_from_bbox(x, y, w, h, 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(w, h)
+    cx, cy = x + w / 2.0, y + h / 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
+        x, y, w, h = ann["bbox"]          # COCO bbox = [x, y, width, height]
+        left, top, side = square_from_bbox(x, y, w, h, img_w, img_h)
+
+        # Perform crop (may be smaller than 'side' at edges)
+        crop = img.crop((left, top, left + side, top + side))
+
+        # 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)
+
+        # 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(out_dir / out_name)
+
+
+def main():
+    args = parse_args()
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+
+    id2fname, annotations = load_coco(args.coco_json)
+
+    # 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 = args.images_dir / id2fname[img_id]
+        if not img_path.is_file():
+            print(f"Warning: image {img_path} not found — skipping.")
+            continue
+        for ann in anns:
+            crop_annotation(img_path, ann, args.output_dir,
+                            pad_color=args.pad_color)
+
+
+if __name__ == "__main__":
+    main()