Skip to content

Instantly share code, notes, and snippets.

@Chris-hughes10
Created January 27, 2022 14:53
Show Gist options
  • Save Chris-hughes10/ba0790c55c375e246904f33a4a9ededf to your computer and use it in GitHub Desktop.
Save Chris-hughes10/ba0790c55c375e246904f33a4a9ededf to your computer and use it in GitHub Desktop.
timm blog - custom parser
from pathlib import Path
from timm.data.parsers.parser import Parser
class ParserImageName(Parser):
def __init__(self, root, class_to_idx=None):
super().__init__()
self.root = Path(root)
self.samples = list(self.root.glob("*.jpg"))
if class_to_idx:
self.class_to_idx = class_to_idx
else:
classes = sorted(
set([self.__extract_label_from_path(p) for p in self.samples]),
key=lambda s: s.lower(),
)
self.class_to_idx = {c: idx for idx, c in enumerate(classes)}
def __extract_label_from_path(self, path):
return "_".join(path.parts[-1].split("_")[0:-1])
def __getitem__(self, index):
path = self.samples[index]
target = self.class_to_idx[self.__extract_label_from_path(path)]
return open(path, "rb"), target
def __len__(self):
return len(self.samples)
def _filename(self, index, basename=False, absolute=False):
filename = self.samples[index][0]
if basename:
filename = filename.parts[-1]
elif not absolute:
filename = filename.absolute()
return filename
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment