Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- import os
- import cv2
- import numpy as np
- import random
- from tqdm.notebook import tqdm # Progress bar
- from IPython.display import display, Markdown
- # Direktori input dan output
- INPUT_FOLDER = "assets/anotated" # Folder hasil anotasi sebelumnya
- OUTPUT_FOLDER = "assets/augmentasi" # Folder hasil augmentasi
- # Pastikan output folder tersedia
- os.makedirs(OUTPUT_FOLDER, exist_ok=True)
- # Fungsi flipping bounding box
- def flip_bounding_box(image_shape, bbox, flip_code):
- height, width = image_shape[:2]
- x_center, y_center, box_width, box_height, class_id = bbox
- if flip_code == 1: # Horizontal flip
- x_center = 1.0 - x_center
- elif flip_code == 0: # Vertical flip
- y_center = 1.0 - y_center
- elif flip_code == -1: # Horizontal + Vertical flip
- x_center = 1.0 - x_center
- y_center = 1.0 - y_center
- return x_center, y_center, box_width, box_height, class_id
- # Fungsi untuk membaca anotasi YOLO dalam setiap folder ID personal
- def load_annotations(folder):
- annotations = {}
- for person_id in os.listdir(folder):
- person_path = os.path.join(folder, person_id)
- if not os.path.isdir(person_path):
- continue # Skip jika bukan folder
- for filename in os.listdir(person_path):
- if filename.endswith('.txt'):
- path = os.path.join(person_path, filename)
- with open(path, "r") as f:
- lines = f.readlines()
- annotations[f"{person_id}/{filename.replace('.txt', '.jpg')}"] = [
- tuple(map(float, line.strip().split()[1:])) + (int(line.strip().split()[0]),)
- for line in lines
- ]
- return annotations
- # Fungsi menyimpan anotasi dalam format YOLO
- def save_augmented_annotations(filename, person_id, output_dir, annotations):
- person_output_folder = os.path.join(output_dir, person_id)
- os.makedirs(person_output_folder, exist_ok=True)
- txt_output_path = os.path.join(person_output_folder, filename.replace('.jpg', '.txt'))
- with open(txt_output_path, "w") as f:
- for x_center, y_center, width, height, class_id in annotations:
- f.write(f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}\n")
- # Fungsi berbagai augmentasi
- def apply_augmentations(image, annotations):
- augmented_images = []
- # 1. Flipping Horizontal
- flipped_h = cv2.flip(image, 1)
- flipped_h_ann = [flip_bounding_box(image.shape, ann, 1) for ann in annotations]
- augmented_images.append((flipped_h, flipped_h_ann, "flipH"))
- # 2. Rotation (-10 to +10 degrees)
- angle = random.randint(-10, 10)
- h, w = image.shape[:2]
- M = cv2.getRotationMatrix2D((w/2, h/2), angle, 1)
- rotated = cv2.warpAffine(image, M, (w, h))
- augmented_images.append((rotated, annotations, f"rotate{angle}"))
- # 3. Gaussian Blur
- blurred = cv2.GaussianBlur(image, (5, 5), 0)
- augmented_images.append((blurred, annotations, "blur"))
- # 4. Grayscale
- grayscale = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
- grayscale = cv2.cvtColor(grayscale, cv2.COLOR_GRAY2BGR)
- augmented_images.append((grayscale, annotations, "gray"))
- # 5. Brightness & Contrast
- alpha = random.uniform(0.8, 1.2) # Contrast
- beta = random.randint(-20, 20) # Brightness
- bright_contrast = cv2.convertScaleAbs(image, alpha=alpha, beta=beta)
- augmented_images.append((bright_contrast, annotations, f"bright{beta}_contrast{alpha:.1f}"))
- # 6. Noise Injection
- noise = np.random.normal(0, 10, image.shape).astype(np.uint8)
- noisy_image = cv2.add(image, noise)
- augmented_images.append((noisy_image, annotations, "noise"))
- # 7. Perspective Transform (Small Distortion)
- pts1 = np.float32([[0, 0], [w, 0], [0, h], [w, h]])
- shift_x = random.randint(-10, 10)
- shift_y = random.randint(-10, 10)
- pts2 = np.float32([[shift_x, shift_y], [w-shift_x, shift_y], [shift_x, h-shift_y], [w-shift_x, h-shift_y]])
- matrix = cv2.getPerspectiveTransform(pts1, pts2)
- perspective = cv2.warpPerspective(image, matrix, (w, h))
- augmented_images.append((perspective, annotations, "perspective"))
- return augmented_images
- # Fungsi utama augmentasi dengan progress bar
- def augment_images():
- person_folders = [f for f in os.listdir(INPUT_FOLDER) if os.path.isdir(os.path.join(INPUT_FOLDER, f))]
- annotations_dict = load_annotations(INPUT_FOLDER)
- if not person_folders:
- display(Markdown("❌ **Folder input kosong atau tidak ada gambar!**"))
- return
- total_files = sum(len(files) for _, _, files in os.walk(INPUT_FOLDER))
- success_count, failed_count = 0, 0
- progress = tqdm(person_folders, desc="🔄 Memproses Augmentasi", unit="folder", leave=False)
- for person_id in progress:
- person_path = os.path.join(INPUT_FOLDER, person_id)
- output_person_folder = os.path.join(OUTPUT_FOLDER, person_id)
- os.makedirs(output_person_folder, exist_ok=True)
- for filename in os.listdir(person_path):
- if filename.endswith(('.jpg', '.jpeg', '.png')):
- img_path = os.path.join(person_path, filename)
- image = cv2.imread(img_path)
- if image is None:
- failed_count += 1
- continue
- annotations = annotations_dict.get(f"{person_id}/{filename}", [])
- # Melakukan augmentasi
- augmented_images = apply_augmentations(image, annotations)
- for aug_image, aug_annotations, aug_type in augmented_images:
- aug_filename = f"{os.path.splitext(filename)[0]}_{aug_type}.jpg"
- cv2.imwrite(os.path.join(output_person_folder, aug_filename), aug_image)
- save_augmented_annotations(aug_filename, person_id, OUTPUT_FOLDER, aug_annotations)
- success_count += 1
- progress.close()
- display(Markdown(f"✅ **Proses augmentasi selesai!**"))
- display(Markdown(f"📂 **Total file input: {total_files}**"))
- # Jalankan augmentasi
- augment_images()
Add Comment
Please, Sign In to add comment
