利用SuperbAISuite和NVIDIATAOToolkit创建高质量的计算机视觉应用

{   "objects": [       {           "id": "7e9fe8ee-50c7-4d4f-9e2c-145d894a8a26",           "class_id": "7b8205ef-b251-450c-b628-e6b9cac1a457",           "class_name": "person",           "annotation_type": "box",           "annotation": {               "multiple": false,               "coord": {                   "x": 275.47,                   "y": 49.27,                   "width": 86.39999999999998,                   "height": 102.25               },               "meta": {},               "difficulty": 0,               "uncertainty": 0.0045           },           "properties": []       },       {           "id": "70257635-801f-4cad-856a-ef0fdbfdf613",           "class_id": "7b8205ef-b251-450c-b628-e6b9cac1a457",           "class_name": "person",           "annotation_type": "box",           "annotation": {               "multiple": false,               "coord": {                   "x": 155.64,                   "y": 40.61,                   "width": 98.34,                   "height": 113.05               },               "meta": {},               "difficulty": 0,               "uncertainty": 0.0127           },           "properties": []       }   ],   "categories": {       "properties": []   },   "difficulty": 1}

class SuiteDataset(Dataset):   """   Instantiate the SuiteDataset object class for training set   """
   def __init__(           self,           team_name: str,           access_key: str,           project_name: str,           export_name: str,           train: bool,           caching_image: bool = True,           transforms: Optional[List[Callable]] = None,           category_names: Optional[List[str]] = None,   ):       """Function to initialize the object class"""       super().__init__()
       # Get project setting and export information through the SDK       # Initialize the Python Client       client = spb.sdk.Client(team_name=team_name, access_key=access_key, project_name=project_name)       # Use get_export       export_info = call_with_retry(client.get_export, name=export_name)       # Download the export compressed file through download_url in Export       export_data = call_with_retry(urlopen, export_info.download_url).read()
       # Load the export compressed file into memory       with ZipFile(BytesIO(export_data), 'r') as export:           label_files = [f for f in export.namelist() if f.startswith('labels/')]           label_interface = json.loads(export.open('project.json', 'r').read())           category_infos = label_interface.get('object_detection', {}).get('object_classes', [])
       cache_dir = None       if caching_image:           cache_dir = f'/tmp/{team_name}/{project_name}'           os.makedirs(cache_dir, exist_ok=True)
       self.client = client       self.export_data = export_data       self.categories = [           {'id': i + 1, 'name': cat['name'], 'type': cat['annotation_type']}           for i, cat in enumerate(category_infos)       ]       self.category_id_map = {cat['id']: i + 1 for i, cat in enumerate(category_infos)}       self.transforms = build_transforms(train, self.categories, transforms, category_names)       self.cache_dir = cache_dir
       # Convert label_files to numpy array and use       self.label_files = np.array(label_files).astype(np.string_)
   def __len__(self):       """Function to return the number of label files"""       return len(self.label_files)
   def __getitem__(self, idx):       """Function to get an item"""       idx = idx if idx >= 0 else len(self) + idx       if idx < 0 or idx >= len(self):           raise IndexError(f'index out of range')
       image_id = idx + 1       label_file = self.label_files[idx].decode('ascii')
       # Load label information corresponding to idx from the export compressed file into memory       with ZipFile(BytesIO(self.export_data), 'r') as export:           label = load_label(export, label_file, self.category_id_map, image_id)
       # Download the image through the Suite sdk based on label_id       try:           image = load_image(self.client, label['label_id'], self.cache_dir)       # Download data in real time using get_data from Suite sdk       except Exception as e:           print(f'Failed to load the {idx}-th image due to {repr(e)}, getting {idx + 1}-th data instead')           return self.__getitem__(idx + 1)
       target = {           'image_id': image_id,           'label_id': label['label_id'],           'annotations': label['annotations'],       }
       if self.transforms is not None:           image, target = self.transforms(image, target)       return image, target

class SuiteCocoDataset(C.CocoDetection):   """   Instantiate the SuiteCocoDataset object class for test set   (by wrapping SuiteDataset to make compatible with torchvision's official COCOEvaluator)   """
   def __init__(           self,           team_name: str,           access_key: str,           project_name: str,           export_name: str,           train: bool,           caching_image: bool = True,           transforms: Optional[List[Callable]] = None,           category_names: Optional[List[str]] = None,           num_init_workers: int = 20,   ):       """Function to initialize the object class"""       super().__init__(img_folder='', ann_file=None, transforms=None)
       # Call the SuiteDataset class       dataset = SuiteDataset(           team_name, access_key, project_name, export_name,           train=False, transforms=[],           caching_image=caching_image, category_names=category_names,       )       self.client = dataset.client       self.cache_dir = dataset.cache_dir
       self.coco = build_coco_dataset(dataset, num_init_workers)       self.ids = list(sorted(self.coco.imgs.keys()))       self._transforms = build_transforms(train, dataset.categories, transforms, category_names)
   def _load_image(self, id: int):       """Function to load an image"""       label_id = self.coco.loadImgs(id)[0]['label_id']       image = load_image(self.client, label_id, self.cache_dir)       return image
   def __getitem__(self, idx):       """Function to get an item"""       try:           return super().__getitem__(idx)       except Exception as e:           print(f'Failed to load the {idx}-th image due to {repr(e)}, getting {idx + 1}-th data instead')           return self.__getitem__(idx + 1)

train_dataset = SuiteDataset(   team_name=args.team_name,   access_key=args.access_key,   project_name=args.project_name,   export_name=args.train_export_name,   caching_image=args.caching_image,   train=True,)test_dataset = SuiteCocoDataset(   team_name=args.team_name,   access_key=args.access_key,   project_name=args.project_name,   export_name=args.test_export_name,   caching_image=args.caching_image,   train=False,   num_init_workers=args.workers,)
train_loader = DataLoader(   train_dataset, num_workers=args.workers,   batch_sampler=G.GroupedBatchSampler(       RandomSampler(train_dataset),       G.create_aspect_ratio_groups(train_dataset, k=3),       args.batch_size,   ),   collate_fn=collate_fn,)test_loader = DataLoader(   test_dataset, num_workers=args.workers,   sampler=SequentialSampler(test_dataset), batch_size=1,   collate_fn=collate_fn,)

pip3 install nvidia-taowget --content-disposition https://api.ngc.nvidia.com/v2/resources/nvidia/tao/tao-getting-started/versions/4.0.1/zip -O getting_started_v4.0.1.zip


$ unzip -u getting_started_v4.0.1.zip  -d ./getting_started_v4.0.1 && rm -rf getting_started_v4.0.1.zip && cd ./getting_started_v4.0.1

$ jupyter notebook --ip 0.0.0.0 --port 8888 --allow-root

http://0.0.0.0:8888

notebooks/tao_launcher_starter_kit/yolo_v4/yolo_v4.ipynb and follow the