Update square_crop.py

square_crop.py

@@ -4,38 +4,15 @@ 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
+import os
 
 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)
@@ -84,11 +61,8 @@ def crop_annotation(img_path, ann, out_dir, pad_color=0):
         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)
+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 = {}
@@ -96,14 +70,12 @@ def main():
         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)
+        img_path = os.path.join(input_dir, id2fname[img_id])
 
+        for ann in anns:
+            try:
+                crop_annotation(img_path, ann, cropped_dir)
+            except error:
+                print(f"Error on {img_path}.")
 
-if __name__ == "__main__":
-    main()
+    return os.listdir(cropped_dir)