commit 2e6f7a8a1f2bfce4118cf9892a457c38cbb4a178 Author: Zuxin Dai Date: Fri Jan 24 10:04:59 2025 -0500 Computer vision YOLO11 model diff --git a/Test_Realtime.py b/Test_Realtime.py new file mode 100644 index 0000000..90dde4c --- /dev/null +++ b/Test_Realtime.py @@ -0,0 +1,94 @@ +import cv2 +from ultralytics import YOLO + +# Camera index (default camera is 0) +camera_index = 0 + +# Load the YOLO model +model = YOLO(r"D:\AIM\lemon\runs\detect\train4\weights\best.pt") # Load custom model + +# Initialize the camera +cap = cv2.VideoCapture(camera_index) +if not cap.isOpened(): + print("Unable to open the camera. Please check the device.") + exit() + +fps = int(cap.get(cv2.CAP_PROP_FPS)) +width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) +height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) +print(f"Camera resolution: {width}x{height}, FPS: {fps} FPS") + +# Dictionary to track the state of each lemon +lemon_states = {} # Format: {ID: "State"} + +# Define class labels +class_labels = { + 0: "Bruised", + 1: "DefectiveLemon", + 2: "GoodLemon", + 3: "NotRipeLemon", + 4: "Rotten" +} + +# Classes that require ID assignment +id_tracked_classes = ["DefectiveLemon", "GoodLemon", "NotRipeLemon"] + +# Set the window to be resizable +cv2.namedWindow("Live Detection", cv2.WINDOW_NORMAL) + +# Process video stream in real-time +while cap.isOpened(): + ret, frame = cap.read() + if not ret: + print("Unable to read camera input. Terminating program.") + break + + # Perform object tracking using BoT-SORT + results = model.track(source=frame, conf=0.5, tracker='botsort.yaml', show=False) + + for result in results: + frame = result.orig_img # Current frame + detections = result.boxes # Detection box information + + for box in detections: + x1, y1, x2, y2 = map(int, box.xyxy[0]) # Detection box coordinates + obj_id = int(box.id) if box.id is not None else -1 # Tracking object ID + class_id = int(box.cls) # Class ID + score = box.conf # Confidence score + label = class_labels.get(class_id, "Unknown") # Get class label + + # Update lemon state and output information for tracked boxes + if obj_id != -1 and label in id_tracked_classes: + if obj_id not in lemon_states: + lemon_states[obj_id] = label + else: + # Once detected as "DefectiveLemon," the state remains "DefectiveLemon" + if lemon_states[obj_id] != "DefectiveLemon": + lemon_states[obj_id] = label + + # Output ID, position, and label + position = f"({x1}, {y1}, {x2}, {y2})" + print(f"ID: {obj_id}, Position: {position}, Label: {lemon_states[obj_id]}") + + # Draw detection boxes and labels (including untracked ones) + if obj_id != -1 and label in id_tracked_classes: + display_text = f"ID {obj_id} | {lemon_states[obj_id]}" + else: + display_text = label # For untracked labels, only show the class + + cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2) + cv2.putText(frame, display_text, (x1, y1 - 10), + cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2) + + # Display the processed video stream + cv2.imshow("Live Detection", frame) + + # Exit the loop when ESC key is pressed + if cv2.waitKey(1) & 0xFF == 27: # 27 is the ASCII value for ESC key + print("ESC key detected. Exiting the program.") + break + +# Release resources +cap.release() +cv2.destroyAllWindows() +print("Camera video processing complete. Program terminated.") diff --git a/Test_image.py b/Test_image.py new file mode 100644 index 0000000..6bfe904 --- /dev/null +++ b/Test_image.py @@ -0,0 +1,66 @@ +import os +import cv2 +import pandas as pd +from ultralytics import YOLO + +# Define input and output directories +input_dir = r'D:\AIM\pecan\test_image\input_image' +output_dir = r'D:\AIM\pecan\test_image\output_image' + +# Create the output directory if it doesn't exist +os.makedirs(output_dir, exist_ok=True) + +# Load a YOLO model (you can specify your model path here) +model = YOLO(r"D:\AIM\pecan\runs\detect\train2\weights\best.pt") # Load a custom-trained model + +# Initialize a threshold and a list for classifications +threshold = 0.5 +all_classifications = [] + +# Iterate over all images in the input directory +for image_filename in os.listdir(input_dir): + if image_filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff')): # Check for image files + image_path = os.path.join(input_dir, image_filename) + image = cv2.imread(image_path) + if image is None: + continue # Skip files that cannot be read + + # Perform inference + results = model(image)[0] + classifications = [] + + # Create a semi-transparent overlay for the bounding boxes + for result in results.boxes.data.tolist(): + x1, y1, x2, y2, score, class_id = result + + if score > threshold: + # Create a semi-transparent overlay for the bounding box area + overlay = image.copy() + cv2.rectangle(overlay, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), -1) + alpha = 0.15 # Transparency factor (0 - transparent, 1 - opaque) + # Apply the overlay + image = cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0) + + # Draw bounding box outline + cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2) + # Prepare label with class name and confidence score + label = f"{results.names[int(class_id)].upper()} {score:.2f}" + # Put text (class label with confidence score) + cv2.putText(image, label, (int(x1), int(y1 - 10)), + cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2, cv2.LINE_AA) + + # Add classification to the list + classifications.append(f"{results.names[int(class_id)]} ({score:.2f})") + + # Save the output image + output_path = os.path.join(output_dir, image_filename) + cv2.imwrite(output_path, image) + + # Save classifications for each image + all_classifications.append({"image": image_filename, "classes": classifications}) + +# Create a DataFrame from the classifications +df = pd.DataFrame(all_classifications) +df.to_excel(os.path.join(output_dir, 'output_classifications.xlsx'), index=False) + +cv2.destroyAllWindows() \ No newline at end of file diff --git a/Test_logic.py b/Test_logic.py new file mode 100644 index 0000000..f8d6a6d --- /dev/null +++ b/Test_logic.py @@ -0,0 +1,52 @@ +import os +import pandas as pd +from ultralytics import YOLO +import cv2 + +# Input and output video paths +video_path = r'D:\AIM\pecan\GH014359.mp4' +video_path_out = r'D:\AIM\pecan\GH014359_out.mp4' + +cap = cv2.VideoCapture(video_path) +ret, frame = cap.read() +H, W, _ = frame.shape +out = cv2.VideoWriter(video_path_out, cv2.VideoWriter_fourcc(*'MP4V'), int(cap.get(cv2.CAP_PROP_FPS)), (W, H)) + +# Load the YOLO model +model = YOLO(r"D:\AIM\pecan\runs\detect\train2\weights\best.pt") # Load a custom model + +threshold = 0.5 +detected_cracked = False # Initialize a flag for detecting cracked pecans + +while ret: + # Perform detection on the current frame + results = model(frame)[0] + + for result in results.boxes.data.tolist(): + x1, y1, x2, y2, score, class_id = result + + if score > threshold: + # Draw bounding boxes and labels + cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 4) + label = results.names[int(class_id)].upper() + cv2.putText(frame, f"{label} {score:.2f}", (int(x1), int(y1 - 10)), + cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 255, 0), 3, cv2.LINE_AA) + + # Check for the "cracked" label + if label == "CRACKED": + detected_cracked = True + + # Write the processed frame to the output video + out.write(frame) + ret, frame = cap.read() + +# Determine the final label based on detections +final_label = "CRACKED" if detected_cracked else "GOOD" + +# Print the final label +print(f"Final Label: {final_label}") + +# Release video resources +cap.release() +out.release() +cv2.destroyAllWindows() diff --git a/Test_logic_track.py b/Test_logic_track.py new file mode 100644 index 0000000..9320d83 --- /dev/null +++ b/Test_logic_track.py @@ -0,0 +1,80 @@ +import os +import cv2 +from ultralytics import YOLO + +# 输入和输出视频路径 +# video_path = r'D:\AIM\pecan\OneDrive_2_2024-11-7\G5 Flex 01 8-5-2024, 1.20.12pm EDT - 8-5-2024, 1.23.22pm EDT.mp4' +# video_path_out = r'D:\AIM\pecan\G5 Flex 01 8-5-2024_out.mp4' +video_path = r'D:\AIM\pecan\GH014359.mp4' +video_path_out = r'D:\AIM\pecan\GH014359_out.mp4' + +# 加载 YOLO 模型 +model = YOLO(r"D:\AIM\pecan\runs\detect\train2\weights\best.pt") # 加载自定义模型 + +# 初始化 VideoWriter 用于保存输出视频 +cap = cv2.VideoCapture(video_path) +fps = int(cap.get(cv2.CAP_PROP_FPS)) +width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) +height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) +fourcc = cv2.VideoWriter_fourcc(*'mp4v') +out = cv2.VideoWriter(video_path_out, fourcc, fps, (width, height)) + +# 字典，用于跟踪每个核桃的状态 +walnut_states = {} # 格式: {ID: "状态"} + +# 定义类别标签 +class_labels = { + 0: "circumferential", + 1: "cracked", + 2: "crushed", + 3: "longitudinal", + 4: "open", + 5: "uncracked" +} + +# 需要分配 ID 的类别 +id_tracked_classes = ["cracked", "uncracked"] + +# 使用 BoT-SORT 进行目标跟踪 +results = model.track(source=video_path, conf=0.5, tracker='botsort.yaml', show=False) + +for result in results: + frame = result.orig_img # 当前帧 + detections = result.boxes # 检测框信息 + + # 处理每个检测框 + for box in detections: + x1, y1, x2, y2 = map(int, box.xyxy[0]) # 检测框坐标 + obj_id = int(box.id) if box.id is not None else -1 # 跟踪目标ID + class_id = int(box.cls) # 类别ID + score = box.conf # 置信度 + + # 获取检测框对应的类别标签 + label = class_labels.get(class_id, "unknown") + + # 仅对需要分配ID的类别更新核桃状态 + if label in id_tracked_classes: + if obj_id not in walnut_states: + walnut_states[obj_id] = label + else: + # 一旦检测到“cracked”，状态保持为“cracked” + if walnut_states[obj_id] != "cracked": + walnut_states[obj_id] = label + + display_text = f"ID {obj_id} | {walnut_states[obj_id]}" + else: + # 非分配ID的类别，仅显示类别标签 + display_text = label + + # 绘制检测框和标签 + cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2) + cv2.putText(frame, display_text, (x1, y1 - 10), + cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2) + + # 将处理好的帧写入输出视频 + out.write(frame) + +# 释放资源 +cap.release() +out.release() +print("视频处理完成，结果已保存。") \ No newline at end of file diff --git a/Test_track_putposition.py b/Test_track_putposition.py new file mode 100644 index 0000000..3e8196b --- /dev/null +++ b/Test_track_putposition.py @@ -0,0 +1,80 @@ +import os +import cv2 +from ultralytics import YOLO + +# Input and output video paths +video_path = r'D:\AIM\lemon\Lemon Videos 1\Lemon Videos\test\7.mp4' +video_path_out = r'D:\AIM\lemon\Lemon Videos 1\Lemon Videos\test\7_out.mp4' +# video_path = r'D:\AIM\lemon\Lemon Videos 1\Lemon Videos\Bad Lemons\8.mp4' +# video_path_out = r'D:\AIM\lemon\Lemon Videos 1\Lemon Videos\Bad Lemons\8_out.mp4' + +# Load the YOLO model +model = YOLO(r"D:\AIM\lemon\runs\detect\train4\weights\best.pt") # Load the custom model + +# Initialize VideoWriter to save the output video +cap = cv2.VideoCapture(video_path) +fps = int(cap.get(cv2.CAP_PROP_FPS)) +width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) +height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) +fourcc = cv2.VideoWriter_fourcc(*'mp4v') +out = cv2.VideoWriter(video_path_out, fourcc, fps, (width, height)) + +# Dictionary to track the state of each lemon +lemon_states = {} # Format: {ID: "State"} + +# Define class labels +class_labels = { + 0: "Bruised", + 1: "DefectiveLemon", + 2: "GoodLemon", + 3: "NotRipeLemon", + 4: "Rotten" +} + +# Classes that require ID assignment +id_tracked_classes = ["DefectiveLemon", "GoodLemon", "NotRipeLemon"] + +# Perform object tracking using BoT-SORT +results = model.track(source=video_path, conf=0.5, tracker='botsort.yaml', show=False) + +for result in results: + frame = result.orig_img # Current frame + detections = result.boxes # Detection box information + + for box in detections: + x1, y1, x2, y2 = map(int, box.xyxy[0]) # Detection box coordinates + obj_id = int(box.id) if box.id is not None else -1 # Tracking object ID + class_id = int(box.cls) # Class ID + score = box.conf # Confidence score + label = class_labels.get(class_id, "Unknown") # Get class label + + # Update lemon state and output information for tracked boxes + if obj_id != -1 and label in id_tracked_classes: + if obj_id not in lemon_states: + lemon_states[obj_id] = label + else: + # Once detected as "DefectiveLemon," the state remains "DefectiveLemon" + if lemon_states[obj_id] != "DefectiveLemon": + lemon_states[obj_id] = label + + # Output ID, position, and label + position = f"({x1}, {y1}, {x2}, {y2})" + print(f"ID: {obj_id}, Position: {position}, Label: {lemon_states[obj_id]}") + + # Draw detection boxes and labels (including untracked ones) + if obj_id != -1 and label in id_tracked_classes: + display_text = f"ID {obj_id} | {lemon_states[obj_id]}" + else: + display_text = label # For untracked labels, only show the class + + cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2) + cv2.putText(frame, display_text, (x1, y1 - 10), + cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2) + + # Write the processed frame to the output video + out.write(frame) + +# Release resources +cap.release() +out.release() +print("Video processing completed, and the result has been saved.") diff --git a/Test_video.py b/Test_video.py new file mode 100644 index 0000000..a8eb8f1 --- /dev/null +++ b/Test_video.py @@ -0,0 +1,52 @@ +import os +import pandas as pd +from ultralytics import YOLO +import cv2 + +video_path = r'D:\AIM\lemon\test.mp4' +video_path_out = r'D:\AIM\lemon\test_out.mp4' + +cap = cv2.VideoCapture(video_path) +ret, frame = cap.read() +H, W, _ = frame.shape +out = cv2.VideoWriter(video_path_out, cv2.VideoWriter_fourcc(*'MP4V'), int(cap.get(cv2.CAP_PROP_FPS)), (W, H)) + +model_path = os.path.join('.', 'runs', 'detect', 'train', 'weights', 'last.pt') + +# Load a model +model = YOLO(r"D:\AIM\lemon\YOLO-Training\YOLOv11_finetune\weights\best.pt") # load a custom model + +threshold = 0.5 +classifications = [] + +while ret: + + results = model(frame)[0] + + for result in results.boxes.data.tolist(): + x1, y1, x2, y2, score, class_id = result + + if score > threshold: + cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 4) + cv2.putText(frame, results.names[int(class_id)].upper(), (int(x1), int(y1 - 10)), + cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 255, 0), 3, cv2.LINE_AA) + + # percentage = + + # if results.names[int(class_id)] in classifications: + + # else: + classifications.append(results.names[int(class_id)]) + + out.write(frame) + ret, frame = cap.read() + +# Create a new DataFrame from the classifications +new_df = pd.DataFrame(classifications) + +# Write the new DataFrame to an Excel file +new_df.to_excel('output_classifications.xlsx', index=False) + +cap.release() +out.release() +cv2.destroyAllWindows() \ No newline at end of file diff --git a/Train.py b/Train.py new file mode 100644 index 0000000..1c5d41c --- /dev/null +++ b/Train.py @@ -0,0 +1,33 @@ +from ultralytics import YOLO +import os + +os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" +os.environ["CUDA_VISIBLE_DEVICES"] = "0" + +def main(): + # 初始化 Weights & Biases + import wandb + wandb.login(key='7cfbcf76a18a8441b04eb5d7adb988e69a79705e') # 替换为你的 API 密钥 + wandb.init(project="YOLO-Training", name="YOLOv11_finetune", mode="online") + + # 加载 YOLO 模型 + model = YOLO('yolo11n.pt') + + # 开始训练 + model.train( + data=r"D:\AIM\lemon\lemon_quality_dataset_YOLO11\data.yaml", # 数据集配置文件路径 + epochs=2, # 训练轮次 + imgsz=640, # 图像大小 + batch=16, # 批量大小 + lr0=0.005, # 初始学习率 + workers=0, # 设置为 0，禁用多进程数据加载 + device=0, # 设备 (0 = 第一块 GPU, 'cpu' = CPU) + project="YOLO-Training", # W&B 项目名称 + name="YOLOv11_finetune", # W&B 实验名称 + exist_ok=True # 如果目录存在是否覆盖 + ) + +if __name__ == '__main__': + import torch.multiprocessing as mp + mp.freeze_support() # 解决 Windows 下的多进程启动问题 + main() diff --git a/Train_model.py b/Train_model.py new file mode 100644 index 0000000..fb16ef3 --- /dev/null +++ b/Train_model.py @@ -0,0 +1,19 @@ +from ultralytics import YOLO +import os + +os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" +os.environ["CUDA_VISIBLE_DEVICES"]="0" + + +if __name__ == '__main__': + model = YOLO('yolo11n.pt') + + results = model.train( + data = "D:\AIM\lemon\AIM.v3i.yolov11\data.yaml", # Path to the dataset (YAML file for COCO128) + epochs = 50, # Number of epochs for fine-tuning + imgsz = 640, # Image size for training + batch = 16, # Batch size + lr0 = 0.005, # Initial learning rate + workers = 4, # Number of data loading workers + device = 0 # Specify the device (0 = first GPU, 'cpu' for CPU) + ) \ No newline at end of file diff --git a/YOLO-Training/YOLOv11_finetune/weights/best.pt b/YOLO-Training/YOLOv11_finetune/weights/best.pt new file mode 100644 index 0000000..e0c2c4c Binary files /dev/null and b/YOLO-Training/YOLOv11_finetune/weights/best.pt differ diff --git a/YOLO-Training/YOLOv11_finetune/weights/last.pt b/YOLO-Training/YOLOv11_finetune/weights/last.pt new file mode 100644 index 0000000..295ad3d Binary files /dev/null and b/YOLO-Training/YOLOv11_finetune/weights/last.pt differ diff --git a/runs/detect/train/weights/best.pt b/runs/detect/train/weights/best.pt new file mode 100644 index 0000000..79a5154 Binary files /dev/null and b/runs/detect/train/weights/best.pt differ diff --git a/runs/detect/train/weights/last.pt b/runs/detect/train/weights/last.pt new file mode 100644 index 0000000..bc9b229 Binary files /dev/null and b/runs/detect/train/weights/last.pt differ diff --git a/runs/detect/train2/weights/best.pt b/runs/detect/train2/weights/best.pt new file mode 100644 index 0000000..455f9db Binary files /dev/null and b/runs/detect/train2/weights/best.pt differ diff --git a/runs/detect/train2/weights/last.pt b/runs/detect/train2/weights/last.pt new file mode 100644 index 0000000..2db4ab2 Binary files /dev/null and b/runs/detect/train2/weights/last.pt differ diff --git a/runs/detect/train3/weights/best.pt b/runs/detect/train3/weights/best.pt new file mode 100644 index 0000000..1936485 Binary files /dev/null and b/runs/detect/train3/weights/best.pt differ diff --git a/runs/detect/train3/weights/last.pt b/runs/detect/train3/weights/last.pt new file mode 100644 index 0000000..bdf5843 Binary files /dev/null and b/runs/detect/train3/weights/last.pt differ diff --git a/runs/detect/train4/weights/best.pt b/runs/detect/train4/weights/best.pt new file mode 100644 index 0000000..6ad6336 Binary files /dev/null and b/runs/detect/train4/weights/best.pt differ diff --git a/runs/detect/train4/weights/last.pt b/runs/detect/train4/weights/last.pt new file mode 100644 index 0000000..a4c1091 Binary files /dev/null and b/runs/detect/train4/weights/last.pt differ diff --git a/runs/detect/train5/weights/best.pt b/runs/detect/train5/weights/best.pt new file mode 100644 index 0000000..7b311cf Binary files /dev/null and b/runs/detect/train5/weights/best.pt differ diff --git a/runs/detect/train5/weights/last.pt b/runs/detect/train5/weights/last.pt new file mode 100644 index 0000000..c59fbda Binary files /dev/null and b/runs/detect/train5/weights/last.pt differ diff --git a/yolo11n.pt b/yolo11n.pt new file mode 100644 index 0000000..45b273b Binary files /dev/null and b/yolo11n.pt differ diff --git a/yolov8n.pt b/yolov8n.pt new file mode 100644 index 0000000..0db4ca4 Binary files /dev/null and b/yolov8n.pt differ