yolo_sam.py
- yolo_sam.display_images(image_path: str, segmented_image_path: str) None
Display the original and segmented images side by side.
This function loads two images (the original and the segmented one), and displays them in a side-by-side layout for comparison. The images are shown without axes, and each image is labeled accordingly.
- Parameters:
image_path (str): The file path to the original image. segmented_image_path (str): The file path to the segmented image.
- Returns:
None: This function only displays the images and does not return any value.
- Example:
display_images(“path/to/original/image.jpg”, “path/to/segmented/image.jpg”)
- yolo_sam.load_train_yolo_model() YOLO
Load and train a YOLO model for object detection.
This function loads the YOLO model from a pre-trained checkpoint, trains it using the provided dataset, and evaluates its performance on the validation set.
- Parameters:
None
- Returns:
YOLO: The trained YOLO model object.
- Example:
model = load_train_yolo_model()
- yolo_sam.main()
- yolo_sam.test_segmentation_sam(image_path: str, image_name: str) None
Perform segmentation using the SAM model and display the original and segmented images.
This function loads the SAM (Segment Anything Model) model, performs segmentation on the provided image, and then displays the original and segmented images side by side for comparison. The segmented image is saved, and the function also visualizes both images.
- Parameters:
image_path (str): The file path to the image to be segmented. image_name (str): The name to save the segmented image under.
- Returns:
None: This function performs segmentation and displays the images but does not return any value.
- Example:
test_segmentation_sam(“path/to/image.jpg”, “segmented_image.jpg”)
- yolo_sam.test_segmentation_yolo_sam(image_path: str, image_name: str) None
Perform segmentation using the YOLO and SAM models, and display the original and segmented images.
This function loads the SAM (Segment Anything Model) model, performs segmentation on the provided image, and then displays the original and segmented images side by side for comparison. The segmented image is saved, and the function also visualizes both images.
- Parameters:
image_path (str): The file path to the image to be segmented. image_name (str): The name to save the segmented image under.
- Returns:
None: This function performs segmentation and displays the images but does not return any value.
- Example:
test_segmentation_yolo_sam(“path/to/image.jpg”, “segmented_image.jpg”)
- yolo_sam.test_yolo_model(model: YOLO, image_path: str) None
Test the YOLO model by performing object detection on a given image.
This function uses the provided YOLO model to detect objects in the specified image and displays the results by showing the image with bounding boxes around detected objects.
- Parameters:
model (YOLO): The trained YOLO model used for object detection. image_path (str): The file path to the image on which object detection will be performed.
- Returns:
None: This function only displays the image with detection results and does not return any value.
- Example:
test_yolo_model(model, “path/to/image.jpg”)