# %%writefile train1.py
# import os
# import sys
# import json
# import boto3
# import torch
# import glob
# import time
# import random
# import datetime
# from ultralytics import YOLO
# from smdebug.pytorch import get_hook
# from smdebug.core.modes import ModeKeys
# from smdebug.pytorch import Hook
# from ultralytics.cfg import get_cfg

# # AWS S3 setup
# s3 = boto3.client("s3")

# # === Clean Logging for CloudWatch & Notebook ===
# class CleanLogger:
#     def write(self, message):
#         sys.__stdout__.write(message.replace("\r", "\n"))
#     def flush(self):
#         sys.__stdout__.flush()

# sys.stdout = CleanLogger()
# sys.stderr = CleanLogger()

# # === Disable Progress Bars & ANSI Color Codes ===
# cfg = get_cfg()
# cfg.progress = False  # Disables tqdm bars

# # SageMaker hyperparameters file path
# HYPERPARAMS_FILE = "/opt/ml/input/config/hyperparameters.json"

# # Path to SageMaker resource configuration file
# RESOURCE_CONFIG_FILE = "/opt/ml/input/config/resourceconfig.json"
# TENSORBOARD_DIR = "/opt/ml/output/tensors"

# # Load hyperparameters
# def load_hyperparameters():
#     if os.path.exists(HYPERPARAMS_FILE):
#         with open(HYPERPARAMS_FILE, "r") as f:
#             hyperparams = json.load(f)
#             print(f"Loaded hyperparameters: {hyperparams}")
#             return hyperparams
#     else:
#         print("No hyperparameters.json found, using defaults.")
#         return {}

# # Load hyperparameters with fallback defaults
# hyperparams = load_hyperparameters()
# epochs = int(hyperparams.get("epochs", 10))
# batch_size = int(hyperparams.get("batch_size", 16))
# learning_rate = float(hyperparams.get("lr", 0.001))
# num_instances = int(hyperparams.get("num_instances", 1))

# # Get Training Job Name
# training_job_name = os.getenv("TRAINING_JOB_NAME", "default_training_job")
# print(f"Training Job Name: {training_job_name}")

# # Retrieve Instance ID from SageMaker
# def get_instance_id():
#     if os.path.exists(RESOURCE_CONFIG_FILE):
#         with open(RESOURCE_CONFIG_FILE, "r") as f:
#             resource_config = json.load(f)
#         current_host = resource_config.get("current_host", "algo-1")
#         hosts = resource_config.get("hosts", ["algo-1"])
#         return hosts.index(current_host)
#     return 0  # Default to 0 if file not found

# instance_id = get_instance_id()
# print(f"Instance {instance_id}/{num_instances} is active")

# # S3 paths
# input_data_bucket = os.getenv("INPUT_DATA_BUCKET", "uh-sagemaker-test-bucket")
# input_data_prefix = os.getenv("INPUT_DATA_PREFIX", "datasets/pavement_marking_combined/")
# model_bucket = os.getenv("MODEL_BUCKET", "uh-sagemaker-test-bucket")
# output_best_model_key = os.getenv("OUTPUT_BEST_MODEL_KEY", "models/PMD_yolov8m_640_v1_best.pt")
# output_train_results_key = f"models/PMD_Model/{training_job_name}/training_results/"
# input_model_bucket = os.getenv("INPUT_MODEL_BUCKET", "uh-general")
# input_yaml_key = os.getenv("INPUT_YAML_KEY", "datasets/pavement_marking_combined/datanew.yaml")
# input_model_key = os.getenv("INPUT_MODEL_KEY", "models/inferencing/pmd/v1/PMD_yolov8m_640_v1.pt")
# model_bucket = os.getenv("MODEL_BUCKET", "uh-sagemaker-test-bucket")

# # S3 paths (again)
# s3_bucket = "uh-sagemaker-test-bucket"
# output_model_s3_prefix = f"models/PMD_Model/{training_job_name}/"
# best_model_s3_key = f"{output_model_s3_prefix}instance_{instance_id}_best.pt"
# train_results_s3_prefix = f"{output_model_s3_prefix}instance_{instance_id}/"

