Skip to content

Commit

Permalink
feat: YOLO_Detection_train 3type_box_detection #22
Browse files Browse the repository at this point in the history
  • Loading branch information
MCG committed Nov 15, 2024
1 parent ce8ded1 commit 22400e6
Show file tree
Hide file tree
Showing 2 changed files with 74 additions and 1 deletion.
48 changes: 47 additions & 1 deletion box_detect/box_plot.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -239,7 +239,7 @@
},
{
"cell_type": "code",
"execution_count": 5,
"execution_count": null,
"metadata": {},
"outputs": [
{
Expand Down Expand Up @@ -313,6 +313,52 @@
"print(f\"Train and validation datasets created.\")\n",
"print(f\"Train files: {len(train_files)}, Validation files: {len(valid_files)}\")\n"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"data.yaml file has been created at: /data/ephemeral/home/MCG/yolo_dataset_split/data.yaml\n"
]
}
],
"source": [
"#yaml파일 생성\n",
"\n",
"import os\n",
"\n",
"# Define directories\n",
"dataset_dir = \"/data/ephemeral/home/MCG/yolo_dataset_split\"\n",
"train_dir = os.path.join(dataset_dir, \"train\")\n",
"valid_dir = os.path.join(dataset_dir, \"valid\")\n",
"yaml_file_path = \"/data/ephemeral/home/MCG/yolo_dataset_split/data.yaml\"\n",
"\n",
"# Define class names\n",
"classes = [\"finger\", \"radius_ulna\", \"others\"]\n",
"\n",
"# Generate YAML content\n",
"yaml_content = f\"\"\"\n",
"train: {train_dir}\n",
"val: {valid_dir}\n",
"\n",
"# Number of classes\n",
"nc: {len(classes)}\n",
"\n",
"# Class names\n",
"names: {classes}\n",
"\"\"\"\n",
"\n",
"# Write to data.yaml\n",
"with open(yaml_file_path, \"w\") as yaml_file:\n",
" yaml_file.write(yaml_content)\n",
"\n",
"print(f\"data.yaml file has been created at: {yaml_file_path}\")\n"
]
}
],
"metadata": {
Expand Down
27 changes: 27 additions & 0 deletions box_detect/yolov11_train.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
from ultralytics import YOLO

# Load a COCO-pretrained YOLO11n model
model = YOLO("yolo11x.pt")

# Train the model on the COCO8 example dataset for 100 epochs
results = model.train(
data="/data/ephemeral/home/MCG/yolo_dataset_split/data.yaml",
epochs=500,
imgsz=2048,
batch=2,
hsv_h=0.0, # Hue shift 비활성화
hsv_s=0.0, # Saturation shift 비활성화
hsv_v=0.2, # Brightness shift 비활성화
degrees=0.2, # 이미지 회전 비활성화
translate=0.0, # 이미지 이동 비활성화
scale=0.2,
shear=0.2, # 이미지 비틀기 비활성화
perspective=0.0, # 원근 변환 비활성화
flipud=0.0, # 상하 뒤집기 비활성화
mosaic=0.0, # Mosaic 비활성화
mixup=0.0, # Mixup 비활성화
copy_paste=0.0 , # Copy-Paste 비활성화
erasing=0.0,
crop_fraction=0.0
)
#유지한 증강: Scale, degrees, shear, hsv_v, fliplr

0 comments on commit 22400e6

Please sign in to comment.