HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Scope ultralytics/CLIP imports (#9449)

Browse source 
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Glenn Jocher 2024-03-31 19:07:59 +02:00 • committed by GitHub
parent 7df821e6ea
commit cd172e9d12
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 20 additions and 20 deletions
Download patch file Download diff file Expand all files Collapse all files

9
ultralytics/models/fastsam/prompt.py
View file

@ -9,7 +9,7 @@ import numpy as np
import torch import torch
from PIL import Image from PIL import Image
from ultralytics.utils import TQDM from ultralytics.utils import TQDM, checks
class FastSAMPrompt: class FastSAMPrompt:
@ -33,9 +33,7 @@ class FastSAMPrompt:
try: try:
import clip import clip
except ImportError: except ImportError:
from ultralytics.utils.checks import check_requirements checks.check_requirements("git+https://github.com/ultralytics/CLIP.git")
check_requirements("git+https://github.com/ultralytics/CLIP.git")
import clip import clip
self.clip = clip self.clip = clip
@ -115,7 +113,8 @@ class FastSAMPrompt:
points (list, optional): Points to be plotted. Defaults to None. points (list, optional): Points to be plotted. Defaults to None.
point_label (list, optional): Labels for the points. Defaults to None. point_label (list, optional): Labels for the points. Defaults to None.
mask_random_color (bool, optional): Whether to use random color for masks. Defaults to True. mask_random_color (bool, optional): Whether to use random color for masks. Defaults to True.
better_quality (bool, optional): Whether to apply morphological transformations for better mask quality. Defaults to True. better_quality (bool, optional): Whether to apply morphological transformations for better mask quality.
Defaults to True.
retina (bool, optional): Whether to use retina mask. Defaults to False. retina (bool, optional): Whether to use retina mask. Defaults to False.
with_contours (bool, optional): Whether to plot contours. Defaults to True. with_contours (bool, optional): Whether to plot contours. Defaults to True.
""" """

31
ultralytics/models/yolo/world/train.py
View file

@ -1,18 +1,12 @@
# Ultralytics YOLO 🚀, AGPL-3.0 license # Ultralytics YOLO 🚀, AGPL-3.0 license
from ultralytics.models import yolo
from ultralytics.nn.tasks import WorldModel
from ultralytics.utils import DEFAULT_CFG, RANK
from ultralytics.data import build_yolo_dataset
from ultralytics.utils.torch_utils import de_parallel
from ultralytics.utils.checks import check_requirements
import itertools import itertools
try: from ultralytics.data import build_yolo_dataset
import clip from ultralytics.models import yolo
except ImportError: from ultralytics.nn.tasks import WorldModel
check_requirements("git+https://github.com/ultralytics/CLIP.git") from ultralytics.utils import DEFAULT_CFG, RANK, checks
import clip from ultralytics.utils.torch_utils import de_parallel
def on_pretrain_routine_end(trainer): def on_pretrain_routine_end(trainer):
@ -22,10 +16,9 @@ def on_pretrain_routine_end(trainer):
names = [name.split("/")[0] for name in list(trainer.test_loader.dataset.data["names"].values())] names = [name.split("/")[0] for name in list(trainer.test_loader.dataset.data["names"].values())]
de_parallel(trainer.ema.ema).set_classes(names, cache_clip_model=False) de_parallel(trainer.ema.ema).set_classes(names, cache_clip_model=False)
device = next(trainer.model.parameters()).device device = next(trainer.model.parameters()).device
text_model, _ = clip.load("ViT-B/32", device=device) trainer.text_model, _ = trainer.clip.load("ViT-B/32", device=device)
for p in text_model.parameters(): for p in trainer.text_model.parameters():
p.requires_grad_(False) p.requires_grad_(False)
trainer.text_model = text_model
class WorldTrainer(yolo.detect.DetectionTrainer): class WorldTrainer(yolo.detect.DetectionTrainer):
@ -48,6 +41,14 @@ class WorldTrainer(yolo.detect.DetectionTrainer):
overrides = {} overrides = {}
super().__init__(cfg, overrides, _callbacks) super().__init__(cfg, overrides, _callbacks)
# Import and assign clip
try:
import clip
except ImportError:
checks.check_requirements("git+https://github.com/ultralytics/CLIP.git")
import clip
self.clip = clip
def get_model(self, cfg=None, weights=None, verbose=True): def get_model(self, cfg=None, weights=None, verbose=True):
"""Return WorldModel initialized with specified config and weights.""" """Return WorldModel initialized with specified config and weights."""
# NOTE: This `nc` here is the max number of different text samples in one image, rather than the actual `nc`. # NOTE: This `nc` here is the max number of different text samples in one image, rather than the actual `nc`.
@ -84,7 +85,7 @@ class WorldTrainer(yolo.detect.DetectionTrainer):
# NOTE: add text features # NOTE: add text features
texts = list(itertools.chain(*batch["texts"])) texts = list(itertools.chain(*batch["texts"]))
text_token = clip.tokenize(texts).to(batch["img"].device) text_token = self.clip.tokenize(texts).to(batch["img"].device)
txt_feats = self.text_model.encode_text(text_token).to(dtype=batch["img"].dtype) # torch.float32 txt_feats = self.text_model.encode_text(text_token).to(dtype=batch["img"].dtype) # torch.float32
txt_feats = txt_feats / txt_feats.norm(p=2, dim=-1, keepdim=True) txt_feats = txt_feats / txt_feats.norm(p=2, dim=-1, keepdim=True)
batch["txt_feats"] = txt_feats.reshape(len(batch["texts"]), -1, txt_feats.shape[-1]) batch["txt_feats"] = txt_feats.reshape(len(batch["texts"]), -1, txt_feats.shape[-1])