nielsr HF staff commited on
Commit
cce7a46
1 Parent(s): 4297151

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +13 -4
README.md CHANGED
@@ -47,12 +47,21 @@ image = Image.open(requests.get(url, stream=True).raw)
47
  feature_extractor = DetrFeatureExtractor.from_pretrained('facebook/detr-resnet-101-panoptic')
48
  model = DetrForSegmentation.from_pretrained('facebook/detr-resnet-101-panoptic')
49
 
 
50
  inputs = feature_extractor(images=image, return_tensors="pt")
 
 
51
  outputs = model(**inputs)
52
- # model predicts COCO classes, bounding boxes, and masks
53
- logits = outputs.logits
54
- bboxes = outputs.pred_boxes
55
- masks = outputs.pred_masks
 
 
 
 
 
 
56
  ```
57
 
58
  Currently, both the feature extractor and model support PyTorch.
 
47
  feature_extractor = DetrFeatureExtractor.from_pretrained('facebook/detr-resnet-101-panoptic')
48
  model = DetrForSegmentation.from_pretrained('facebook/detr-resnet-101-panoptic')
49
 
50
+ # prepare inputs for the model
51
  inputs = feature_extractor(images=image, return_tensors="pt")
52
+
53
+ # forward pass
54
  outputs = model(**inputs)
55
+
56
+ # use the `post_process_panoptic` method of `DetrFeatureExtractor` to convert to COCO format
57
+ processed_sizes = torch.as_tensor(inputs["pixel_values"].shape[-2:]).unsqueeze(0)
58
+ result = feature_extractor.post_process_panoptic(outputs, processed_sizes)[0]
59
+
60
+ # the segmentation is stored in a special-format png
61
+ panoptic_seg = Image.open(io.BytesIO(result["png_string"]))
62
+ panoptic_seg = numpy.array(panoptic_seg, dtype=numpy.uint8)
63
+ # retrieve the ids corresponding to each mask
64
+ panoptic_seg_id = rgb_to_id(panoptic_seg)
65
  ```
66
 
67
  Currently, both the feature extractor and model support PyTorch.