Rhis is my code when I am trying to crop and save detected object with YOLOV 8 :
I am new for YOLOV show first time I code this on Google Colab notebook:
from google.colab import drive
drive.mount('/content/drive')
!pip install torch torchvision opencv-python
!pip install ultralytics --quiet
import torch
torch.cuda.is_available()
import cv2
import numpy as np
import os
from google.colab.patches import cv2_imshow
from ultralytics import YOLO
# Load your model
model_path = '/content/drive/MyDrive/Data/runs/detect/train7/weights/best.pt'
model = YOLO(model_path)
# Root directory containing subdirectories of images
root_dir = '/content/drive/My Drive/content/drive/MyDrive/Data/out_test' # Update with the correct path to your root directory
# Directory to save cropped images
save_dir = '/content/drive/My Drive/cropped_images'
os.makedirs(save_dir, exist_ok=True)
def crop_objects(img, results, save_dir, img_name):
# Get the bounding boxes
boxes = results[0].boxes.xyxy # Extracting bounding boxes from results
cropped_images = []
for i, box in enumerate(boxes):
x1, y1, x2, y2 = map(int, box[:4])
# Crop the image
cropped_img = img[y1:y2, x1:x2]
cropped_images.append(cropped_img)
# Save the cropped image
output_path = os.path.join(save_dir, f'{img_name}_cropped_{i}.jpg')
cv2.imwrite(output_path, cropped_img)
print(f'Cropped image saved to {output_path}')
return cropped_images
# Walk through all subdirectories and process images
for subdir, dirs, files in os.walk(root_dir):
for file in files:
img_path = os.path.join(subdir, file)
if img_path.lower().endswith(('png', 'jpg', 'jpeg')):
print(f"Processing image: {img_path}")
img = cv2.imread(img_path)
if img is None:
print(f"Could not read image: {img_path}")
continue
img_name = os.path.splitext(file)[0]
results = model(img) # Run inference on the image
# Debugging: Check results
print(f"Results for {img_path}: {results}")
crop_objects(img, results, save_dir, img_name)
best.pt is my model after Ii train.
How it is working but no error, but I don't have cropped images? what I can do to fix this code?
My output is cropped images in a folder.
I’d suggest you do some troubleshooting first. Try plt.imshow(cropped_images[0]) (importing matplotlib.pyplot as plt) to see whether the cropped images are even being correctly created or not. If there is a correctly cropped output, make sure that the path to which you are writing with cv2.imwrite() is correct, or else the cropped images will not save. If there is not a correct output, double check that the coordinate values you are inputting are correct, the path to your data for the model is correct, and try displaying the output to confirm that the model is indeed even detecting anything in your data. You could also play around with this to make sure your model is detecting anything:
if not len(boxes) == 0:
print(“At least there’s a detection!”)
In the past when I have faced this issue, the above tricks have worked for me. If that doesn’t help, do let me know. All the best!
Edit: As for predicting on and cropping multiple images from multiple folders, I would strongly suggest that you move all your images to a single directory for ease of access. Then, you could easily iterate through all the images. Here is a simplified example implementation:
from ultralytics import YOLO
import cv2
import os
model = YOLO("path/to/best.pt")
root_dir = "path/to/directory/with/all/images"
save_dir = "path/to/save/directory"
# define your crop_objects function here
for img_file in os.listdir(root_dir):
img_path = os.path.join(root_dir, img_file)
img = cv2.imread(img_path)
results = model(img_path)
crop_objects(img, results, save_dir, img_file)
Hope this helps!