Upload 2 files

config.py

@@ -0,0 +1,13 @@
+# config.py
+
+# Variables for renamer.py
+BASE_DIR = r"E:\studio Dropbox\studio\projects\TheFold\3models\SachaQuenby\1prep\2clean"
+CROP_CATEGORY = ["1closeup", "2mid", "3full"]
+STYLE_CATEGORIES = ["4fashion", "5natural", "6beauty", "7catwalk"]
+
+# Variables for cropper.py
+INPUT_FOLDER = r"E:\studio Dropbox\studio\projects\TheFold\3models\SachaQuenby\1prep\2crop"
+OUTPUT_FOLDER = r"E:\studio Dropbox\studio\projects\TheFold\3models\SachaQuenby\1prep\2crop\cropped"
+YOLO_DETECTED_FOLDER = r"path"
+ASPECT_RATIOS = [(1024, 1280), (1280, 1024), (1024, 1536), (1536, 1024), (1024, 1024)]
+SAVE_TO_YOLO_DETECTED_FOLDER = False

cropper-2.py

@@ -0,0 +1,89 @@
+import torch
+import cv2
+import os
+import threading
+from PIL import Image
+import sys
+sys.path.insert(0, 'E:/studio Dropbox/studio/ai/libs/notes')
+
+print(sys.executable)
+
+# Import variables from config.py
+from dataset_prep.config import INPUT_FOLDER, OUTPUT_FOLDER, YOLO_DETECTED_FOLDER, ASPECT_RATIOS, SAVE_TO_YOLO_DETECTED_FOLDER
+
+# Define the locks
+counter_lock = threading.Lock()
+model_lock = threading.Lock()
+
+def resize_bbox_to_dimensions(bbox, target_width, target_height, img_width, img_height):
+    x1, y1, x2, y2 = bbox
+    current_width = x2 - x1
+    current_height = y2 - y1
+    desired_aspect_ratio = target_width / target_height
+    current_aspect_ratio = current_width / current_height
+
+    print(f"Original bbox: {bbox}")
+    print(f"Current aspect ratio: {current_aspect_ratio}")
+    print(f"Desired aspect ratio: {desired_aspect_ratio}")
+
+    if current_aspect_ratio < desired_aspect_ratio:
+        new_width = desired_aspect_ratio * current_height
+        x1 -= (new_width - current_width) / 2
+        x2 += (new_width - current_width) / 2
+    elif current_aspect_ratio > desired_aspect_ratio:
+        new_height = current_width / desired_aspect_ratio
+        y1 -= (new_height - current_height) / 2
+        y2 += (new_height - current_height) / 2
+
+    x1 = max(x1, 0)
+    y1 = max(y1, 0)
+    x2 = min(x2, img_width)
+    y2 = min(y2, img_height)
+
+    new_bbox = [int(x1), int(y1), int(x2), int(y2)]
+    print(f"New bbox: {new_bbox}")
+
+    return new_bbox
+
+def process_files(filelist):
+    global image_processed_counter
+    with counter_lock:
+        model = torch.hub.load('WongKinYiu/yolov7', 'custom', 'yolov7-e6e.pt', force_reload=False, trust_repo=True)
+    for filename in filelist:
+        try:  # Start of the try block
+            img_path = os.path.join(INPUT_FOLDER, filename)
+            image = cv2.imread(img_path)
+            if image is None:
+                raise ValueError(f"Could not read image {filename}")
+            img_width, img_height = image.shape[1], image.shape[0]
+
+            with model_lock:
+                results = model(img_path)
+            detections = results.pandas().xyxy[0]
+
+            person_detected = detections[detections['name'] == 'person']
+            print(f"Person detected: {not person_detected.empty}")
+            if not person_detected.empty:
+                x1, y1, x2, y2 = person_detected.iloc[0][['xmin', 'ymin', 'xmax', 'ymax']].astype(int)
+
+                for target_width, target_height in ASPECT_RATIOS:
+                    new_x1, new_y1, new_x2, new_y2 = resize_bbox_to_dimensions([x1, y1, x2, y2], target_width, target_height, img_width, img_height)
+                    new_x1, new_y1 = max(new_x1, 0), max(new_y1, 0)
+                    new_x2, new_y2 = min(new_x2, img_width), min(new_y2, img_height)
+                    cropped_img = image[new_y1:new_y2, new_x1:new_x2]
+
+                    # Create a folder for each aspect ratio if it doesn't exist
+                    aspect_ratio_folder = f"{target_width}_{target_height}"
+                    aspect_ratio_path = os.path.join(OUTPUT_FOLDER, aspect_ratio_folder)
+                    os.makedirs(aspect_ratio_path, exist_ok=True)
+
+                    # Save the cropped image to the corresponding folder
+                    output_filename = os.path.join(aspect_ratio_path, f"cropped_{filename}")
+                    cv2.imwrite(output_filename, cropped_img)
+                  
+        except Exception as e:
+            print(f"An error occurred while processing file {filename}: {e}")
+
+if __name__ == "__main__":
+    filelist = os.listdir(INPUT_FOLDER)
+    process_files(filelist)