# # Local paths
# output_model_dir = "/opt/ml/model/"
# local_model_path = f"{output_model_dir}/best.pt"
# # (Ignore the stray number '654654331173' below)
# # Local paths
# local_train_images_dir = "/opt/ml/input/data/train/images/"
# local_train_labels_dir = "/opt/ml/input/data/train/labels/"
# local_val_images_dir = "/opt/ml/input/data/val/images/"
# local_val_labels_dir = "/opt/ml/input/data/val/labels/"
# output_model_dir = "/opt/ml/model/"
# train_results_dir = f"runs/segment/"
# local_model_path = "/opt/ml/input/model/PMD_yolov8m_640_v1.pt"
# data_config_path = "/opt/ml/input/config/datanew.yaml"

# # Ensure directories exist
# os.makedirs(local_train_images_dir, exist_ok=True)
# os.makedirs(local_train_labels_dir, exist_ok=True)
# os.makedirs(local_val_images_dir, exist_ok=True)
# os.makedirs(local_val_labels_dir, exist_ok=True)
# os.makedirs(output_model_dir, exist_ok=True)

# # 🔹 List S3 files
# def list_s3_files(bucket, prefix):
#     files = []
#     paginator = s3.get_paginator('list_objects_v2')
#     for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
#         for obj in page.get('Contents', []):
#             files.append(obj['Key'])
#     return files

# # Function to Download Partitioned Images and Full Labels
# def download_data(split):
#     """Each instance downloads a subset of images but all labels for the dataset split (train/val)."""
#     img_prefix = f"{input_data_prefix}{split}/images/"
#     label_prefix = f"{input_data_prefix}{split}/labels/"

#     all_images = list_s3_files(input_data_bucket, img_prefix)
#     all_labels = list_s3_files(input_data_bucket, label_prefix)

#     if not all_images or not all_labels:
#         raise RuntimeError(f"No images or labels found in S3 for {split} dataset!")

#     # Sort to ensure alignment
#     all_images.sort()
#     all_labels.sort()

#     # Partition images among instances
#     total_images = len(all_images)
#     files_per_instance = total_images // num_instances
#     start_idx = instance_id * files_per_instance
#     end_idx = total_images if instance_id == num_instances - 1 else start_idx + files_per_instance
#     assigned_images = all_images[start_idx:end_idx]

#     print(f"Instance {instance_id}/{num_instances} → Processing {len(assigned_images)} {split} images.")

#     # Define local directories based on split
#     local_img_dir = local_train_images_dir if split == "train" else local_val_images_dir
#     local_label_dir = local_train_labels_dir if split == "train" else local_val_labels_dir

#     image_download_count = 0
#     label_download_count = 0

#     # Download assigned images
#     for img_s3_key in assigned_images:
#         img_filename = os.path.basename(img_s3_key)
#         local_img_path = os.path.join(local_img_dir, img_filename)
#         s3.download_file(input_data_bucket, img_s3_key, local_img_path)
#         image_download_count += 1

#     # Download all labels (each instance gets full labels)
#     for label_s3_key in all_labels:
#         label_filename = os.path.basename(label_s3_key)
#         local_label_path = os.path.join(local_label_dir, label_filename)
#         s3.download_file(input_data_bucket, label_s3_key, local_label_path)
#         label_download_count += 1

#     print(f"Instance {instance_id} Download Summary ({split}): {image_download_count} images, {label_download_count} labels")

# # Download Train & Validation Data
# download_data("train")
# print("Downloaded Training images")
# download_data("val")
# print("Downloaded Validation images")

# # Ensure model and yaml exist
# os.makedirs(os.path.dirname(local_model_path), exist_ok=True)
# if not os.path.exists(local_model_path) or os.path.getsize(local_model_path) == 0:
#     print(f"Model not found! Downloading: {input_model_key}...")
#     s3.download_file(input_model_bucket, input_model_key, local_model_path)

