Visualizing YOLOv5 Segmentation Data

[DylDevs]

Search before asking

• I have searched the YOLOv5 issues and discussions and found no similar questions.

Question

Greetings, I would like to use YOLOv5 Segmentation for one of my projects. However, before I collect and annotate data, I want to get a proof of concept that v5-seg can do everything that I need to do.
I have the model loaded to my needs with this code:

MODEL_NAME = "yolov5s-seg.pt"
MODEL_PATH = os.path.dirname(__file__) + f"/{MODEL_NAME}"

# Load YOLO model
model = torch.hub.load('ultralytics/yolov5', 'custom', path=MODEL_PATH)
model.conf = 0.65

When it comes to retrieving data for the model, I think I did it correctly.

while True:
    # Get the current frame capture using DXcam (Screen Capture)
    frame = camera.grab(region=(capture_x, capture_y, capture_x + capture_width, capture_y + capture_height))
    if frame is None: # Failsafe
        continue

    # Make a copy of the frame for YOLO inference
    yolo_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    # Convert the image to a tensor and add a batch dimension
    img_tensor = torch.from_numpy(yolo_frame).permute(2, 0, 1).float().div(255.0).unsqueeze(0)

    # Perform inference
    results = model(img_tensor)

    # Parse the results
    if len(results) > 0 and len(results[0]) > 0:  # Check if there are any detections
        # Assuming results[0] is a tensor with the following columns: [x1, y1, x2, y2, confidence, class, *masks]
        masks = results[0][:, 6:]  # Extract masks (assuming masks are from the 7th column onwards)
        boxes = results[0][:, :4].cpu().numpy()  # Bounding boxes
        scores = results[0][:, 4].cpu().numpy()  # Confidence scores
        classes = results[0][:, 5].cpu().numpy()  # Class predictions

        # Convert the masks to the same size as the original image
        if masks is not None:
            masks = masks.cpu().numpy()  # Convert masks to numpy array

        # Iterate through masks and draw them on the image
        for mask in masks:
            mask_resized = cv2.resize(mask, (frame.shape[1], frame.shape[0]), interpolation=cv2.INTER_LINEAR)
            mask_binary = (mask_resized > 0.5).astype('uint8') * 255
            contours, _ = cv2.findContours(mask_binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            cv2.drawContours(frame, contours, -1, (0, 255, 0), 2)

    # Display the results in an OpenCV window
    cv2.imshow('Vehicle Detection', frame)

    # Wait for a key press, with a short delay to allow for video playback
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cv2.destroyAllWindows()

However, no matter what image I give it, the output is always this:

I am not sure of what i'm doing wrong, is it model loading? Is it Inference? Is it plotting the data? From my knowledge this should work just fine.
Any help on the matter is appreciated!

Additional

Also, another small question. My GPU is not really good, so I tend to train my models in Google Colab. My issue is that my model will get to 150 - 200 epochs and then the runtime will disconnect because of Inactivity. I feel like i'm checking in on it and clicking around often enough for it to not crash (every ~10 minutes or so). I understand this can be easily fixed by getting a Colab subscription, but I would like to avoid paying for Colab if at all possible. Inout on this issue is also appreciated!

[glenn-jocher]

Hello,

Thank you for reaching out with your questions about YOLOv5 segmentation and training on Google Colab.

YOLOv5 Segmentation Issue

From your description, it seems like your model setup and inference code are generally correct. However, the issue might be related to how the masks are being processed or displayed. Here are a few things you might want to check:

1. Model Output: Ensure that the model is indeed outputting mask data. You can add a print statement to check the shape and content of the masks tensor right after inference.
2. Mask Processing: Verify that the mask resizing and thresholding are being done correctly. Sometimes, discrepancies in dimensions or incorrect threshold values can lead to invisible or incorrect masks.
3. Visualization: Check if the contours are being drawn on the frame correctly. You might want to try a simpler form of visualization, like overlaying the mask directly on the image with some transparency, to see if the masks are being generated at all.

Google Colab Disconnection

Regarding the issue with Google Colab disconnecting during long training sessions, here are a couple of suggestions:

• Code Activity: Try to include some form of output or logging in your training loop that updates more frequently. This might help in keeping the session alive.
• Colab Extensions: There are browser extensions available that can help keep your Colab session active by simulating interaction. These can be a temporary workaround if frequent manual interaction isn't helping.

If the problem persists, reviewing the exact output or errors, if any, could provide more insights into what might be going wrong.

Feel free to share any updates or additional information. We're here to help!

[DylDevs]

@glenn-jocher, Thanks for the quick response!

I have been unsuccessful in getting the YOLOv5-seg model to work as expected. I notice that i get this error when loading the model:

WARNING  YOLOv5 SegmentationModel is not yet AutoShape compatible. You will not be able to run inference with this model.

I think this has something to do with it. Is there a way that you know of to load the model for inference in Python? If so, it would help if you could provide some example code to help me get started.

Also, about the Colab issue, I got a browser extension that seemingly fixed the issue!

[glenn-jocher]

@DylDevs hello,

Thank you for the update and I'm glad to hear that the browser extension resolved your Colab issue!

Regarding the YOLOv5-seg model warning, it indicates that the model isn't compatible with the AutoShape feature, which simplifies the inference process. You can still run inference manually by handling the input and output tensors directly. Here’s a basic example to help you get started with manual inference:

import torch
from PIL import Image
from torchvision.transforms import functional as F

# Load your model (ensure it's in the correct directory)
model = torch.load('path_to_yolov5s-seg.pt')
model.eval()  # Set the model to evaluation mode

# Load an image
image = Image.open('path_to_your_image.jpg')
image = F.to_tensor(image).unsqueeze(0)  # Transform image to tensor and add batch dimension

# Perform inference
with torch.no_grad():
    results = model(image)

# Process results here (e.g., extracting masks)
# Note: You'll need to adapt this part based on how your model outputs data

This code snippet manually handles the image transformation and model inference. Make sure to adapt the result processing part according to the specific output format of your segmentation model.

If you encounter any more issues or have further questions, feel free to ask. Happy coding!

[DylDevs]

@glenn-jocher
Thanks for the rsponse again. I tried your suggestion, and I got this error from it:

Traceback (most recent call last):
  File "c:\Users\Dylan\Documents\Coding\Python\ETS2 Vehicle Detection\yolov5Seg\run - Copy.py", line 17, in <module>
    model = torch.load(MODEL_PATH)
            ^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Dylan\AppData\Local\Programs\Python\Python311\Lib\site-packages\torch\serialization.py", line 1026, in load
    return _load(opened_zipfile,
           ^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Dylan\AppData\Local\Programs\Python\Python311\Lib\site-packages\torch\serialization.py", line 1438, in _load
    result = unpickler.load()
             ^^^^^^^^^^^^^^^^
  File "C:\Users\Dylan\AppData\Local\Programs\Python\Python311\Lib\site-packages\torch\serialization.py", line 1431, in find_class
    return super().find_class(mod_name, name)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ModuleNotFoundError: No module named 'models.yolo'

Any ideas as to how I can fix it?

[glenn-jocher]

Hello @DylDevs,

Thank you for your patience and for providing the error details. The error you're encountering, ModuleNotFoundError: No module named 'models.yolo', typically occurs when trying to load a model that was saved with custom classes or modules that aren't available in your current environment.

To resolve this, you should load the model using the torch.hub.load method, which ensures all dependencies are correctly handled. Here’s how you can do it:

1. Ensure you have the correct YOLOv5 repository and dependencies installed:

   pip install -r https://raw.githubusercontent.com/ultralytics/yolov5/master/requirements.txt

2. Load the model using torch.hub.load:

   import torch
   
   # Load the YOLOv5-seg model from PyTorch Hub
   model = torch.hub.load('ultralytics/yolov5', 'custom', path='path_to_yolov5s-seg.pt')
   model.eval()  # Set the model to evaluation mode
   
   # Example inference
   from PIL import Image
   from torchvision.transforms import functional as F
   
   # Load an image
   image = Image.open('path_to_your_image.jpg')
   image = F.to_tensor(image).unsqueeze(0)  # Transform image to tensor and add batch dimension
   
   # Perform inference
   with torch.no_grad():
       results = model(image)
   
   # Process results here (e.g., extracting masks)
   # Note: You'll need to adapt this part based on how your model outputs data

This approach should help you avoid the ModuleNotFoundError by leveraging the torch.hub.load method, which manages dependencies and model loading more robustly.

If you have any further questions or run into other issues, feel free to ask. We're here to help! 😊

[DylDevs]

@glenn-jocher This code does the same thing as I started with.

I have figured out that results is a list of tensors. How should I decode this to get the actual output from the model. I would like to get this sorted out as quickly as possible so i can get this implemented in my project.

Thanks!

[glenn-jocher]

Hello @DylDevs,

Thank you for your patience and for the additional details. Let's dive into decoding the results from the YOLOv5-seg model to get the actual output.

When you perform inference with the YOLOv5-seg model, the results are typically a list of tensors. Each tensor contains information about the detected objects, including bounding boxes, confidence scores, class labels, and segmentation masks.

Here's a step-by-step guide to decode and visualize the results:

1. Perform Inference:

   import torch
   from PIL import Image
   from torchvision.transforms import functional as F
   import cv2
   import numpy as np
   
   # Load the YOLOv5-seg model from PyTorch Hub
   model = torch.hub.load('ultralytics/yolov5', 'custom', path='path_to_yolov5s-seg.pt')
   model.eval()  # Set the model to evaluation mode
   
   # Load an image
   image = Image.open('path_to_your_image.jpg')
   image = F.to_tensor(image).unsqueeze(0)  # Transform image to tensor and add batch dimension
   
   # Perform inference
   with torch.no_grad():
       results = model(image)

2. Decode the Results:

   # Assuming results[0] contains the output for the first image in the batch
   output = results[0]
   
   # Extract bounding boxes, confidence scores, class labels, and masks
   boxes = output[:, :4].cpu().numpy()  # Bounding boxes
   scores = output[:, 4].cpu().numpy()  # Confidence scores
   class_ids = output[:, 5].cpu().numpy()  # Class labels
   masks = output[:, 6:].cpu().numpy()  # Segmentation masks
   
   # Convert the masks to the same size as the original image
   masks = masks.squeeze()  # Remove unnecessary dimensions
   masks = np.array([cv2.resize(mask, (image.width, image.height)) for mask in masks])

3. Visualize the Results:

   # Convert the image back to a format suitable for OpenCV
   image_cv = cv2.cvtColor(np.array(image.squeeze().permute(1, 2, 0) * 255, dtype=np.uint8), cv2.COLOR_RGB2BGR)
   
   for i, mask in enumerate(masks):
       # Threshold the mask to create a binary mask
       mask_binary = (mask > 0.5).astype(np.uint8) * 255
   
       # Find contours and draw them on the image
       contours, _ = cv2.findContours(mask_binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
       cv2.drawContours(image_cv, contours, -1, (0, 255, 0), 2)
   
       # Draw bounding boxes and labels
       x1, y1, x2, y2 = boxes[i]
       cv2.rectangle(image_cv, (int(x1), int(y1)), (int(x2), int(y2)), (255, 0, 0), 2)
       label = f"Class: {int(class_ids[i])}, Score: {scores[i]:.2f}"
       cv2.putText(image_cv, label, (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
   
   # Display the image with detections
   cv2.imshow('YOLOv5 Segmentation', image_cv)
   cv2.waitKey(0)
   cv2.destroyAllWindows()

This should help you decode the results and visualize the segmentation masks along with bounding boxes and class labels. If you have any further questions or need additional assistance, feel free to ask. We're here to help! 😊

[DylDevs]

With a little bit of modification, that code worked, thanks!

[glenn-jocher]

Hello @DylDevs,

I'm glad to hear that the code modifications helped! 🎉 If you encounter any further issues or have additional questions, feel free to reach out.

To ensure we can assist you most effectively, please make sure to provide a minimum reproducible code example if you run into any new bugs or issues. This helps us reproduce the problem on our end and investigate a solution more efficiently. You can find more details on how to create a minimum reproducible example here: Minimum Reproducible Example.

Additionally, always ensure you're using the latest versions of torch and the YOLOv5 repository from Ultralytics GitHub. This ensures you have the latest features and bug fixes.

If you have any more questions or need further assistance, don't hesitate to ask. The YOLO community and the Ultralytics team are here to help!