mirrors
/
DINOv
mirror of https://github.com/UX-Decoder/DINOv.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Update dinov.py

main
Feng Li 2024-04-06 19:49:24 +08:00 committed by GitHub
parent 813f7ea681
commit 7ee80e88cf
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
dinov/architectures/dinov.py
View File

@ -953,6 +953,7 @@ class DINOv(nn.Module):
id_start = 0
id_start_list = [0]
empty_flag = False
batch_examples_by_category = {}
for targets_per_image in targets:
gt_boxes = targets_per_image.gt_boxes if torch.is_tensor(
targets_per_image.gt_boxes) else targets_per_image.gt_boxes.tensor
@ -1039,7 +1040,7 @@ class DINOv(nn.Module):
unique_category_examples_by_category = getIdx(gt_classes, id_start)
id_start += num_mask
id_start_list.append(id_start)
batch_examples_by_category = {}
for k, v in unique_category_examples_by_category.items():
if k in batch_examples_by_category.keys():
batch_examples_by_category[k] = torch.cat(
@ -1481,4 +1482,4 @@ class DINOv(nn.Module):
@register_model
def get_segmentation_model(cfg, **kwargs):
return DINOv(cfg)
return DINOv(cfg)