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.
77 lines
2.4 KiB
Python
77 lines
2.4 KiB
Python
# Copied from https://github.com/mlcommons/training/blob/637c82f9e699cd6caf108f92efb2c1d446b630e0/single_stage_detector/ssd/transforms.py
|
|
|
|
import torch
|
|
import torchvision
|
|
|
|
from torch import nn, Tensor
|
|
from torchvision.transforms import functional as F
|
|
from torchvision.transforms import transforms as T
|
|
from typing import List, Tuple, Dict, Optional
|
|
|
|
from PIL import Image
|
|
Image.MAX_IMAGE_PIXELS = None
|
|
from typing import Any
|
|
|
|
try:
|
|
import accimage
|
|
except ImportError:
|
|
accimage = None
|
|
|
|
@torch.jit.unused
|
|
def _is_pil_image(img: Any) -> bool:
|
|
if accimage is not None:
|
|
return isinstance(img, (Image.Image, accimage.Image))
|
|
else:
|
|
return isinstance(img, Image.Image)
|
|
|
|
def get_image_size_tensor(img: Tensor) -> List[int]:
|
|
# Returns (w, h) of tensor image
|
|
torchvision.transforms._functional_tensor._assert_image_tensor(img)
|
|
return [img.shape[-1], img.shape[-2]]
|
|
|
|
@torch.jit.unused
|
|
def get_image_size_pil(img: Any) -> List[int]:
|
|
if _is_pil_image(img):
|
|
return list(img.size)
|
|
raise TypeError("Unexpected type {}".format(type(img)))
|
|
|
|
def get_image_size(img: Tensor) -> List[int]:
|
|
"""Returns the size of an image as [width, height].
|
|
Args:
|
|
img (PIL Image or Tensor): The image to be checked.
|
|
Returns:
|
|
List[int]: The image size.
|
|
"""
|
|
if isinstance(img, torch.Tensor):
|
|
return get_image_size_tensor(img)
|
|
|
|
return get_image_size_pil(img)
|
|
|
|
class Compose(object):
|
|
def __init__(self, transforms):
|
|
self.transforms = transforms
|
|
|
|
def __call__(self, image, target):
|
|
for t in self.transforms:
|
|
image, target = t(image, target)
|
|
return image, target
|
|
|
|
|
|
class RandomHorizontalFlip(T.RandomHorizontalFlip):
|
|
def forward(self, image: Tensor,
|
|
target: Optional[Dict[str, Tensor]] = None) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
|
|
if torch.rand(1) < self.p:
|
|
image = F.hflip(image)
|
|
if target is not None:
|
|
width, _ = get_image_size(image)
|
|
target["boxes"][:, [0, 2]] = width - target["boxes"][:, [2, 0]]
|
|
if "masks" in target:
|
|
target["masks"] = target["masks"].flip(-1)
|
|
return image, target
|
|
|
|
|
|
class ToTensor(nn.Module):
|
|
def forward(self, image: Tensor,
|
|
target: Optional[Dict[str, Tensor]] = None) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
|
|
image = F.to_tensor(image)
|
|
return image, target |