# os.makedirs(os.path.dirname(data_config_path), exist_ok=True)
# if not os.path.exists(data_config_path) or os.path.getsize(data_config_path) == 0:
#     print(f"YAML not found! Downloading: {input_yaml_key}...")
#     s3.download_file(input_data_bucket, input_yaml_key, data_config_path)

# # Define TensorBoard Directory
# tensorboard_dir = "/opt/ml/output/tensorboard"
# os.makedirs(tensorboard_dir, exist_ok=True)

# # Correct Debugger Hook Configuration
# hook = Hook(
#     out_dir="/opt/ml/output/tensors",  # Ensure correct output directory
#     export_tensorboard=True,           # Enable TensorBoard support
#     tensorboard_dir="/opt/ml/output/tensorboard"
# )
# print("Debugger Hook Initialized. Saving tensors to /opt/ml/output/tensors")

# # Attach Hook to Training Mode
# hook.set_mode(ModeKeys.TRAIN)
# hook = get_hook()
# if hook:
#     hook.set_mode(ModeKeys.TRAIN)

# # Oversampling and Data Augmentation for Minority Classes
# def oversample_minority_classes(labels_dir, images_dir, target_count=100):
#     """
#     Duplicate  images and label files for classes 
#     that are underrepresented, so that each class has at least 'target_count' samples.
#     """
#     import shutil
#     from random import choice
    
#     # First, build a mapping from class ID to list of label files
#     files_by_class = {}
#     for label_file in glob.glob(os.path.join(labels_dir, "*.txt")):
#         with open(label_file, "r") as f:
#             for line in f:
#                 parts = line.strip().split()
#                 if parts:
#                     class_id = int(parts[0])
#                     files_by_class.setdefault(class_id, []).append(label_file) 
    
#     # For each class, if count < target_count, duplicate files until reaching target_count
#     for class_id, file_list in files_by_class.items():
#         current_count = len(file_list)
#         if current_count < target_count:
#             deficit = target_count - current_count
#             print(f"Class {class_id} is underrepresented (has {current_count} samples). Duplicating {deficit} files...")
#             for i in range(deficit):
#                 file_to_copy = choice(file_list)
#                 base_name = os.path.basename(file_to_copy)
#                 new_label_name = f"dup_{class_id}_{i}_{base_name}"
#                 new_label_path = os.path.join(labels_dir, new_label_name)
#                 shutil.copy(file_to_copy, new_label_path)
                
#                 # Also duplicate the corresponding image (assuming same base name with .jpg extension)
#                 image_base = os.path.splitext(base_name)[0] + ".jpg"
#                 src_image_path = os.path.join(images_dir, image_base)
#                 if os.path.exists(src_image_path):
#                     new_image_name = f"dup_{class_id}_{i}_{image_base}"
#                     new_image_path = os.path.join(images_dir, new_image_name)
#                     shutil.copy(src_image_path, new_image_path)
#                 else:
#                     print(f"Warning: Image {src_image_path} not found for label {file_to_copy}")
# ##############################################################################

# # Call oversampling on the training data before training starts
# print("Starting oversampling of minority classes on training data...")
# oversample_minority_classes(local_train_labels_dir, local_train_images_dir, target_count=100)
# print("Oversampling complete.")

# # Training & Model Preparation
# def train_yolo():
#     """Train YOLOv8 model with oversampled training data."""
#     model = YOLO(local_model_path)
    
#     print("Using oversampled training data for class balancing.")
    
#     model.train(
#         data=data_config_path,
#         epochs=epochs,
#         imgsz=640,
#         batch=batch_size,
#         device="cuda" if torch.cuda.is_available() else "cpu",
#         weight_decay=1e-5,
#         patience=5,
#         lr0=learning_rate,
#         save=True
#     )
    
#     # Ensure Debugger Hook Saves Data
#     if hook is not None:
#         hook.save_tensor("dummy_tensor", torch.tensor([1.0]))
    
#     latest_train_dir = max(glob.glob("runs/segment/train*"), key=os.path.getctime, default=None)
#     if latest_train_dir:
#         best_model_path = os.path.join(latest_train_dir, "weights", "best.pt")
#         if os.path.exists(best_model_path):
#             print(f"Best model saved at {best_model_path}")
#             return best_model_path, latest_train_dir
#     raise RuntimeError("Best model not found!")

