From 60686d4c1729170c050baa3817199dbf406da4cc Mon Sep 17 00:00:00 2001 From: Forough Koohi <87902202+foroughkoohi@users.noreply.github.com> Date: Sat, 7 Jun 2025 15:39:08 +0330 Subject: [PATCH 1/3] zip commit --- 1_dataset_generation.py | 366 ++++++++++++++++++++++++++++++++++++++++ 2_data_analysis.py | 104 ++++++++++++ 3_data_augmentation.py | 67 ++++++++ 4_yolo_annotation.py | 128 ++++++++++++++ 5_model_training.py | 277 ++++++++++++++++++++++++++++++ 5 files changed, 942 insertions(+) create mode 100644 1_dataset_generation.py create mode 100644 2_data_analysis.py create mode 100644 3_data_augmentation.py create mode 100644 4_yolo_annotation.py create mode 100644 5_model_training.py diff --git a/1_dataset_generation.py b/1_dataset_generation.py new file mode 100644 index 0000000..a613974 --- /dev/null +++ b/1_dataset_generation.py @@ -0,0 +1,366 @@ +import carla +import random +import time +import os +import numpy as np +import cv2 +from datetime import datetime +from collections import defaultdict +import queue + +# Initial settings +DATA_DIR = os.path.join(os.path.dirname(__file__), "data") +os.makedirs(os.path.join(DATA_DIR, "rgb"), exist_ok=True) +os.makedirs(os.path.join(DATA_DIR, "semantics"), exist_ok=True) +os.makedirs(os.path.join(DATA_DIR, "analysis"), exist_ok=True) + +# Weather conditions +weather_presets = { + 'Clear_Day': carla.WeatherParameters.ClearNoon, + 'Clear_Night': carla.WeatherParameters.ClearNight, + 'Heavy_Rain': carla.WeatherParameters.HardRainNoon, + 'Foggy': carla.WeatherParameters(fog_density=80.0) +} + +# Class mapping for instance segmentation +class_mapping = { + 114: "Car", # vehicle.car + 115: "Pedestrian", # walker.pedestrian + 122: "TrafficLight", # traffic.traffic_light + 177: "Bus" # vehicle.bus +} + +def save_image(image, path): + array = np.frombuffer(image.raw_data, dtype=np.dtype("uint8")) + array = np.reshape(array, (image.height, image.width, 4)) + + if 'semantics' in path: + seg_data = array[:, :, :3] + np.save(path.replace('.png', '.npy'), seg_data) + else: + rgb_data = array[:, :, :3][:, :, ::-1] + cv2.imwrite(path, rgb_data) + +def extract_instance_masks(seg_array): + """Extract instance masks from instance segmentation data""" + # Red channel: class ID, Green and Blue channels: instance ID + class_ids = seg_array[:, :, 0] + instance_ids = (seg_array[:, :, 1].astype(np.uint16) << 8) + seg_array[:, :, 2].astype(np.uint16) + + # Create masks for each unique instance + unique_instances = np.unique(instance_ids) + masks = [] + for inst_id in unique_instances: + if inst_id == 0: # Skip background + continue + mask = (instance_ids == inst_id).astype(np.uint8) + class_id = class_ids[mask > 0][0] # Assume all pixels in mask share the same class + masks.append((class_id, mask)) + return masks + +def create_dense_traffic_scene(world, traffic_manager, blueprint_library, zone_center, radius=40): + """Create a dense scene with cars, buses, pedestrians and traffic lights""" + vehicles_list = [] + walkers_list = [] + lights_list = [] + + vehicle_types = { + "vehicle.tesla.model3": 4, + "vehicle.audi.a2": 3, + "vehicle.dodge.charger_police": 2, + "vehicle.toyota.prius": 3, + "vehicle.volkswagen.t2": 1, + "vehicle.mercedes.sprinter": 2, + "vehicle.carlamotors.carlacola": 1 + } + + all_spawn_points = world.get_map().get_spawn_points() + random.shuffle(all_spawn_points) + used_spawn_points = set() + + # Spawn vehicles + for vehicle_type, weight in vehicle_types.items(): + is_bus = any(bus_id in vehicle_type for bus_id in ["t2", "sprinter", "carlacola"]) + spawn_count = weight * (2 if is_bus else 3) + for _ in range(spawn_count): + try: + spawn_point = min( + (sp for sp in all_spawn_points if sp not in used_spawn_points), + key=lambda sp: sp.location.distance(zone_center.location) + ) + used_spawn_points.add(spawn_point) + bp = blueprint_library.find(vehicle_type) + if bp.has_attribute('color'): + color = random.choice(bp.get_attribute('color').recommended_values) + bp.set_attribute('color', color) + vehicle = world.try_spawn_actor(bp, spawn_point) + if vehicle is not None: + vehicles_list.append(vehicle) + vehicle.set_autopilot(True, traffic_manager.get_port()) + if is_bus: + traffic_manager.distance_to_leading_vehicle(vehicle, random.uniform(4.0, 6.0)) + traffic_manager.vehicle_percentage_speed_difference(vehicle, random.uniform(-10, 0)) + else: + traffic_manager.distance_to_leading_vehicle(vehicle, random.uniform(1.0, 3.0)) + traffic_manager.vehicle_percentage_speed_difference(vehicle, random.uniform(-20, 20)) + traffic_manager.update_vehicle_lights(vehicle, True) + except Exception as e: + print(f"Vehicle spawn error: {e}") + continue + + # Create pedestrian groups + pedestrian_groups = 4 + pedestrians_per_group = 6 + for _ in range(pedestrian_groups): + group_center = carla.Location( + x=zone_center.location.x + random.uniform(-radius/2, radius/2), + y=zone_center.location.y + random.uniform(-radius/2, radius/2), + z=zone_center.location.z + ) + for _ in range(pedestrians_per_group): + try: + offset = random.uniform(1, 4) + angle = random.uniform(0, 2 * np.pi) + x = group_center.x + offset * np.cos(angle) + y = group_center.y + offset * np.sin(angle) + spawn_point = carla.Transform( + carla.Location(x=x, y=y, z=group_center.z + 1.0) + ) + bp = random.choice(blueprint_library.filter("walker.pedestrian.*")) + walker = world.try_spawn_actor(bp, spawn_point) + if walker is not None: + walkers_list.append(walker) + walker_control = carla.WalkerControl() + walker_control.speed = random.uniform(0.6, 2.0) + direction = carla.Vector3D( + x=group_center.x - x, + y=group_center.y - y, + z=0 + ) + direction.make_unit_vector() + walker_control.direction = direction + walker.apply_control(walker_control) + except Exception as e: + print(f"Walker spawn error: {e}") + continue + + # Configure traffic lights + try: + traffic_lights = world.get_actors().filter('traffic.traffic_light*') + for traffic_light in traffic_lights: + if traffic_light.get_location().distance(zone_center.location) < radius: + states = [carla.TrafficLightState.Red, + carla.TrafficLightState.Yellow, + carla.TrafficLightState.Green] + traffic_light.set_state(random.choice(states)) + traffic_light.set_green_time(random.uniform(4.0, 8.0)) + traffic_light.set_red_time(random.uniform(2.0, 6.0)) + traffic_light.set_yellow_time(random.uniform(1.0, 3.0)) + lights_list.append(traffic_light) + except Exception as e: + print(f"Traffic light error: {e}") + + return vehicles_list, walkers_list, lights_list + +def analyze_dataset(data_dir): + print("\n=== Dataset Analysis ===") + + rgb_files = sorted(os.listdir(os.path.join(data_dir, "rgb"))) + semantic_files = sorted([f for f in os.listdir(os.path.join(data_dir, "semantics")) if f.endswith('.npy')]) + + print(f"Total Samples: {len(rgb_files)}\n") + + weather_counts = defaultdict(int) + for f in rgb_files: + weather = f.split('_')[0] + weather_counts[weather] += 1 + + print("Samples per weather condition:") + for weather, count in weather_counts.items(): + print(f"{weather}: {count} samples") + + print("\nTotal Annotated Instances:") + instance_counts = defaultdict(int) + for sem_file in semantic_files: + seg_data = np.load(os.path.join(data_dir, "semantics", sem_file)) + masks = extract_instance_masks(seg_data) + for class_id, _ in masks: + if class_id in class_mapping: + instance_counts[class_id] += 1 + + for cls_id, count in instance_counts.items(): + print(f"{class_mapping[cls_id]} ({cls_id}): {count} instances") + +def analyze_instances(segmentation_files): + instance_counts = defaultdict(int) + class_pixel_counts = defaultdict(int) + + for seg_file in segmentation_files: + seg_data = np.load(seg_file) + masks = extract_instance_masks(seg_data) + + for class_id, mask in masks: + class_name = class_mapping[class_id] + instance_counts[class_name] += 1 + class_pixel_counts[class_name] += np.sum(mask) + + for class_name, count in instance_counts.items(): + avg_size = class_pixel_counts[class_name] / count if count > 0 else 0 + print(f"{class_name}:") + print(f" - Total instances: {count}") + print(f" - Total pixels: {class_pixel_counts[class_name]}") + print(f" - Average size: {avg_size:.1f} pixels") + +def validate_segmentation_data(seg_data): + """Validate segmentation data""" + if seg_data is None: + return False + + class_ids = np.unique(seg_data[:, :, 0]) + valid_classes = [cls_id for cls_id in class_ids if cls_id in class_mapping] + + if len(valid_classes) == 0: + print("No valid classes found in segmentation data") + return False + + return True + +def main(): + client = carla.Client("localhost", 2000) + client.set_timeout(10.0) + world = client.get_world() + + ego_vehicle = None + camera_rgb = None + camera_seg = None + vehicles_list = [] + walkers_list = [] + lights_list = [] + + try: + original_settings = world.get_settings() + settings = world.get_settings() + settings.synchronous_mode = True + settings.fixed_delta_seconds = 0.05 + world.apply_settings(settings) + + blueprint_library = world.get_blueprint_library() + ego_bp = blueprint_library.find("vehicle.tesla.model3") + + max_tries = 20 + spawn_points = world.get_map().get_spawn_points() + random.shuffle(spawn_points) + + for attempt in range(max_tries): + try: + spawn_point = spawn_points[attempt % len(spawn_points)] + ego_vehicle = world.spawn_actor(ego_bp, spawn_point) + print(f"Successfully spawned ego vehicle on attempt {attempt+1}") + break + except RuntimeError as e: + print(f"Failed attempt {attempt+1}/{max_tries} to spawn ego vehicle: {e}") + if attempt == max_tries - 1: + raise RuntimeError("Could not spawn ego vehicle") + + traffic_manager = client.get_trafficmanager(8000) + traffic_manager.set_synchronous_mode(True) + traffic_manager.set_global_distance_to_leading_vehicle(1.5) + traffic_manager.global_percentage_speed_difference(10.0) + ego_vehicle.set_autopilot(True, traffic_manager.get_port()) + + # Set up cameras + camera_rgb = world.spawn_actor( + blueprint_library.find("sensor.camera.rgb"), + carla.Transform(carla.Location(x=1.5, z=2.4)), + attach_to=ego_vehicle + ) + camera_seg = world.spawn_actor( + blueprint_library.find("sensor.camera.instance_segmentation"), + carla.Transform(carla.Location(x=1.5, z=2.4)), + attach_to=ego_vehicle + ) + + global rgb_data, seg_data + rgb_data = None + seg_data = None + + def rgb_callback(image): + global rgb_data + rgb_data = image + + def seg_callback(image): + global seg_data + seg_data = image + + camera_rgb.listen(rgb_callback) + camera_seg.listen(seg_callback) + + frame_count = {w: 0 for w in weather_presets.keys()} + all_spawn_points = world.get_map().get_spawn_points() + + for weather_name, weather in weather_presets.items(): + world.set_weather(weather) + print(f"\nWeather set to: {weather_name}") + while frame_count[weather_name] < 1: + for vehicle in vehicles_list: + if vehicle and vehicle.is_alive: + vehicle.destroy() + for walker in walkers_list: + if walker and walker.is_alive: + walker.destroy() + vehicles_list.clear() + walkers_list.clear() + + valid_spawn_points = [sp for sp in all_spawn_points if sp.location.z > 0] + if not valid_spawn_points: + print("No valid spawn points available!") + break + + zone_center = random.choice(valid_spawn_points) + ego_vehicle.set_transform(zone_center) + + new_vehicles, new_walkers, new_lights = create_dense_traffic_scene( + world, traffic_manager, blueprint_library, zone_center + ) + vehicles_list.extend(new_vehicles) + walkers_list.extend(new_walkers) + lights_list.extend(new_lights) + + for _ in range(30): + world.tick() + + if rgb_data and seg_data: + filename = f"{weather_name}_{frame_count[weather_name]:05d}" + save_image(rgb_data, os.path.join(DATA_DIR, "rgb", f"rgb_{filename}.png")) + save_image(seg_data, os.path.join(DATA_DIR, "semantics", f"seg_{filename}.png")) + frame_count[weather_name] += 1 + print(f"Saved: {filename} | Count: {frame_count[weather_name]}") + + time.sleep(3.0) + + print("Data collection completed") + analyze_dataset(DATA_DIR) + + finally: + print("Cleaning up actors...") + if camera_rgb: + camera_rgb.stop() + if camera_seg: + camera_seg.stop() + if ego_vehicle: + ego_vehicle.set_autopilot(False) + ego_vehicle.destroy() + for vehicle in vehicles_list: + if vehicle and vehicle.is_alive: + vehicle.destroy() + for walker in walkers_list: + if walker and walker.is_alive: + walker.destroy() + world.apply_settings(original_settings) + + seg_data = np.load("path_to_segmentation_mask.npy") + print("Unique class IDs:", np.unique(seg_data[:, :, 0])) + print("Unique instance IDs:", np.unique((seg_data[:, :, 1].astype(np.uint16) << 8) | seg_data[:, :, 2])) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/2_data_analysis.py b/2_data_analysis.py new file mode 100644 index 0000000..b2bc359 --- /dev/null +++ b/2_data_analysis.py @@ -0,0 +1,104 @@ +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +from collections import defaultdict + +# Paths to your data directories +DATA_DIR = os.path.join(os.path.dirname(__file__), "data") +RGB_DIR = os.path.join(DATA_DIR, "rgb") +SEM_DIR = os.path.join(DATA_DIR, "semantics") + +# Ensure directories exist +if not os.path.isdir(RGB_DIR) or not os.path.isdir(SEM_DIR): + raise FileNotFoundError(f"RGB or semantics directory not found under {DATA_DIR}") + +# Mapping of class IDs to class names +CLASS_MAPPING = {0: "Pedestrian", 1: "Traffic Light", 2: "Car", 3: "Bus"} + +# Debug: inspect unique class and instance IDs across all instance segmentation masks +mask_files = [f for f in os.listdir(SEM_DIR) if f.lower().endswith('.npy')] +unique_class_ids_all = set() +unique_instance_ids_all = set() +non_empty_count = 0 +for fname in mask_files: + mask = np.load(os.path.join(SEM_DIR, fname)) # Shape: (height, width, 3) + class_ids = mask[:, :, 0] # R channel for class IDs + instance_ids = (mask[:, :, 1].astype(np.uint16) << 8) + mask[:, :, 2] # G << 8 + B for instance IDs + unique_class_ids = set(np.unique(class_ids)) + unique_instance_ids = set(np.unique(instance_ids)) + unique_class_ids_all.update(unique_class_ids) + unique_instance_ids_all.update(unique_instance_ids) + if len(unique_instance_ids - {0}) > 0: + non_empty_count += 1 +print(f"Debug: Total mask files = {len(mask_files)}, Non-empty masks = {non_empty_count}") +print(f"Debug: All unique class IDs across masks = {sorted(unique_class_ids_all)}") +print(f"Debug: All unique instance IDs across masks = {sorted(unique_instance_ids_all)}") +print("If class IDs don't match CLASS_MAPPING keys, adjust CLASS_MAPPING or mask generation.") + +# 1) Count number of samples per weather condition +weather_counts = defaultdict(int) +for fname in os.listdir(RGB_DIR): + if fname.lower().startswith("rgb_") and fname.lower().endswith(('.png', '.jpg', '.jpeg')): + parts = fname.split("_") + weather = "_".join(parts[1:-1]) + weather_counts[weather] += 1 + +# Convert to DataFrame +df_weather = pd.DataFrame(sorted(weather_counts.items()), columns=["Weather", "NumSamples"]) + +# 2) Count object instances per class across all instance segmentation masks +instance_counts = defaultdict(int) +for fname in mask_files: + mask_path = os.path.join(SEM_DIR, fname) + mask = np.load(mask_path) # Shape: (height, width, 3) + class_ids = mask[:, :, 0] # R channel for class IDs + instance_ids = (mask[:, :, 1].astype(np.uint16) << 8) + mask[:, :, 2] # G << 8 + B for instance IDs + unique_instances = np.unique(instance_ids) + for inst_id in unique_instances: + if inst_id == 0: # Skip background + continue + inst_mask = (instance_ids == inst_id) + if inst_mask.any(): + cls_id = class_ids[inst_mask][0] + if cls_id in CLASS_MAPPING: + instance_counts[CLASS_MAPPING[cls_id]] += 1 + +# Convert to DataFrame +df_instances = pd.DataFrame(sorted(instance_counts.items()), columns=["Class", "NumInstances"]) + +# 3) Total annotated instances +total_instances = df_instances["NumInstances"].sum() if not df_instances.empty else 0 + +# 4) Print summary +print("\n=== Dataset Analysis Report ===") +print("Samples per Weather Condition:") +print(df_weather.to_string(index=False)) +print("\nInstances per Class:") +if df_instances.empty: + print("No instances found. Check CLASS_MAPPING IDs and mask generation.") +else: + print(df_instances.to_string(index=False)) +print(f"\nTotal Annotated Instances: {total_instances}\n") + +# 5) Visualizations +plt.figure(figsize=(8, 5)) +plt.bar(df_weather['Weather'], df_weather['NumSamples']) +plt.title('Number of Samples per Weather Condition') +plt.xlabel('Weather') +plt.ylabel('Num Samples') +plt.xticks(rotation=45) +plt.tight_layout() +plt.show() + +plt.figure(figsize=(8, 5)) +if not df_instances.empty: + plt.bar(df_instances['Class'], df_instances['NumInstances']) + plt.title('Total Object Instances per Class') + plt.xlabel('Class') + plt.ylabel('Num Instances') + plt.xticks(rotation=45) + plt.tight_layout() + plt.show() +else: + print('Skipping instance plot as no instances were found.') \ No newline at end of file diff --git a/3_data_augmentation.py b/3_data_augmentation.py new file mode 100644 index 0000000..30b2df7 --- /dev/null +++ b/3_data_augmentation.py @@ -0,0 +1,67 @@ +import os +import random +import numpy as np +import cv2 +from pathlib import Path +from tqdm import tqdm +import albumentations as A +import matplotlib.pyplot as plt + +# Paths +DATA_DIR = Path(os.path.join(os.path.dirname(__file__), "data")) +RGB_DIR = DATA_DIR / "rgb" +SEM_DIR = DATA_DIR / "semantics" +AUG_RGB_DIR = DATA_DIR / "aug_rgb" +AUG_SEM_DIR = DATA_DIR / "aug_semantics" +AUG_RGB_DIR.mkdir(parents=True, exist_ok=True) +AUG_SEM_DIR.mkdir(parents=True, exist_ok=True) + +# Augmentation pipeline for instance segmentation +AUGMENTATIONS = A.Compose([ + A.HorizontalFlip(p=0.5), # Flipping horizontally with 50% probability + A.RandomBrightnessContrast(p=0.5), # Adjusting brightness and contrast with 50% probability + A.Rotate(limit=15, p=0.5), # Rotating the image within a limit of 15 degrees with 50% probability + A.VerticalFlip(p=0.5), # Flipping vertically with 50% probability + A.RandomGamma(p=0.5) # Adjusting gamma with 50% probability +], additional_targets={'mask': 'image'}) + +# List all valid pairs +image_files = sorted([f for f in RGB_DIR.glob("*.png")]) +mask_files = sorted([f for f in SEM_DIR.glob("*.npy")]) + +assert len(image_files) == len(mask_files), "Mismatch between RGB and mask files." + +# Number of augmentations per image +AUG_PER_IMAGE = 2 + +# Perform augmentation +for img_path, mask_path in tqdm(zip(image_files, mask_files), total=len(image_files)): + img = cv2.imread(str(img_path)) + mask = np.load(mask_path) # Shape: (height, width, 3) + + for i in range(AUG_PER_IMAGE): + augmented = AUGMENTATIONS(image=img, mask=mask) + aug_img = augmented['image'] + aug_mask = augmented['mask'] + + # Save augmented files + stem = img_path.stem + f"_aug{i}" + cv2.imwrite(str(AUG_RGB_DIR / f"{stem}.png"), aug_img) + np.save(str(AUG_SEM_DIR / f"{stem}.npy"), aug_mask) + +print("\nData augmentation completed.") + +# Optional: Visualize one sample +sample_img = cv2.imread(str(image_files[0])) +sample_mask = np.load(mask_files[0]) +aug_example = AUGMENTATIONS(image=sample_img, mask=sample_mask) + +fig, axs = plt.subplots(1, 2, figsize=(10, 5)) +axs[0].imshow(cv2.cvtColor(sample_img, cv2.COLOR_BGR2RGB)) +axs[0].set_title("Original") +axs[0].axis("off") +axs[1].imshow(cv2.cvtColor(aug_example['image'], cv2.COLOR_BGR2RGB)) +axs[1].set_title("Augmented") +axs[1].axis("off") +plt.tight_layout() +plt.show() \ No newline at end of file diff --git a/4_yolo_annotation.py b/4_yolo_annotation.py new file mode 100644 index 0000000..dd91dd4 --- /dev/null +++ b/4_yolo_annotation.py @@ -0,0 +1,128 @@ +import os +import numpy as np +import cv2 + +# Directory setup +DATA_DIR = os.path.join(os.path.dirname(__file__), "data") +SEMANTICS_DIRS = [os.path.join(DATA_DIR, "semantics"), os.path.join(DATA_DIR, "aug_semantics")] +RGB_DIRS = { + "semantics": os.path.join(DATA_DIR, "rgb"), + "aug_semantics": os.path.join(DATA_DIR, "aug_rgb") +} + +ANNOT_DIR_NORMAL = os.path.join(DATA_DIR, "yolo_annotation/normal") +ANNOT_DIR_AUG = os.path.join(DATA_DIR, "yolo_annotation/augmented") +os.makedirs(ANNOT_DIR_NORMAL, exist_ok=True) +os.makedirs(ANNOT_DIR_AUG, exist_ok=True) + +# Valid class mappings +valid_classes = { + 114: "Car", # vehicle.car + 115: "Pedestrian", # walker.pedestrian + 122: "TrafficLight", # traffic.traffic_light + 177: "Bus" # vehicle.bus +} + +def generate_yolo_v11_annotations(instance_array, image_shape): + """ + Generate YOLOv11 annotations from instance segmentation mask. + + Args: + instance_array: 3-channel array (R: class ID, G and B: instance ID) + image_shape: Shape of the corresponding RGB image (height, width, channels) + + Returns: + List of annotation lines in YOLO format + """ + height, width = image_shape[:2] + annotations = [] + + # Extract class and instance IDs from the three-channel mask + class_ids = instance_array[:, :, 0] + instance_ids = (instance_array[:, :, 1].astype(np.uint16) << 8) + instance_array[:, :, 2] + + # Find unique instance IDs (excluding 0, which is background) + unique_instance_ids = np.unique(instance_ids) + unique_instance_ids = unique_instance_ids[unique_instance_ids != 0] + + for inst_id in unique_instance_ids: + # Create binary mask for this instance + binary_mask = (instance_ids == inst_id).astype(np.uint8) + + # Find external contours + contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + + if contours: + # Select the largest contour by area + largest_contour = max(contours, key=cv2.contourArea) + + if len(largest_contour) >= 3 and cv2.contourArea(largest_contour) > 50: + # Get class ID (all pixels in an instance should share the same class) + cls_id = class_ids[binary_mask > 0][0] + + if cls_id in valid_classes: + # Normalize contour points + normalized_pts = [(float(x) / width, float(y) / height) for [[x, y]] in largest_contour] + # Create annotation line + line = f"{cls_id} " + " ".join(f"{x:.6f} {y:.6f}" for x, y in normalized_pts) + annotations.append(line) + + return annotations + +def process_directory(sem_dir, rgb_dir, is_aug=False): + """ + Process a directory of semantic masks to generate YOLO annotations. + + Args: + sem_dir: Path to semantics directory + rgb_dir: Path to corresponding RGB directory + is_aug: Boolean indicating if processing augmented data + """ + for fname in os.listdir(sem_dir): + if not fname.endswith(".npy"): + continue + + sem_path = os.path.join(sem_dir, fname) + stem = os.path.splitext(fname)[0] + + if is_aug: + # For augmented data, mask file name like "rgb_something_aug0.npy" + img_name = stem + ".png" # Corresponding image is "rgb_something_aug0.png" + label_path = os.path.join(ANNOT_DIR_AUG, stem + ".txt") + else: + # For normal data, mask file name like "seg_something.npy" + img_name = "rgb_" + stem.replace("seg_", "") + ".png" + label_path = os.path.join(ANNOT_DIR_NORMAL, stem + ".txt") + + img_path = os.path.join(rgb_dir, img_name) + + if not os.path.exists(img_path): + print(f"Image not found: {img_path}") + continue + + # Load three-channel instance segmentation mask + sem_mask = np.load(sem_path) + image = cv2.imread(img_path) + annotations = generate_yolo_v11_annotations(sem_mask, image.shape) + + # Always write annotation file, even if empty + with open(label_path, "w") as f: + for line in annotations: + f.write(line + "\n") + + if annotations: + print(f"[✔] Wrote: {label_path} with {len(annotations)} annotations") + else: + print(f"[✔] Empty annotation file written for: {fname}") + +def main(): + """Process all semantics directories to generate YOLOv11 annotations.""" + for sem_path in SEMANTICS_DIRS: + base = os.path.basename(sem_path) + is_aug = "aug" in base + rgb_path = RGB_DIRS[base] + process_directory(sem_path, rgb_path, is_aug=is_aug) + print("✅ YOLOv11-compatible annotations generated.") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/5_model_training.py b/5_model_training.py new file mode 100644 index 0000000..edc5a43 --- /dev/null +++ b/5_model_training.py @@ -0,0 +1,277 @@ +import os +import shutil +from pathlib import Path +import yaml +import torch +from sklearn.model_selection import train_test_split +from ultralytics import YOLO +import numpy as np +import time + +# Directory setup +BASE_DIR = Path(__file__).parent / "data" # Set base directory to the local "data" folder +DATA_DIR = BASE_DIR +IMAGES_DIR = DATA_DIR / "rgb" +AUG_IMAGES_DIR = DATA_DIR / "aug_rgb" +LABELS_DIR = DATA_DIR / "yolo_annotation" / "normal" +AUG_LABELS_DIR = DATA_DIR / "yolo_annotation" / "augmented" +OUTPUT_DIR = BASE_DIR / "output" +MODEL_DIR = BASE_DIR / "models" +WORKING_LABELS_DIR = BASE_DIR / "labels" / "normal" +WORKING_AUG_LABELS_DIR = BASE_DIR / "labels" / "augmented" + +# Create directories if they don't exist +for dir_path in [WORKING_LABELS_DIR, WORKING_AUG_LABELS_DIR, OUTPUT_DIR, MODEL_DIR]: + dir_path.mkdir(parents=True, exist_ok=True) + +# Copy label files to working directories +for file in LABELS_DIR.glob("*.txt"): + shutil.copy(file, WORKING_LABELS_DIR) +for file in AUG_LABELS_DIR.glob("*.txt"): + shutil.copy(file, WORKING_AUG_LABELS_DIR) + +# Class and weather definitions +CLASSES = { + 114: "Car", # vehicle.car + 115: "Pedestrian", # walker.pedestrian + 122: "TrafficLight", # traffic.traffic_light + 177: "Bus" # vehicle.bus +} +WEATHER_CONDITIONS = ["Clear_Day", "Clear_Night", "Heavy_Rain", "Foggy"] + +def verify_setup(): + """Verify that all required directories exist.""" + required_dirs = [DATA_DIR, IMAGES_DIR, AUG_IMAGES_DIR, LABELS_DIR, AUG_LABELS_DIR] + for dir_path in required_dirs: + if not dir_path.exists(): + raise FileNotFoundError(f"Directory not found: {dir_path}") + +def fix_labels(labels_dir, valid_classes): + """Remove invalid class IDs from label files.""" + for label_file in labels_dir.glob("*.txt"): + with open(label_file, "r") as f: + lines = f.readlines() + + fixed_lines = [] + for line in lines: + try: + class_id = int(line.split()[0]) + if class_id in valid_classes: + fixed_lines.append(line) + else: + print(f"Invalid class {class_id} in file {label_file}") + except (IndexError, ValueError): + print(f"Invalid line format in {label_file}: {line.strip()}") + + with open(label_file, "w") as f: + f.writelines(fixed_lines) + +def get_weather(img_path): + """Extract weather condition from image file name.""" + stem = img_path.stem + parts = stem.split("_") + if parts[-1].startswith("aug"): + weather = "_".join(parts[1:-2]) # For augmented data + else: + weather = "_".join(parts[1:-1]) # For original data + return weather + +def prepare_dataset(): + """Prepare image-label pairs for training.""" + fix_labels(WORKING_LABELS_DIR, CLASSES.keys()) + fix_labels(WORKING_AUG_LABELS_DIR, CLASSES.keys()) + + # Separate original and augmented data + original_images = list(IMAGES_DIR.glob("*.png")) + aug_images = list(AUG_IMAGES_DIR.glob("*.png")) + original_labels = list(WORKING_LABELS_DIR.glob("*.txt")) + aug_labels = list(WORKING_AUG_LABELS_DIR.glob("*.txt")) + + # Pair original data + orig_img_keys = {img.stem.replace("rgb_", ""): img for img in original_images} + orig_lbl_keys = {lbl.stem.replace("seg_", ""): lbl for lbl in original_labels} + original_pairs = [(orig_img_keys[key], orig_lbl_keys[key]) + for key in orig_img_keys if key in orig_lbl_keys] + + # Pair augmented data + aug_img_keys = {img.stem: img for img in aug_images} + aug_lbl_keys = {lbl.stem: lbl for lbl in aug_labels} + aug_pairs = [(aug_img_keys[key], aug_lbl_keys[key]) + for key in aug_img_keys if key in aug_lbl_keys] + + # Combine pairs + all_pairs = original_pairs + aug_pairs + if not all_pairs: + raise ValueError("No matching image-label pairs found!") + + images = [img for img, _ in all_pairs] + labels = [lbl for _, lbl in all_pairs] + print(f"Found {len(images)} valid image-label pairs") + return images, labels + +def split_dataset(images, labels): + """Split dataset into train, validation, and test sets.""" + train_val_imgs, test_imgs, train_val_lbls, test_lbls = train_test_split( + images, labels, test_size=0.2, random_state=42 + ) + train_imgs, val_imgs, train_lbls, val_lbls = train_test_split( + train_val_imgs, train_val_lbls, test_size=0.25, random_state=42 + ) + + print(f"Train: {len(train_imgs)}, Validation: {len(val_imgs)}, Test: {len(test_imgs)}") + return train_imgs, val_imgs, test_imgs, train_lbls, val_lbls, test_lbls + +def organize_dataset(train_imgs, val_imgs, test_imgs, train_lbls, val_lbls, test_lbls): + """Organize dataset into train, val, and test directories.""" + splits = { + "train": (train_imgs, train_lbls), + "val": (val_imgs, val_lbls), + "test": (test_imgs, test_lbls) + } + dataset_dir = Path("/dataset") + for split, (imgs, lbls) in splits.items(): + img_dir = dataset_dir / "images" / split + lbl_dir = dataset_dir / "labels" / split + img_dir.mkdir(parents=True, exist_ok=True) + lbl_dir.mkdir(parents=True, exist_ok=True) + + for img, lbl in zip(imgs, lbls): + shutil.copy(img, img_dir / img.name) + shutil.copy(lbl, lbl_dir / lbl.name) + + return dataset_dir + +def create_dataset_yaml(dataset_dir): + """Create dataset.yaml file for YOLO training.""" + yaml_data = { + "path": str(dataset_dir.absolute()), + "train": "images/train", + "val": "images/val", + "test": "images/test", + "nc": len(CLASSES), + "names": list(CLASSES.values()) + } + yaml_path = dataset_dir / "dataset.yaml" + with open(yaml_path, "w") as f: + yaml.dump(yaml_data, f, sort_keys=False) + return yaml_path + +def train_model(model_path, data_yaml): + """Train YOLO model with specified parameters.""" + start_time = time.time() + model = YOLO(model_path) + model_name = Path(model_path).stem + results = model.train( + data=data_yaml, + task='segment', # Explicitly specify segmentation task + epochs=1, + imgsz=640, + batch=8, + device=0 if torch.cuda.is_available() else "cpu", + project="/data/models", + name=f"{model_name}_segment", + exist_ok=True, + pretrained=True, + verbose=True, + patience=10, + lr0=0.001, + optimizer="AdamW", + plots=True + ) + training_time = time.time() - start_time + return model, results, training_time + +def evaluate_model(model, test_imgs, weather_conditions, data_yaml): + """Evaluate model on test set.""" + metrics = {} + # Perform evaluation on the entire test set + results = model.val(data=data_yaml, split="test") + + # General metrics + overall_metrics = { + "mAP50": results.seg.map50 if hasattr(results, 'seg') else 0.0, + "mAP50-95": results.seg.map if hasattr(results, 'seg') else 0.0, + "inference_time": results.speed["inference"], + "per_class_map": results.seg.ap_class_by_class if hasattr(results, 'seg') else [] + } + + # Assign metrics to each weather condition (temporary until per-condition evaluation is implemented) + for condition in weather_conditions: + condition_imgs = [img for img in test_imgs if get_weather(img) == condition] + if condition_imgs: + metrics[condition] = overall_metrics + else: + metrics[condition] = {"note": "No images for this condition"} + + return metrics + +def save_results(model_name, metrics, training_time): + """Save training and evaluation results to a file.""" + results_file = OUTPUT_DIR / f"{model_name}_detailed_results.txt" + with open(results_file, "w") as f: + f.write(f"Model: {model_name}\n") + f.write(f"Training Time: {training_time:.2f} seconds\n\n") + f.write("Performance Across Weather Conditions:\n") + f.write("-" * 50 + "\n") + + for condition, condition_metrics in metrics.items(): + f.write(f"\nWeather Condition: {condition}\n") + if "note" in condition_metrics: + f.write(f"{condition_metrics['note']}\n") + else: + f.write(f"mAP50: {condition_metrics['mAP50']:.4f}\n") + f.write(f"mAP50-95: {condition_metrics['mAP50-95']:.4f}\n") + f.write(f"Inference Time: {condition_metrics['inference_time']:.4f} ms\n") + f.write("\nPer-Class Performance:\n") + for class_id, class_map in enumerate(condition_metrics['per_class_map']): + f.write(f"{CLASSES[class_id]}: {class_map:.4f}\n") + f.write("-" * 50 + "\n") + +def main(): + """Main function to prepare dataset, train, and evaluate models.""" + try: + verify_setup() + + # Prepare and split dataset + images, labels = prepare_dataset() + train_imgs, val_imgs, test_imgs, train_lbls, val_lbls, test_lbls = split_dataset(images, labels) + + # Organize dataset and create YAML + dataset_dir = organize_dataset(train_imgs, val_imgs, test_imgs, train_lbls, val_lbls, test_lbls) + data_yaml = create_dataset_yaml(dataset_dir) + + # Model paths + model_paths = [ + "yolo11n-seg.pt", + "yolo11m-seg.pt", + "yolo11l-seg.pt" + ] + + for model_path in model_paths: + try: + print(f"\nTraining Model: {model_path}") + model, results, training_time = train_model(model_path, data_yaml) + + print(f"Evaluating Model: {model_path}") + metrics = evaluate_model(model, test_imgs, WEATHER_CONDITIONS, data_yaml) + save_results(Path(model_path).stem, metrics, training_time) + + except Exception as e: + print(f"Error with {model_path}: {e}") + if 'metrics' in locals(): + print(f"\nResults Summary for {model_path}:") + print("-" * 50) + for condition, condition_metrics in metrics.items(): + print(f"\nWeather Condition: {condition}") + if "note" not in condition_metrics: + print(f"mAP50: {condition_metrics['mAP50']:.4f}") + print(f"mAP50-95: {condition_metrics['mAP50-95']:.4f}") + print(f"Inference Time: {condition_metrics['inference_time']:.2f} ms") + print(f"\nTotal Training Time: {training_time:.2f} seconds") + print("-" * 50) + + except Exception as e: + print(f"Error: {e}") + +if __name__ == "__main__": + main() \ No newline at end of file From 300f1c31ba68e866dc4fcdcc91dfb252c4c1a18d Mon Sep 17 00:00:00 2001 From: Forough Koohi <87902202+foroughkoohi@users.noreply.github.com> Date: Sat, 7 Jun 2025 16:06:22 +0330 Subject: [PATCH 2/3] Sync data-code --- 1-data_set_generation.py | 337 +++++++++++++++++++++++++++++++++++++++ 2-dataAnalysis.py | 96 +++++++++++ 3-dataAugmentation.py | 68 ++++++++ 4-yoloAnotation.py | 89 +++++++++++ 5-modelTraining.py | 288 +++++++++++++++++++++++++++++++++ 5 files changed, 878 insertions(+) create mode 100644 1-data_set_generation.py create mode 100644 2-dataAnalysis.py create mode 100644 3-dataAugmentation.py create mode 100644 4-yoloAnotation.py create mode 100644 5-modelTraining.py diff --git a/1-data_set_generation.py b/1-data_set_generation.py new file mode 100644 index 0000000..262c5b1 --- /dev/null +++ b/1-data_set_generation.py @@ -0,0 +1,337 @@ +import carla +import random +import time +import os +import numpy as np +import cv2 +from datetime import datetime +from collections import defaultdict +import queue + +# Initial settings +DATA_DIR = os.path.join(os.path.dirname(__file__), "data") +os.makedirs(os.path.join(DATA_DIR, "rgb"), exist_ok=True) +os.makedirs(os.path.join(DATA_DIR, "semantics"), exist_ok=True) +os.makedirs(os.path.join(DATA_DIR, "analysis"), exist_ok=True) + +# Weather conditions +weather_presets = { + + 'Clear_Day': carla.WeatherParameters.ClearNoon, + 'Clear_Night': carla.WeatherParameters.ClearNight, + 'Heavy_Rain': carla.WeatherParameters.HardRainNoon, + 'Foggy': carla.WeatherParameters(fog_density=80.0) + +} + +# Class mapping for semantic segmentation +class_mapping = {0: "Pedestrian", 1: "Traffic Light", 2: "Car", 3: "Bus"} + +def save_image(image, path): + """Save image data to file (RGB or semantic class IDs)""" + array = np.frombuffer(image.raw_data, dtype=np.dtype("uint8")) + array = np.reshape(array, (image.height, image.width, 4))[:, :, :3] + + if 'semantics' in path: + # Save only class IDs from red channel + class_ids = array[:, :, 2] + np.save(path.replace('.png', '.npy'), class_ids) + else: + # Save RGB image + cv2.imwrite(path, array) + +def create_dense_traffic_scene(world, traffic_manager, blueprint_library, zone_center, radius=40): + """Create a dense scene with cars, buses, pedestrians and traffic lights""" + vehicles_list = [] + walkers_list = [] + lights_list = [] + + vehicle_types = { + # Cars + "vehicle.tesla.model3": 4, + "vehicle.audi.a2": 3, + "vehicle.dodge.charger_police": 2, + "vehicle.toyota.prius": 3, + # Buses + "vehicle.volkswagen.t2": 1, + "vehicle.mercedes.sprinter": 2, + "vehicle.carlamotors.carlacola": 1 + } + + # Use pre-defined spawn points from the map + all_spawn_points = world.get_map().get_spawn_points() + random.shuffle(all_spawn_points) + used_spawn_points = set() + + # Spawn vehicles + for vehicle_type, weight in vehicle_types.items(): + is_bus = any(bus_id in vehicle_type for bus_id in ["t2", "sprinter", "carlacola"]) + spawn_count = weight * (2 if is_bus else 3) + for _ in range(spawn_count): + try: + # Find the spawn point closest to zone_center + spawn_point = min( + (sp for sp in all_spawn_points if sp not in used_spawn_points), + key=lambda sp: sp.location.distance(zone_center.location) + ) + used_spawn_points.add(spawn_point) + bp = blueprint_library.find(vehicle_type) + if bp.has_attribute('color'): + color = random.choice(bp.get_attribute('color').recommended_values) + bp.set_attribute('color', color) + vehicle = world.try_spawn_actor(bp, spawn_point) + if vehicle is not None: + vehicles_list.append(vehicle) + vehicle.set_autopilot(True, traffic_manager.get_port()) + if is_bus: + traffic_manager.distance_to_leading_vehicle(vehicle, random.uniform(4.0, 6.0)) + traffic_manager.vehicle_percentage_speed_difference(vehicle, random.uniform(-10, 0)) + else: + traffic_manager.distance_to_leading_vehicle(vehicle, random.uniform(1.0, 3.0)) + traffic_manager.vehicle_percentage_speed_difference(vehicle, random.uniform(-20, 20)) + traffic_manager.update_vehicle_lights(vehicle, True) + except Exception as e: + print(f"Vehicle spawn error: {e}") + continue + + # Create pedestrian groups + pedestrian_groups = 4 + pedestrians_per_group = 6 + for _ in range(pedestrian_groups): + group_center = carla.Location( + x=zone_center.location.x + random.uniform(-radius/2, radius/2), + y=zone_center.location.y + random.uniform(-radius/2, radius/2), + z=zone_center.location.z + ) + for _ in range(pedestrians_per_group): + try: + offset = random.uniform(1, 4) + angle = random.uniform(0, 2 * np.pi) + x = group_center.x + offset * np.cos(angle) + y = group_center.y + offset * np.sin(angle) + spawn_point = carla.Transform( + carla.Location(x=x, y=y, z=group_center.z + 1.0) + ) + bp = random.choice(blueprint_library.filter("walker.pedestrian.*")) + walker = world.try_spawn_actor(bp, spawn_point) + if walker is not None: + walkers_list.append(walker) + walker_control = carla.WalkerControl() + walker_control.speed = random.uniform(0.6, 2.0) + direction = carla.Vector3D( + x=group_center.x - x, + y=group_center.y - y, + z=0 + ) + direction.make_unit_vector() + walker_control.direction = direction + walker.apply_control(walker_control) + except Exception as e: + print(f"Walker spawn error: {e}") + continue + + # Configure traffic lights + try: + traffic_lights = world.get_actors().filter('traffic.traffic_light*') + for traffic_light in traffic_lights: + if traffic_light.get_location().distance(zone_center.location) < radius: + states = [carla.TrafficLightState.Red, + carla.TrafficLightState.Yellow, + carla.TrafficLightState.Green] + traffic_light.set_state(random.choice(states)) + traffic_light.set_green_time(random.uniform(4.0, 8.0)) + traffic_light.set_red_time(random.uniform(2.0, 6.0)) + traffic_light.set_yellow_time(random.uniform(1.0, 3.0)) + lights_list.append(traffic_light) + except Exception as e: + print(f"Traffic light error: {e}") + + return vehicles_list, walkers_list, lights_list + +def analyze_dataset(data_dir): + """Analyze dataset and generate reports""" + print("\n=== Dataset Analysis ===") + + # Count samples + rgb_files = os.listdir(os.path.join(data_dir, "rgb")) + semantic_files = [f for f in os.listdir(os.path.join(data_dir, "semantics")) + if f.endswith('.npy')] + + print(f"Total Samples: {len(rgb_files)}\n") + + # Count samples per weather + weather_counts = defaultdict(int) + for f in rgb_files: + weather = f.split('_')[0] + weather_counts[weather] += 1 + + print("Samples per weather condition:") + for weather, count in weather_counts.items(): + print(f"{weather}: {count} samples") + + print("\nTotal Annotated Instances:") + # Analyze semantic segmentation data + class_pixels = defaultdict(int) + for sem_file in semantic_files: + sem_data = np.load(os.path.join(data_dir, "semantics", sem_file)) + unique, counts = np.unique(sem_data, return_counts=True) + for cls_id, count in zip(unique, counts): + if cls_id in class_mapping: + class_pixels[cls_id] += count + + for cls_id, pixels in class_pixels.items(): + print(f"{class_mapping[cls_id]} ({cls_id}): {pixels} pixels") + +def main(): + client = carla.Client("localhost", 2000) + client.set_timeout(10.0) + world = client.get_world() + + # Initialize variables + ego_vehicle = None + camera_rgb = None + camera_seg = None + vehicles_list = [] + walkers_list = [] + lights_list = [] + + try: + # World settings + original_settings = world.get_settings() + settings = world.get_settings() + settings.synchronous_mode = True + settings.fixed_delta_seconds = 0.05 + world.apply_settings(settings) + + # Create ego vehicle + blueprint_library = world.get_blueprint_library() + ego_bp = blueprint_library.find("vehicle.tesla.model3") + + # Try to spawn ego vehicle + max_tries = 20 + spawn_points = world.get_map().get_spawn_points() + random.shuffle(spawn_points) + + for attempt in range(max_tries): + try: + spawn_point = spawn_points[attempt % len(spawn_points)] + ego_vehicle = world.spawn_actor(ego_bp, spawn_point) + print(f"Successfully spawned ego vehicle on attempt {attempt+1}") + break + except RuntimeError as e: + print(f"Failed attempt {attempt+1}/{max_tries} to spawn ego vehicle: {e}") + if attempt == max_tries - 1: + raise RuntimeError("Could not spawn ego vehicle") + + # Set up traffic manager + traffic_manager = client.get_trafficmanager(8000) + traffic_manager.set_synchronous_mode(True) + traffic_manager.set_global_distance_to_leading_vehicle(1.5) + traffic_manager.global_percentage_speed_difference(10.0) + ego_vehicle.set_autopilot(True, traffic_manager.get_port()) + + # Set up cameras + camera_rgb = world.spawn_actor( + blueprint_library.find("sensor.camera.rgb"), + carla.Transform(carla.Location(x=1.5, z=2.4)), + attach_to=ego_vehicle + ) + camera_seg = world.spawn_actor( + blueprint_library.find("sensor.camera.semantic_segmentation"), + carla.Transform(carla.Location(x=1.5, z=2.4)), + attach_to=ego_vehicle + ) + + # Set up image capture + global rgb_data, seg_data + rgb_data = None + seg_data = None + + def rgb_callback(image): + global rgb_data + rgb_data = image + + def seg_callback(image): + global seg_data + seg_data = image + + camera_rgb.listen(rgb_callback) + camera_seg.listen(seg_callback) + + # Data collection loop + frame_count = {w: 0 for w in weather_presets.keys()} + all_spawn_points = world.get_map().get_spawn_points() + + for weather_name, weather in weather_presets.items(): + world.set_weather(weather) + print(f"\nWeather set to: {weather_name}") + while frame_count[weather_name] < 100: + # Clean up previous actors + for vehicle in vehicles_list: + if vehicle and vehicle.is_alive: + vehicle.destroy() + for walker in walkers_list: + if walker and walker.is_alive: + walker.destroy() + vehicles_list.clear() + walkers_list.clear() + + # Select a valid spawn point + valid_spawn_points = [ + sp for sp in all_spawn_points + if sp.location.z > 0 + ] + if not valid_spawn_points: + print("No valid spawn points available!") + break + + zone_center = random.choice(valid_spawn_points) + + # Move ego vehicle to the selected spawn point + ego_vehicle.set_transform(zone_center) + + # Create new scene + new_vehicles, new_walkers, new_lights = create_dense_traffic_scene( + world, traffic_manager, blueprint_library, zone_center + ) + vehicles_list.extend(new_vehicles) + walkers_list.extend(new_walkers) + lights_list.extend(new_lights) + + # Let the scene stabilize + for _ in range(30): + world.tick() + + # Capture and save images + if rgb_data and seg_data: + filename = f"{weather_name}_{frame_count[weather_name]:05d}" + save_image(rgb_data, os.path.join(DATA_DIR, "rgb", f"rgb_{filename}.png")) + save_image(seg_data, os.path.join(DATA_DIR, "semantics", f"seg_{filename}.png")) + frame_count[weather_name] += 1 + print(f"Saved: {filename} | Count: {frame_count[weather_name]}") + + # Delay between captures + time.sleep(3.0) + + print("Data collection completed") + analyze_dataset(DATA_DIR) + + finally: + print("Cleaning up actors...") + if camera_rgb: + camera_rgb.stop() + if camera_seg: + camera_seg.stop() + if ego_vehicle: + ego_vehicle.set_autopilot(False) + ego_vehicle.destroy() + for vehicle in vehicles_list: + if vehicle and vehicle.is_alive: + vehicle.destroy() + for walker in walkers_list: + if walker and walker.is_alive: + walker.destroy() + world.apply_settings(original_settings) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/2-dataAnalysis.py b/2-dataAnalysis.py new file mode 100644 index 0000000..3c4e92c --- /dev/null +++ b/2-dataAnalysis.py @@ -0,0 +1,96 @@ +import os +import numpy as np +import pandas as pd +import cv2 +import matplotlib.pyplot as plt +from collections import defaultdict + +# Paths to your data directories +DATA_DIR = os.path.join(os.path.dirname(__file__), "data") +RGB_DIR = os.path.join(DATA_DIR, "rgb") +SEM_DIR = os.path.join(DATA_DIR, "semantics") + +# Ensure directories exist +if not os.path.isdir(RGB_DIR) or not os.path.isdir(SEM_DIR): + raise FileNotFoundError(f"RGB or semantics directory not found under {DATA_DIR}") + +# Mapping of semantic mask IDs to class names +CLASS_MAPPING = {0: "Pedestrian", 1: "Traffic Light", 2: "Car", 3: "Bus"} + +# Debug: inspect unique IDs across all semantic masks +mask_files = [f for f in os.listdir(SEM_DIR) if f.lower().endswith('.npy')] +unique_ids_all = set() +non_empty_count = 0 +for fname in mask_files: + mask = np.load(os.path.join(SEM_DIR, fname)) + ids = set(np.unique(mask)) + unique_ids_all.update(ids) + if len(ids - {0}) > 0: + non_empty_count += 1 +print(f"Debug: Total mask files = {len(mask_files)}, Non-empty masks = {non_empty_count}") +print(f"Debug: All unique IDs across masks = {sorted(unique_ids_all)}") +print("If these IDs don't include your CLASS_MAPPING keys, adjust CLASS_MAPPING or mask generation.") + +# 1) Count number of samples per weather condition +weather_counts = defaultdict(int) +for fname in os.listdir(RGB_DIR): + if fname.lower().startswith("rgb_") and fname.lower().endswith(('.png', '.jpg', '.jpeg')): + parts = fname.split("_") + weather = "_".join(parts[1:-1]) + weather_counts[weather] += 1 + +# Convert to DataFrame +df_weather = pd.DataFrame(sorted(weather_counts.items()), columns=["Weather", "NumSamples"]) + +# 2) Count object instances per class across all semantic masks +instance_counts = defaultdict(int) +for fname in os.listdir(SEM_DIR): + if fname.lower().endswith('.npy'): + mask_path = os.path.join(SEM_DIR, fname) + mask = np.load(mask_path) + for cls_id, cls_name in CLASS_MAPPING.items(): + binary = (mask == cls_id).astype(np.uint8) + if not binary.any(): + continue + num_labels, _ = cv2.connectedComponents(binary, connectivity=8) + instances = max(0, num_labels - 1) + instance_counts[cls_name] += instances + +# Convert to DataFrame +df_instances = pd.DataFrame(sorted(instance_counts.items()), columns=["Class", "NumInstances"]) + +# 3) Total annotated instances +total_instances = df_instances["NumInstances"].sum() if not df_instances.empty else 0 + +# 4) Print summary +print("\n=== Dataset Analysis Report ===") +print("Samples per Weather Condition:") +print(df_weather.to_string(index=False)) +print("\nInstances per Class:") +if df_instances.empty: + print("No instances found. Check CLASS_MAPPING IDs and mask generation.") +else: + print(df_instances.to_string(index=False)) +print(f"\nTotal Annotated Instances: {total_instances}\n") + +# 5) Visualizations +plt.figure(figsize=(8, 5)) +plt.bar(df_weather['Weather'], df_weather['NumSamples']) +plt.title('Number of Samples per Weather Condition') +plt.xlabel('Weather') +plt.ylabel('Num Samples') +plt.xticks(rotation=45) +plt.tight_layout() +plt.show() + +plt.figure(figsize=(8, 5)) +if not df_instances.empty: + plt.bar(df_instances['Class'], df_instances['NumInstances']) + plt.title('Total Object Instances per Class') + plt.xlabel('Class') + plt.ylabel('Num Instances') + plt.xticks(rotation=45) + plt.tight_layout() + plt.show() +else: + print('Skipping instance plot as no instances were found.') diff --git a/3-dataAugmentation.py b/3-dataAugmentation.py new file mode 100644 index 0000000..2d44a90 --- /dev/null +++ b/3-dataAugmentation.py @@ -0,0 +1,68 @@ +import os +import random +import numpy as np +import cv2 +from pathlib import Path +from tqdm import tqdm +import albumentations as A +import matplotlib.pyplot as plt + +# Paths +DATA_DIR = Path(os.path.join(os.path.dirname(__file__), "data")) +RGB_DIR = DATA_DIR / "rgb" +SEM_DIR = DATA_DIR / "semantics" +AUG_RGB_DIR = DATA_DIR / "aug_rgb" +AUG_SEM_DIR = DATA_DIR / "aug_semantics" +AUG_RGB_DIR.mkdir(parents=True, exist_ok=True) +AUG_SEM_DIR.mkdir(parents=True, exist_ok=True) + +# Augmentation pipeline +AUGMENTATIONS = A.Compose([ + A.HorizontalFlip(p=0.5), # Flipping horizontally with 50% probability + A.RandomBrightnessContrast(p=0.5), # Adjusting brightness and contrast with 50% probability + A.Rotate(limit=15, p=0.5), # Rotating the image within a limit of 15 degrees with 50% probability + A.VerticalFlip(p=0.5), # Flipping vertically with 50% probability + A.RandomGamma(p=0.5) # Adjusting gamma with 50% probability +]) + +# List all valid pairs +image_files = sorted([f for f in RGB_DIR.glob("*.png")]) +mask_files = sorted([f for f in SEM_DIR.glob("*.npy")]) + +assert len(image_files) == len(mask_files), "Mismatch between RGB and mask files." + +# Number of augmentations per image +AUG_PER_IMAGE = 2 + +# Perform augmentation +for img_path, mask_path in tqdm(zip(image_files, mask_files), total=len(image_files)): + img = cv2.imread(str(img_path)) + mask = np.load(mask_path).astype(np.uint8) + + for i in range(AUG_PER_IMAGE): + augmented = AUGMENTATIONS(image=img, mask=mask) + aug_img = augmented['image'] + aug_mask = augmented['mask'] + + # Save augmented files + stem = img_path.stem + f"_aug{i}" + cv2.imwrite(str(AUG_RGB_DIR / f"{stem}.png"), aug_img) + np.save(str(AUG_SEM_DIR / f"{stem}.npy"), aug_mask) + +print("\nData augmentation completed.") + +# Optional: Visualize one sample +sample_img = cv2.imread(str(image_files[0])) +sample_mask = np.load(mask_files[0]) +aug_example = AUGMENTATIONS(image=sample_img, mask=sample_mask) + +fig, axs = plt.subplots(1, 2, figsize=(10, 5)) +axs[0].imshow(cv2.cvtColor(sample_img, cv2.COLOR_BGR2RGB)) +axs[0].set_title("Original") +axs[0].axis("off") +axs[1].imshow(cv2.cvtColor(aug_example['image'], cv2.COLOR_BGR2RGB)) +axs[1].imshow(aug_example['mask'], alpha=0.4, cmap='jet') +axs[1].set_title("Augmented") +axs[1].axis("off") +plt.tight_layout() +plt.show() diff --git a/4-yoloAnotation.py b/4-yoloAnotation.py new file mode 100644 index 0000000..e38469c --- /dev/null +++ b/4-yoloAnotation.py @@ -0,0 +1,89 @@ +import os +import numpy as np +import cv2 + +DATA_DIR = os.path.join(os.path.dirname(__file__), "data") +SEMANTICS_DIRS = [os.path.join(DATA_DIR, "semantics"), os.path.join(DATA_DIR, "aug_semantics")] +RGB_DIRS = { + "semantics": os.path.join(DATA_DIR, "rgb"), + "aug_semantics": os.path.join(DATA_DIR, "aug_rgb") +} + +ANNOT_DIR_NORMAL = os.path.join(DATA_DIR, "yolo_annotation/normal") +ANNOT_DIR_AUG = os.path.join(DATA_DIR, "yolo_annotation/augmented") +os.makedirs(ANNOT_DIR_NORMAL, exist_ok=True) +os.makedirs(ANNOT_DIR_AUG, exist_ok=True) + +valid_classes = {0: "Pedestrian", 1: "Traffic Light", 2: "Car", 3: "Bus"} + +def generate_yolo_v11_annotations(instance_array, image_shape): + height, width = image_shape[:2] + annotations = [] + + for cls_id in valid_classes: + binary_mask = (instance_array == cls_id).astype(np.uint8) + if not binary_mask.any(): + continue + + n_labels, labels = cv2.connectedComponents(binary_mask, connectivity=8) + + for inst_id in range(1, n_labels): + inst_mask = (labels == inst_id).astype(np.uint8) + contours, _ = cv2.findContours(inst_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + + for cnt in contours: + if len(cnt) < 3 or cv2.contourArea(cnt) < 50: + continue + + normalized_pts = [(x / width, y / height) for [[x, y]] in cnt] + line = f"{cls_id} " + " ".join(f"{x:.6f} {y:.6f}" for x, y in normalized_pts) + annotations.append(line) + + return annotations + +def process_directory(sem_dir, rgb_dir, is_aug=False): + for fname in os.listdir(sem_dir): + if not fname.endswith(".npy"): + continue + + sem_path = os.path.join(sem_dir, fname) + stem = os.path.splitext(fname)[0] + + + if is_aug: + img_name = stem.replace("seg_", "") + ".png" + label_path = os.path.join(ANNOT_DIR_AUG, stem + ".txt") + else: + img_name = "rgb_" + stem.replace("seg_", "") + ".png" + label_path = os.path.join(ANNOT_DIR_NORMAL, stem + ".txt") + + img_path = os.path.join(rgb_dir, img_name) + + if not os.path.exists(img_path): + print(f"Image not found: {img_path}") + continue + + sem_mask = np.load(sem_path) + image = cv2.imread(img_path) + annotations = generate_yolo_v11_annotations(sem_mask, image.shape) + + if annotations: + with open(label_path, "w") as f: + for line in annotations: + f.write(line + "\n") + print(f"[✔] Wrote: {label_path}") + else: + if os.path.exists(label_path): + os.remove(label_path) + print(f"[✘] No valid annotations for: {fname}") + +def main(): + for sem_path in SEMANTICS_DIRS: + base = os.path.basename(sem_path) + is_aug = "aug" in base + rgb_path = RGB_DIRS[base] + process_directory(sem_path, rgb_path, is_aug=is_aug) + print("✅ YOLOv11-compatible annotations generated.") + +if __name__ == "__main__": + main() diff --git a/5-modelTraining.py b/5-modelTraining.py new file mode 100644 index 0000000..796cae3 --- /dev/null +++ b/5-modelTraining.py @@ -0,0 +1,288 @@ +import os +import shutil +from pathlib import Path +import yaml +import torch +from sklearn.model_selection import train_test_split +from ultralytics import YOLO +import numpy as np +import time +from PIL import Image + + +ROOT_DIR = Path().absolute() +DATA_DIR = Path(os.path.join(os.path.dirname(__file__), "data")) +IMAGES_DIR = DATA_DIR / "rgb" +AUG_IMAGES_DIR = DATA_DIR / "aug_rgb" +LABELS_DIR = DATA_DIR / "yolo_annotation/normal" +AUG_LABELS_DIR = DATA_DIR / "yolo_annotation/augmented" +OUTPUT_DIR = ROOT_DIR / "output" +MODEL_DIR = ROOT_DIR / "models" + + +for dir_path in [OUTPUT_DIR, MODEL_DIR]: + dir_path.mkdir(exist_ok=True) + + +CLASSES = {0: "Pedestrian", 1: "Traffic Light", 2: "Car", 3: "Bus"} +WEATHER_CONDITIONS = ["clear", "foggy", "rainy", "snowy"] + +def verify_setup(): + + required_dirs = [DATA_DIR, IMAGES_DIR, AUG_IMAGES_DIR, LABELS_DIR, AUG_LABELS_DIR] + for dir_path in required_dirs: + if not dir_path.exists(): + raise FileNotFoundError(f"Directory not found: {dir_path}") + +def fix_labels(labels_dir, valid_classes): + for label_file in labels_dir.glob("*.txt"): + with open(label_file, "r") as f: + lines = f.readlines() + + fixed_lines = [] + for line in lines: + class_id = int(line.split()[0]) + if class_id in valid_classes: + fixed_lines.append(line) + else: + print(f"Invalid class {class_id} in file {label_file}") + + + with open(label_file, "w") as f: + f.writelines(fixed_lines) + + +fix_labels(LABELS_DIR, range(len(CLASSES))) +fix_labels(AUG_LABELS_DIR, range(len(CLASSES))) + +def prepare_dataset(): + fix_labels(LABELS_DIR, CLASSES.keys()) + fix_labels(AUG_LABELS_DIR, CLASSES.keys()) + + images = list(IMAGES_DIR.glob("*.png")) + list(AUG_IMAGES_DIR.glob("*.png")) + labels = list(LABELS_DIR.glob("*.txt")) + list(AUG_LABELS_DIR.glob("*.txt")) + + image_stems = {img.stem.replace("rgb_", ""): img for img in images} + label_stems = {lbl.stem.replace("rgb_", ""): lbl for lbl in labels} + + paired_images = [] + paired_labels = [] + for stem in image_stems: + if stem in label_stems: + paired_images.append(image_stems[stem]) + paired_labels.append(label_stems[stem]) + else: + print(f"Warning: No matching label for image {image_stems[stem]}") + + if not paired_images: + raise ValueError("No matching image-label pairs found!") + + print(f"Found {len(paired_images)} valid image-label pairs") + return paired_images, paired_labels + +def split_dataset(images, labels): + + train_val_imgs, test_imgs, train_val_lbls, test_lbls = train_test_split( + images, labels, test_size=0.2, random_state=42 + ) + train_imgs, val_imgs, train_lbls, val_lbls = train_test_split( + train_val_imgs, train_val_lbls, test_size=0.25, random_state=42 + ) + + print(f"Train: {len(train_imgs)}, Validation: {len(val_imgs)}, Test: {len(test_imgs)}") + return train_imgs, val_imgs, test_imgs, train_lbls, val_lbls, test_lbls + +def organize_dataset(train_imgs, val_imgs, test_imgs, train_lbls, val_lbls, test_lbls): + + splits = { + "train": (train_imgs, train_lbls), + "val": (val_imgs, val_lbls), + "test": (test_imgs, test_lbls) + } + dataset_dir = DATA_DIR / "dataset" + for split, (imgs, lbls) in splits.items(): + img_dir = dataset_dir / "images" / split + lbl_dir = dataset_dir / "labels" / split + img_dir.mkdir(parents=True, exist_ok=True) + lbl_dir.mkdir(parents=True, exist_ok=True) + + for img, lbl in zip(imgs, lbls): + shutil.copy(img, img_dir / img.name) + shutil.copy(lbl, lbl_dir / lbl.name) + + return dataset_dir + +def create_dataset_yaml(dataset_dir): + yaml_data = { + "path": str(dataset_dir.absolute()), + "train": "images/train", + "val": "images/val", + "test": "images/test", + "nc": len(CLASSES), + "names": list(CLASSES.values()) + } + yaml_path = dataset_dir / "dataset.yaml" + with open(yaml_path, "w") as f: + yaml.dump(yaml_data, f, sort_keys=False) + return yaml_path + +def train_model(model_type, data_yaml): + start_time = time.time() + model = YOLO(model_type) + results = model.train( + data=data_yaml, + epochs=5, + imgsz=640, + batch=8, + device=0 if torch.cuda.is_available() else "cpu", + project=str(OUTPUT_DIR), + name=f"{model_type.split('.')[0]}_detect", + exist_ok=True, + pretrained=True, + verbose=True, + patience=10, + lr0=0.001, + optimizer="AdamW", + plots=True + ) + training_time = time.time() - start_time + return model, results, training_time + +def evaluate_model(model, test_imgs, weather_conditions): + metrics = {} + for condition in weather_conditions: + + temp_dir = DATA_DIR / "temp" / condition + temp_dir.mkdir(parents=True, exist_ok=True) + + + condition_imgs = [img for img in test_imgs if condition in str(img)] + for img in condition_imgs: + shutil.copy(img, temp_dir / img.name) + + + results = model.val( + data=str(DATA_DIR / "dataset" / "dataset.yaml"), + split="test" + ) + + + shutil.rmtree(temp_dir) + + metrics[condition] = { + "mAP50": results.box.map50, + "mAP50-95": results.box.map, + "inference_time": results.speed["inference"], + "per_class_map": results.box.ap_class_by_class + } + return metrics + +def save_results(model_name, metrics, training_time): + results_file = OUTPUT_DIR / f"{model_name}_detailed_results.txt" + with open(results_file, "w") as f: + f.write(f"Model: {model_name}\n") + f.write(f"Training Time: {training_time:.2f} seconds\n\n") + f.write("Performance Across Weather Conditions:\n") + f.write("-" * 50 + "\n") + + for condition, condition_metrics in metrics.items(): + f.write(f"\nWeather Condition: {condition}\n") + f.write(f"mAP50: {condition_metrics['mAP50']:.4f}\n") + f.write(f"mAP50-95: {condition_metrics['mAP50-95']:.4f}\n") + f.write(f"Inference Time: {condition_metrics['inference_time']:.4f} ms\n") + + f.write("\nPer-Class Performance:\n") + for class_id, class_map in enumerate(condition_metrics['per_class_map']): + f.write(f"{CLASSES[class_id]}: {class_map:.4f}\n") + f.write("-" * 50 + "\n") + +def fix_coordinates(coords, img_width, img_height): + """Helper function to fix and validate coordinates""" + x_center, y_center = coords[0], coords[1] + width, height = coords[2], coords[3] + + # Ensure coordinates are within image bounds + x_center = max(0, min(x_center, img_width)) + y_center = max(0, min(y_center, img_height)) + width = max(0, min(width, img_width - x_center)) + height = max(0, min(height, img_height - y_center)) + + # Normalize coordinates + x_norm = x_center / img_width + y_norm = y_center / img_height + w_norm = width / img_width + h_norm = height / img_height + + return x_norm, y_norm, w_norm, h_norm + +def convert_npy_to_txt(npy_dir, output_dir, class_id=0): + output_dir.mkdir(parents=True, exist_ok=True) + for npy_file in npy_dir.glob("*.npy"): + data = np.load(npy_file) + txt_file = output_dir / f"{npy_file.stem.replace('seg_', 'rgb_')}.txt" + + # Read corresponding image to get dimensions + img_file = IMAGES_DIR / f"{npy_file.stem.replace('seg_', 'rgb_')}.png" + if not img_file.exists(): + print(f"Warning: No matching image found for {npy_file}") + continue + + from PIL import Image + with Image.open(img_file) as img: + img_width, img_height = img.size + + with open(txt_file, "w") as f: + for obj in data: + if len(obj) >= 4: + # Use the helper function to fix coordinates + x_norm, y_norm, w_norm, h_norm = fix_coordinates( + obj[:4], img_width, img_height + ) + if all(0 <= coord <= 1 for coord in [x_norm, y_norm, w_norm, h_norm]): + f.write(f"{class_id} {x_norm:.6f} {y_norm:.6f} {w_norm:.6f} {h_norm:.6f}\n") + else: + print(f"Warning: Invalid normalized coordinates in {npy_file}") + else: + print(f"Warning: Invalid data format in {npy_file}") + + +def main(): + try: + verify_setup() + convert_npy_to_txt(Path(DATA_DIR / "semantics"), LABELS_DIR) + + images, labels = prepare_dataset() + train_imgs, val_imgs, test_imgs, train_lbls, val_lbls, test_lbls = split_dataset(images, labels) + + dataset_dir = organize_dataset(train_imgs, val_imgs, test_imgs, train_lbls, val_lbls, test_lbls) + data_yaml = create_dataset_yaml(dataset_dir) + + model_types = ["yolov8n.pt", "yolov8m.pt", "yolov8l.pt"] + for model_type in model_types: + try: + print(f"\nModel Training: {model_type}") + model, results, training_time = train_model(model_type, data_yaml) + + print(f"Model Evaluation: {model_type}") + metrics = evaluate_model(model, test_imgs, WEATHER_CONDITIONS) + save_results(model_type.split(".")[0], metrics, training_time) + + print(f"\nResults Summary for {model_type}:") + print("-" * 50) + for condition, condition_metrics in metrics.items(): + print(f"\nWeather Condition: {condition}") + print(f"mAP50: {condition_metrics['mAP50']:.4f}") + print(f"mAP50-95: {condition_metrics['mAP50-95']:.4f}") + print(f"Inference Time: {condition_metrics['inference_time']:.2f} ms") + print(f"\nTotal Training Time: {training_time:.2f} seconds") + print("-" * 50) + + except Exception as e: + print(f"Error with {model_type}: {e}") + continue + + except Exception as e: + print(f"Error: {e}") + +if __name__ == "__main__": + main() \ No newline at end of file From 7e9e19d53bd7aeaab44ed622a56e19d27d27644f Mon Sep 17 00:00:00 2001 From: Forough Koohi <87902202+foroughkoohi@users.noreply.github.com> Date: Sat, 7 Jun 2025 17:23:31 +0330 Subject: [PATCH 3/3] =?UTF-8?q?=D9=81=D8=A7=DB=8C=D9=84=D9=87=D8=A7=DB=8C?= =?UTF-8?q?=20=D8=AE=D8=B1=D9=88=D8=AC=DB=8C=20=D9=88=20=D9=81=D8=A7=DB=8C?= =?UTF-8?q?=D9=84=D9=87=D8=A7=DB=8C=20=D8=AF=DB=8C=D8=AA=D8=A7=D8=B3=D8=AA?= =?UTF-8?q?=20=D8=B3=D8=A7=D8=AE=D8=AA=D9=87=20=D8=B4=D8=AF=D9=87?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- Dataset_and_Outputs_link.txt | 1 + 1 file changed, 1 insertion(+) create mode 100644 Dataset_and_Outputs_link.txt diff --git a/Dataset_and_Outputs_link.txt b/Dataset_and_Outputs_link.txt new file mode 100644 index 0000000..e9d7e3b --- /dev/null +++ b/Dataset_and_Outputs_link.txt @@ -0,0 +1 @@ +https://drive.google.com/file/d/1GqT0EXcsesLxQEa1_RkvRSgSbZnw8kbd/view?usp=drive_link \ No newline at end of file