fix: added image width/height and fixed multiproc to square

sdcat/cluster/cluster.py

@@ -252,7 +252,7 @@ def cluster_vits(
             info('ROI crops already exist. Creating square crops in parallel using {multiprocessing.cpu_count()} processes...')
             with multiprocessing.Pool(num_processes) as pool:
                 args = [(row, 224) for index, row in df_dets.iterrows()]
-                pool.starmap(crop_square_image, args)
+                pool.starmap(square_image, args)
         else:
             # Crop and squaring the images in parallel using multiprocessing to speed up the processing
             info(f'Cropping {len(df_dets)} detections in parallel using {multiprocessing.cpu_count()} processes...')

sdcat/cluster/commands.py

@@ -14,6 +14,7 @@
 import pandas as pd
 import pytz
 import torch
+from PIL.Image import Image
 
 from sdcat import common_args
 from sdcat.config import config as cfg
@@ -318,6 +319,12 @@ def run_cluster_roi(roi_dir, save_dir, device, config_ini, alpha, cluster_select
     # Sort the dataframe by image_path to make sure the images are in order for start_image and end_image filtering
     df = df.sort_values(by='image_path')
 
+    # Add the image_width and image_height columns to the dataframe
+    for index, row in df.iterrows():
+        im_size = Image.open(row['image_path']).size
+        df.at[index, 'image_width'] = im_size[0]
+        df.at[index, 'image_height'] = im_size[1]
+
     # Create a unique crop name for each detection with a unique id
     crop_path = save_dir / 'crops'
     crop_path.mkdir(parents=True, exist_ok=True)