# # Train the model
# start_time = time.time()
# best_model_path, latest_train_dir = train_yolo()
# end_time = time.time()
# print(f"Training complete! ⏳ Time taken: {round(end_time - start_time, 2)}s")

# # 🔹 Upload Trained Model & Results to S3
# def upload_to_s3(local_path, bucket, s3_key):
#     """Upload a file to S3."""
#     s3.upload_file(local_path, bucket, s3_key)
#     print(f"Uploaded: {local_path} -> s3://{bucket}/{s3_key}")

# def upload_folder_to_s3(local_folder, bucket, s3_prefix):
#     """Upload an entire folder to S3."""
#     for root, _, files in os.walk(local_folder):
#         for file in files:
#             local_path = os.path.join(root, file)
#             s3_key = os.path.join(s3_prefix, os.path.relpath(local_path, local_folder))
#             upload_to_s3(local_path, bucket, s3_key)

# # Upload best model and training results
# upload_to_s3(best_model_path, s3_bucket, best_model_s3_key)
# upload_folder_to_s3(latest_train_dir, s3_bucket, train_results_s3_prefix)

# def federated_averaging(model_paths):
#     """Federated Averaging: Load multiple models and merge their weights."""
#     models = [YOLO(path) for path in model_paths]
#     merged_model = models[0]  # Use the first model as the base

#     state_dicts = [m.model.state_dict() for m in models]  # Extract state dicts

#     new_state_dict = {}
#     for key in state_dicts[0].keys():
#         weights = [sd[key].float() for sd in state_dicts]  # Convert all to float32 tensors
#         stacked_weights = torch.stack(weights, dim=0)  # Ensure it's a tensor
#         avg_weight = torch.mean(stacked_weights, dim=0)  # Compute element-wise mean
#         new_state_dict[key] = avg_weight

#     merged_model.model.load_state_dict(new_state_dict)  # Load averaged weights

#     return merged_model

# if instance_id == 0:
#     print("Master instance collecting models for Federated Averaging...")

#     # Wait for all models to be uploaded
#     time.sleep(60)  # Wait for other instances to complete uploads

#     # List all models in S3
#     response = s3.list_objects_v2(Bucket=s3_bucket, Prefix=output_model_s3_prefix)
#     model_keys = [obj["Key"] for obj in response.get("Contents", []) if obj["Key"].endswith("best.pt")]

#     if len(model_keys) < num_instances:
#         print(f"Not all models uploaded yet! Found {len(model_keys)}/{num_instances}. Retrying in 60s...")
#         time.sleep(60)  # Retry after 60s
#         response = s3.list_objects_v2(Bucket=s3_bucket, Prefix=output_model_s3_prefix)
#         model_keys = [obj["Key"] for obj in response.get("Contents", []) if obj["Key"].endswith("best.pt")]

#     # Download all models
#     model_paths = []
#     for key in model_keys:
#         local_path = f"/tmp/{os.path.basename(key)}"
#         s3.download_file(s3_bucket, key, local_path)
#         model_paths.append(local_path)

#     if len(model_paths) < num_instances:
#         raise RuntimeError(f"Some models are missing. Found {len(model_paths)}/{num_instances}")

#     print("All models collected. Performing Federated Averaging...")
#     merged_model = federated_averaging(model_paths)
#     print("Training job completed with Federated Averaging!")

#     # Save merged model
#     final_model_path = "/opt/ml/model/final_best.pt"
#     merged_model.save(final_model_path)
#     print(f"Federated Averaged Model saved at {final_model_path}")

#     # Upload final model to S3
#     final_model_s3_key = f"models/PMD_Model/Best_Model/PMD_yolov8m_640_v1_best.pt"
#     upload_to_s3(final_model_path, s3_bucket, final_model_s3_key)
#     print(f"Final model uploaded to: s3://{s3_bucket}/{final_model_s3_key}")
#     print(f"🔹 Class Weights Used for Training: Oversampling/Augmentation approach")



can you give me some reference link for the oversampling approach i had used in this script