search

voxel51 / fiftyone

Refine high-quality datasets and visual AI models
https://fiftyone.ai
Apache License 2.0
8.72k stars 550 forks source link

Fix YOLO-NAS inference #4429

Closed brimoor closed 4 months ago

brimoor commented 4 months ago

Resolves https://github.com/voxel51/fiftyone/issues/4428

Context

@jacobmarks as far as I can tell, TorchYoloNasModel was previously broken because its _predict_all() method returned a single Detections instance, but the TorchImageModel contract requires _predict_all() to return a list of Detections.

Note that, as currently implemented, TorchYoloNasModel does not support batching because it produces a model with ragged_batches == True:

from super_gradients.training import models
import fiftyone.utils.super_gradients as fousg

sg_model = models.get("yolo_nas_m", pretrained_weights="coco")
fo_model = fousg._convert_yolo_nas_detection_model(sg_model)

print(fo_model.ragged_batches)  # True

which means that apply_model() will not allow a batch_size to be provided: https://github.com/voxel51/fiftyone/blob/6faa5d72f61a566f58b85c4a8dbd8df39b9f4e18/fiftyone/core/models.py#L1807

If the underlying SuperGradients model instance does in fact support batching, then the TorchYoloNasModel wrapper would need to be tweaked to allow this to occur.

In the meantime, this PR fixes inference by implementing predict() and predict_all() correctly.

Example usage

The following now works as expected:

pip install super-gradients
import fiftyone as fo
import fiftyone.zoo as foz

from super_gradients.training import models

dataset = foz.load_zoo_dataset("quickstart", max_samples=25)
dataset.select_fields().keep_fields()

model = models.get("yolo_nas_m", pretrained_weights="coco")

dataset.apply_model(model, label_field="yolo_nas", confidence_thresh=0.7)

session = fo.launch_app(dataset)

Summary by CodeRabbit

coderabbitai[bot] commented 4 months ago

Walkthrough

The refactoring in fiftyone/utils/super_gradients.py simplifies the _convert_yolo_nas_detection_model function and modifies the TorchYoloNasModel class, improving model initialization and prediction handling. These changes are aimed at resolving issues with YOLO-NAS inference, ensuring correct detection visualization in FiftyOne's Web UI.

Changes

Files/Modules Change Summary
fiftyone/utils/super_gradients.py Refactored _convert_yolo_nas_detection_model to simplify config_model setup and modified TorchYoloNasModel to handle model loading and prediction differently.

Sequence Diagram(s) (Beta)

sequenceDiagram
    participant User
    participant FiftyOne
    participant TorchYoloNasModel
    participant SuperGradients

    User->>FiftyOne: Apply YOLO-NAS model on dataset
    FiftyOne->>TorchYoloNasModel: Initialize model
    TorchYoloNasModel->>SuperGradients: Load YOLO-NAS model
    SuperGradients-->>TorchYoloNasModel: Return model
    TorchYoloNasModel-->>FiftyOne: Model initialized
    FiftyOne->>TorchYoloNasModel: Predict on images
    TorchYoloNasModel->>SuperGradients: Perform prediction
    SuperGradients-->>TorchYoloNasModel: Return predictions
    TorchYoloNasModel-->>FiftyOne: Return predictions
    FiftyOne-->>User: Display predictions in Web UI

Assessment against linked issues

Objective Addressed Explanation
Resolve YOLO-NAS inference issue in FiftyOne's Web UI (#4428) ✅
Ensure detections appear correctly in labels_batch and labels (#4428) ✅

In code's embrace, a bug was found, With tweaks and care, it’s now unbound. Predictions clear, on screen they show, The YOLO-NAS, in perfect flow. FiftyOne's joy, a fix profound.

Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media?

Share - [X](https://twitter.com/intent/tweet?text=I%20just%20used%20%40coderabbitai%20for%20my%20code%20review%2C%20and%20it%27s%20fantastic%21%20It%27s%20free%20for%20OSS%20and%20offers%20a%20free%20trial%20for%20the%20proprietary%20code.%20Check%20it%20out%3A&url=https%3A//coderabbit.ai) - [Mastodon](https://mastodon.social/share?text=I%20just%20used%20%40coderabbitai%20for%20my%20code%20review%2C%20and%20it%27s%20fantastic%21%20It%27s%20free%20for%20OSS%20and%20offers%20a%20free%20trial%20for%20the%20proprietary%20code.%20Check%20it%20out%3A%20https%3A%2F%2Fcoderabbit.ai) - [Reddit](https://www.reddit.com/submit?title=Great%20tool%20for%20code%20review%20-%20CodeRabbit&text=I%20just%20used%20CodeRabbit%20for%20my%20code%20review%2C%20and%20it%27s%20fantastic%21%20It%27s%20free%20for%20OSS%20and%20offers%20a%20free%20trial%20for%20proprietary%20code.%20Check%20it%20out%3A%20https%3A//coderabbit.ai) - [LinkedIn](https://www.linkedin.com/sharing/share-offsite/?url=https%3A%2F%2Fcoderabbit.ai&mini=true&title=Great%20tool%20for%20code%20review%20-%20CodeRabbit&summary=I%20just%20used%20CodeRabbit%20for%20my%20code%20review%2C%20and%20it%27s%20fantastic%21%20It%27s%20free%20for%20OSS%20and%20offers%20a%20free%20trial%20for%20proprietary%20code)
Tips ### Chat There are 3 ways to chat with [CodeRabbit](https://coderabbit.ai): - Review comments: Directly reply to a review comment made by CodeRabbit. Example: - `I pushed a fix in commit .` - `Generate unit testing code for this file.` - `Open a follow-up GitHub issue for this discussion.` - Files and specific lines of code (under the "Files changed" tab): Tag `@coderabbitai` in a new review comment at the desired location with your query. Examples: - `@coderabbitai generate unit testing code for this file.` - `@coderabbitai modularize this function.` - PR comments: Tag `@coderabbitai` in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples: - `@coderabbitai generate interesting stats about this repository and render them as a table.` - `@coderabbitai show all the console.log statements in this repository.` - `@coderabbitai read src/utils.ts and generate unit testing code.` - `@coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.` Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. ### CodeRabbit Commands (invoked as PR comments) - `@coderabbitai pause` to pause the reviews on a PR. - `@coderabbitai resume` to resume the paused reviews. - `@coderabbitai review` to trigger an incremental review. This is useful when automatic reviews are disabled for the repository. - `@coderabbitai full review` to full the review from scratch and review all the files again. - `@coderabbitai summary` to regenerate the summary of the PR. - `@coderabbitai resolve` resolve all the CodeRabbit review comments. - `@coderabbitai help` to get help. Additionally, you can add `@coderabbitai ignore` anywhere in the PR description to prevent this PR from being reviewed. ### CodeRabbit Configration File (`.coderabbit.yaml`) - You can programmatically configure CodeRabbit by adding a `.coderabbit.yaml` file to the root of your repository. - Please see the [configuration documentation](https://docs.coderabbit.ai/guides/configure-coderabbit) for more information. - If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: `# yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json` ### Documentation and Community - Visit our [Documentation](https://coderabbit.ai/docs) for detailed information on how to use CodeRabbit. - Join our [Discord Community](https://discord.com/invite/GsXnASn26c) to get help, request features, and share feedback. - Follow us on [X/Twitter](https://twitter.com/coderabbitai) for updates and announcements.