mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-14 08:23:59 +08:00
firestar5683
22a8575199
Keep official openpilot ancestry and dependency licenses. Store resources directly in Git, remove LFS setup requirements, and leave the downloadable big model outside the source tree.
89 lines
2.9 KiB
Python
89 lines
2.9 KiB
Python
# Copied from https://github.com/mlcommons/training/blob/637c82f9e699cd6caf108f92efb2c1d446b630e0/single_stage_detector/ssd/coco_utils.py
|
|
|
|
import os
|
|
import torch
|
|
import torchvision
|
|
|
|
from test.external.mlperf_retinanet import transforms as T
|
|
|
|
class ConvertCocoPolysToMask(object):
|
|
def __init__(self, filter_iscrowd=True):
|
|
self.filter_iscrowd = filter_iscrowd
|
|
|
|
def __call__(self, image, target):
|
|
w, h = image.size
|
|
|
|
image_id = target["image_id"]
|
|
image_id = torch.tensor([image_id])
|
|
|
|
anno = target["annotations"]
|
|
|
|
if self.filter_iscrowd:
|
|
anno = [obj for obj in anno if obj['iscrowd'] == 0]
|
|
|
|
boxes = [obj["bbox"] for obj in anno]
|
|
# guard against no boxes via resizing
|
|
boxes = torch.as_tensor(boxes, dtype=torch.float32).reshape(-1, 4)
|
|
boxes[:, 2:] += boxes[:, :2]
|
|
boxes[:, 0::2].clamp_(min=0, max=w)
|
|
boxes[:, 1::2].clamp_(min=0, max=h)
|
|
|
|
classes = [obj["category_id"] for obj in anno]
|
|
classes = torch.tensor(classes, dtype=torch.int64)
|
|
|
|
keypoints = None
|
|
if anno and "keypoints" in anno[0]:
|
|
keypoints = [obj["keypoints"] for obj in anno]
|
|
keypoints = torch.as_tensor(keypoints, dtype=torch.float32)
|
|
num_keypoints = keypoints.shape[0]
|
|
if num_keypoints:
|
|
keypoints = keypoints.view(num_keypoints, -1, 3)
|
|
|
|
keep = (boxes[:, 3] > boxes[:, 1]) & (boxes[:, 2] > boxes[:, 0])
|
|
boxes = boxes[keep]
|
|
classes = classes[keep]
|
|
|
|
target = {}
|
|
target["boxes"] = boxes
|
|
target["labels"] = classes
|
|
target["image_id"] = image_id
|
|
|
|
# for conversion to coco api
|
|
area = torch.tensor([obj["area"] for obj in anno])
|
|
iscrowd = torch.tensor([obj["iscrowd"] for obj in anno])
|
|
target["area"] = area
|
|
target["iscrowd"] = iscrowd
|
|
|
|
return image, target
|
|
|
|
class CocoDetection(torchvision.datasets.CocoDetection):
|
|
def __init__(self, img_folder, ann_file, transforms):
|
|
super(CocoDetection, self).__init__(img_folder, ann_file)
|
|
self._transforms = transforms
|
|
|
|
def __getitem__(self, idx):
|
|
img, target = super(CocoDetection, self).__getitem__(idx)
|
|
image_id = self.ids[idx]
|
|
target = dict(image_id=image_id, annotations=target)
|
|
if self._transforms is not None:
|
|
img, target = self._transforms(img, target)
|
|
return img, target
|
|
|
|
def get_openimages(name, root, image_set, transforms):
|
|
PATHS = {
|
|
"train": os.path.join(root, "train"),
|
|
"val": os.path.join(root, "validation"),
|
|
}
|
|
|
|
t = [ConvertCocoPolysToMask(filter_iscrowd=False)]
|
|
|
|
if transforms is not None:
|
|
t.append(transforms)
|
|
transforms = T.Compose(t)
|
|
|
|
img_folder = os.path.join(PATHS[image_set], "data")
|
|
ann_file = os.path.join(PATHS[image_set], "labels", f"{name}.json")
|
|
|
|
dataset = CocoDetection(img_folder, ann_file, transforms=transforms)
|
|
|
|
return dataset |