Saini16/Blood_Cell_Object_Detection
1
1"""2This script is meant to be run in Google Colab for fine-tuning the YOLOv10 model on the BCCD dataset.3It contains all the steps needed for training and should be run before deploying the application.4"""5 6import os7import glob8import zipfile9import requests10import xml.etree.ElementTree as ET11from pathlib import Path12import shutil13import ultralytics14from ultralytics import YOLO15import numpy as np16import time17 18def download_bccd_dataset():19 """20 Downloads the BCCD dataset from the GitHub repository.21 Returns the path to the dataset directory.22 """23 # Install dependencies if needed24 os.system('pip install ultralytics gdown')25 26 # Clone the repository27 os.system('git clone https://github.com/Shenggan/BCCD_Dataset.git')28 29 # Verify download30 dataset_dir = Path('BCCD_Dataset')31 if not dataset_dir.exists():32 print("Failed to download dataset using git. Trying alternative download...")33 # Alternative download method using direct download links34 os.makedirs('BCCD_Dataset/BCCD', exist_ok=True)35 url = "https://github.com/Shenggan/BCCD_Dataset/archive/refs/heads/master.zip"36 r = requests.get(url, allow_redirects=True)37 with open('bccd_dataset.zip', 'wb') as f:38 f.write(r.content)39 40 # Extract the zipfile41 with zipfile.ZipFile('bccd_dataset.zip', 'r') as zip_ref:42 zip_ref.extractall('.')43 44 # Move contents to the expected location45 extracted_dir = Path('BCCD_Dataset-master')46 if extracted_dir.exists():47 # Copy contents to the BCCD_Dataset directory48 for item in extracted_dir.glob('*'):49 if item.is_dir():50 shutil.copytree(item, dataset_dir / item.name)51 else:52 shutil.copy(item, dataset_dir / item.name)53 54 print("Dataset downloaded successfully.")55 return dataset_dir56 57def setup_dataset_for_yolo(dataset_path):58 """59 Prepares the BCCD dataset for YOLO format.60 Args:61 dataset_path: Path to the downloaded dataset62 Returns:63 Path to the processed dataset64 """65 yolo_dir = Path('BCCD_YOLO')66 os.makedirs(yolo_dir, exist_ok=True)67 68 # Create directory structure69 for split in ['train', 'val', 'test']:70 os.makedirs(yolo_dir / split / 'images', exist_ok=True)71 os.makedirs(yolo_dir / split / 'labels', exist_ok=True)72 73 # Map sources to destinations74 splits = {75 'train': dataset_path / 'BCCD' / 'train',76 'val': dataset_path / 'BCCD' / 'val',77 'test': dataset_path / 'BCCD' / 'test'78 }79 80 # Process each split81 for split_name, split_dir in splits.items():82 image_files = list(split_dir.glob('*.jpg'))83 for img_file in image_files:84 # Copy image85 shutil.copy(img_file, yolo_dir / split_name / 'images' / img_file.name)86 87 # Convert annotation88 xml_file = split_dir / f"{img_file.stem}.xml"89 if xml_file.exists():90 txt_file = yolo_dir / split_name / 'labels' / f"{img_file.stem}.txt"91 convert_annotations(xml_file, txt_file)92 93 return yolo_dir94 95def convert_annotations(xml_path, txt_path):96 """97 Converts XML annotations to YOLO format TXT files.98 Args:99 xml_path: Path to XML annotation file100 txt_path: Path to output TXT file101 """102 tree = ET.parse(xml_path)103 root = tree.getroot()104 105 # Get image dimensions106 size = root.find('size')107 img_width = int(size.find('width').text)108 img_height = int(size.find('height').text)109 110 # Map class names to IDs111 class_map = {'RBC': 0, 'WBC': 1, 'Platelets': 2}112 113 with open(txt_path, 'w') as f:114 for obj in root.findall('object'):115 cls_name = obj.find('name').text116 if cls_name not in class_map:117 continue118 119 cls_id = class_map[cls_name]120 121 # Get bounding box coordinates122 bbox = obj.find('bndbox')123 x_min = float(bbox.find('xmin').text)124 y_min = float(bbox.find('ymin').text)125 x_max = float(bbox.find('xmax').text)126 y_max = float(bbox.find('ymax').text)127 128 # Convert to YOLO format: center_x, center_y, width, height129 x_center = (x_min + x_max) / (2.0 * img_width)130 y_center = (y_min + y_max) / (2.0 * img_height)131 width = (x_max - x_min) / img_width132 height = (y_max - y_min) / img_height133 134 # Write to file135 f.write(f"{cls_id} {x_center} {y_center} {width} {height}\n")136 137def create_dataset_yaml(dataset_path):138 """139 Creates the YAML file required by YOLOv10 for training.140 Args:141 dataset_path: Path to the processed dataset142 """143 yaml_content = f"""144# YOLOv10 dataset config for BCCD145path: {dataset_path.absolute()} # Root directory146train: train/images # Train images relative to path147val: val/images # Validation images relative to path148test: test/images # Test images relative to path149 150# Classes151names:152 0: RBC153 1: WBC154 2: Platelets155 156# Number of classes157nc: 3158"""159 160 yaml_path = dataset_path / 'bccd.yaml'161 with open(yaml_path, 'w') as f:162 f.write(yaml_content)163 164 return yaml_path165 166def train_model(dataset_path):167 """168 Trains YOLOv10 on the BCCD dataset.169 Args:170 dataset_path: Path to the processed dataset171 Returns:172 Path to the trained model173 """174 # Create YAML config file175 yaml_path = create_dataset_yaml(dataset_path)176 177 # Import required modules178 import torch179 180 # Load a pretrained YOLOv10 model181 # Note: Use 'yolov10n.pt' for faster training, or 'yolov10s.pt' for better accuracy182 model = YOLO('yolov10n.pt') # Nano model183 184 # Train the model185 device = '0' if torch.cuda.is_available() else 'cpu'186 print(f"Training on device: {device}")187 188 results = model.train(189 data=str(yaml_path),190 epochs=50, # Number of epochs191 imgsz=640, # Image size192 batch=16, # Batch size193 patience=15, # Early stopping patience194 device=device,195 project='BCCD_Training',196 name='yolov10_bccd',197 seed=42,198 workers=8 if torch.cuda.is_available() else 1199 )200 201 # Get the path to the best model202 best_model_path = Path('BCCD_Training/yolov10_bccd/weights/best.pt')203 204 # Export the model to other formats if needed205 model.export(format='onnx')206 207 # Copy model to Google Drive if running in Colab208 try:209 from google.colab import drive210 drive_path = Path('/content/drive/MyDrive/BCCD_Model')211 drive_path.mkdir(exist_ok=True, parents=True)212 213 model_save_path = drive_path / 'yolov10_bccd.pt'214 shutil.copy(best_model_path, model_save_path)215 print(f"Model saved to Google Drive at {model_save_path}")216 except:217 print("Not running in Colab or couldn't mount Google Drive.")218 219 return best_model_path220 221def main():222 """223 Main function to execute the fine-tuning process.224 """225 start_time = time.time()226 227 print("Step 1: Downloading BCCD dataset...")228 dataset_path = download_bccd_dataset()229 230 print("Step 2: Setting up dataset in YOLO format...")231 yolo_dataset_path = setup_dataset_for_yolo(dataset_path)232 233 print("Step 3: Training YOLOv10 model...")234 trained_model_path = train_model(yolo_dataset_path)235 236 elapsed_time = (time.time() - start_time) / 60237 print(f"Training completed in {elapsed_time:.2f} minutes.")238 print(f"Trained model saved at: {trained_model_path}")239 240 return trained_model_path241 242if __name__ == "__main__":243 main()