Fix/update example

#3
by asmith26 - opened
Files changed (1) hide show
  1. README.md +7 -7
README.md CHANGED
@@ -47,24 +47,24 @@ from PIL import Image
47
  import torch
48
  import numpy
49
 
50
- from transformers import DetrFeatureExtractor, DetrForSegmentation
51
- from transformers.models.detr.feature_extraction_detr import rgb_to_id
52
 
53
  url = "http://images.cocodataset.org/val2017/000000039769.jpg"
54
  image = Image.open(requests.get(url, stream=True).raw)
55
 
56
- feature_extractor = DetrFeatureExtractor.from_pretrained("facebook/detr-resnet-50-panoptic")
57
  model = DetrForSegmentation.from_pretrained("facebook/detr-resnet-50-panoptic")
58
 
59
  # prepare image for the model
60
- inputs = feature_extractor(images=image, return_tensors="pt")
61
 
62
  # forward pass
63
  outputs = model(**inputs)
64
 
65
- # use the `post_process_panoptic` method of `DetrFeatureExtractor` to convert to COCO format
66
  processed_sizes = torch.as_tensor(inputs["pixel_values"].shape[-2:]).unsqueeze(0)
67
- result = feature_extractor.post_process_panoptic(outputs, processed_sizes)[0]
68
 
69
  # the segmentation is stored in a special-format png
70
  panoptic_seg = Image.open(io.BytesIO(result["png_string"]))
@@ -73,7 +73,7 @@ panoptic_seg = numpy.array(panoptic_seg, dtype=numpy.uint8)
73
  panoptic_seg_id = rgb_to_id(panoptic_seg)
74
  ```
75
 
76
- Currently, both the feature extractor and model support PyTorch.
77
 
78
  ## Training data
79
 
 
47
  import torch
48
  import numpy
49
 
50
+ from transformers import DetrForSegmentation, DetrImageProcessor
51
+ from transformers.image_transforms import rgb_to_id
52
 
53
  url = "http://images.cocodataset.org/val2017/000000039769.jpg"
54
  image = Image.open(requests.get(url, stream=True).raw)
55
 
56
+ image_processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50-panoptic")
57
  model = DetrForSegmentation.from_pretrained("facebook/detr-resnet-50-panoptic")
58
 
59
  # prepare image for the model
60
+ inputs = image_processor(images=image, return_tensors="pt")
61
 
62
  # forward pass
63
  outputs = model(**inputs)
64
 
65
+ # use the `post_process_panoptic` method of `DetrImageProcessor` to convert to COCO format
66
  processed_sizes = torch.as_tensor(inputs["pixel_values"].shape[-2:]).unsqueeze(0)
67
+ result = image_processor.post_process_panoptic(outputs, processed_sizes)[0]
68
 
69
  # the segmentation is stored in a special-format png
70
  panoptic_seg = Image.open(io.BytesIO(result["png_string"]))
 
73
  panoptic_seg_id = rgb_to_id(panoptic_seg)
74
  ```
75
 
76
+ Currently, both the image processor and model support PyTorch.
77
 
78
  ## Training data
79