""" Module for preprocessing and data augmentation of the BCCD dataset. """ import os import shutil import cv2 import numpy as np import xml.etree.ElementTree as ET from tqdm import tqdm import albumentations as A from PIL import Image def preprocess_dataset(dataset_path): """ Preprocesses the BCCD dataset images. Args: dataset_path: Path to the BCCD dataset Returns: Path to the preprocessed dataset """ print("Preprocessing dataset...") # Create output directory output_dir = "preprocessed_dataset" os.makedirs(output_dir, exist_ok=True) os.makedirs(os.path.join(output_dir, "images"), exist_ok=True) os.makedirs(os.path.join(output_dir, "annotations"), exist_ok=True) # Get paths image_path = os.path.join(dataset_path, "BCCD", "JPEGImages") annot_path = os.path.join(dataset_path, "BCCD", "Annotations") # Get all image files image_files = [f for f in os.listdir(image_path) if f.endswith(('.jpg', '.jpeg', '.png'))] for file in tqdm(image_files, desc="Preprocessing images"): # Read image img = cv2.imread(os.path.join(image_path, file)) if img is None: continue # Apply preprocessing # 1. Convert to RGB img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 2. Normalize img = img / 255.0 # 3. Resize to a standard size if needed img = cv2.resize(img, (640, 640)) # 4. Convert back to 0-255 range and BGR for saving img = (img * 255).astype(np.uint8) img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) # Save the preprocessed image cv2.imwrite(os.path.join(output_dir, "images", file), img) # Copy the annotation file base_name = os.path.splitext(file)[0] xml_file = os.path.join(annot_path, base_name + ".xml") if os.path.exists(xml_file): shutil.copy(xml_file, os.path.join(output_dir, "annotations", base_name + ".xml")) print(f"Preprocessing completed. Saved to {output_dir}") return output_dir def augment_dataset(dataset_path): """ Applies data augmentation to the preprocessed dataset. Args: dataset_path: Path to the preprocessed dataset Returns: Path to the augmented dataset """ print("Augmenting dataset...") # Create output directory output_dir = "augmented_dataset" os.makedirs(output_dir, exist_ok=True) os.makedirs(os.path.join(output_dir, "images"), exist_ok=True) os.makedirs(os.path.join(output_dir, "annotations"), exist_ok=True) # Copy original data first image_path = os.path.join(dataset_path, "images") annot_path = os.path.join(dataset_path, "annotations") for file in os.listdir(image_path): shutil.copy(os.path.join(image_path, file), os.path.join(output_dir, "images", file)) for file in os.listdir(annot_path): shutil.copy(os.path.join(annot_path, file), os.path.join(output_dir, "annotations", file)) # Define augmentation pipeline augmentations = [ A.Compose([ A.HorizontalFlip(p=1.0), A.BBoxParams(format='pascal_voc', label_fields=['class_labels']) ]), A.Compose([ A.RandomBrightnessContrast(p=1.0), A.BBoxParams(format='pascal_voc', label_fields=['class_labels']) ]), A.Compose([ A.Rotate(limit=20, p=1.0), A.BBoxParams(format='pascal_voc', label_fields=['class_labels']) ]), A.Compose([ A.RandomSizedBBoxSafeCrop(width=640, height=640, p=1.0), A.BBoxParams(format='pascal_voc', label_fields=['class_labels']) ]) ] # Get all image files image_files = [f for f in os.listdir(image_path) if f.endswith(('.jpg', '.jpeg', '.png'))] for file in tqdm(image_files, desc="Augmenting images"): # Read image img_path = os.path.join(image_path, file) img = cv2.imread(img_path) if img is None: continue # Convert to RGB img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # Read annotation base_name = os.path.splitext(file)[0] xml_path = os.path.join(annot_path, base_name + ".xml") if not os.path.exists(xml_path): continue # Parse XML to get bounding boxes tree = ET.parse(xml_path) root = tree.getroot() bboxes = [] class_labels = [] for obj in root.findall('object'): cls = obj.find('name').text bbox = obj.find('bndbox') xmin = int(float(bbox.find('xmin').text)) ymin = int(float(bbox.find('ymin').text)) xmax = int(float(bbox.find('xmax').text)) ymax = int(float(bbox.find('ymax').text)) bboxes.append([xmin, ymin, xmax, ymax]) class_labels.append(cls) # Apply each augmentation for i, aug in enumerate(augmentations): # Apply augmentation try: augmented = aug(image=img, bboxes=bboxes, class_labels=class_labels) aug_img = augmented['image'] aug_bboxes = augmented['bboxes'] aug_labels = augmented['class_labels'] # Skip if no bounding boxes are left after augmentation if len(aug_bboxes) == 0: continue # Create a new XML file for the augmented image new_file_name = f"{base_name}_aug_{i}.jpg" new_xml_name = f"{base_name}_aug_{i}.xml" # Save the augmented image aug_img_bgr = cv2.cvtColor(aug_img, cv2.COLOR_RGB2BGR) cv2.imwrite(os.path.join(output_dir, "images", new_file_name), aug_img_bgr) # Create a new XML annotation create_xml_annotation( os.path.join(output_dir, "annotations", new_xml_name), new_file_name, aug_img.shape[1], # width aug_img.shape[0], # height aug_bboxes, aug_labels ) except Exception as e: print(f"Error augmenting {file} with {i}: {e}") continue print(f"Augmentation completed. Saved to {output_dir}") return output_dir def create_xml_annotation(path, filename, width, height, bboxes, labels): """ Creates an XML annotation file in Pascal VOC format. Args: path: Path to save the XML file filename: Image filename width: Image width height: Image height bboxes: List of bounding boxes [xmin, ymin, xmax, ymax] labels: List of class labels """ root = ET.Element("annotation") # Add image information folder = ET.SubElement(root, "folder") folder.text = "images" file_node = ET.SubElement(root, "filename") file_node.text = filename size = ET.SubElement(root, "size") width_node = ET.SubElement(size, "width") width_node.text = str(width) height_node = ET.SubElement(size, "height") height_node.text = str(height) depth = ET.SubElement(size, "depth") depth.text = "3" # Add object information for (xmin, ymin, xmax, ymax), label in zip(bboxes, labels): obj = ET.SubElement(root, "object") name = ET.SubElement(obj, "name") name.text = label bndbox = ET.SubElement(obj, "bndbox") xmin_node = ET.SubElement(bndbox, "xmin") xmin_node.text = str(int(xmin)) ymin_node = ET.SubElement(bndbox, "ymin") ymin_node.text = str(int(ymin)) xmax_node = ET.SubElement(bndbox, "xmax") xmax_node.text = str(int(xmax)) ymax_node = ET.SubElement(bndbox, "ymax") ymax_node.text = str(int(ymax)) # Write to file tree = ET.ElementTree(root) tree.write(path)