Update app.py

app.py

@@ -1,36 +1,67 @@
-import requests
-import os, io
-import gradio as gr
-# from PIL import Image
+# import requests
+# import os, io
+# import gradio as gr
+# # from PIL import Image
+
+# # API_URL = "https://api-inference.huggingface.co/models/facebook/detr-resnet-50-panoptic"
+
+# SECRET_TOKEN = os.getenv("SECRET_TOKEN")
+# API_URL = "https://api-inference.huggingface.co/models/facebook/detr-resnet-50-dc5-panoptic"
+# headers = {"Authorization": f'Bearer {SECRET_TOKEN}'}
+
+
 
-# API_URL = "https://api-inference.huggingface.co/models/facebook/detr-resnet-50-panoptic"
+# def image_classifier(inp):
+#     return {'cat': 0.3, 'dog': 0.7}
 
-SECRET_TOKEN = os.getenv("SECRET_TOKEN")
-API_URL = "https://api-inference.huggingface.co/models/facebook/detr-resnet-50-dc5-panoptic"
-headers = {"Authorization": f'Bearer {SECRET_TOKEN}'}
+# def query(filename):
+#     with open(filename, "rb") as f:
+#         data = f.read()
+#     response = requests.post(API_URL, headers=headers, data=data)
+#     return response.json()
+
+# def rb(img):
+#     # initialiaze io to_bytes converter
+#     img_byte_arr = io.BytesIO()
+#     # define quality of saved array
+#     img.save(img_byte_arr, format='JPEG', subsampling=0, quality=100)
+#     # converts image array to bytesarray
+#     img_byte_arr = img_byte_arr.getvalue()
+#     response = requests.post(API_URL, headers=headers, data=img_byte_arr)
+#     return response.json()
+
+
+# inputs = gr.inputs.Image(type="pil", label="Upload an image")
+# demo = gr.Interface(fn=rb, inputs=inputs, outputs="json")
+# demo.launch()
+
+import io
+import requests
+from PIL import Image
+import torch
+import numpy
 
+from transformers import DetrFeatureExtractor, DetrForSegmentation
+from transformers.models.detr.feature_extraction_detr import rgb_to_id
 
+url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+image = Image.open(requests.get(url, stream=True).raw)
 
-def image_classifier(inp):
-    return {'cat': 0.3, 'dog': 0.7}
+feature_extractor = DetrFeatureExtractor.from_pretrained("facebook/detr-resnet-50-panoptic")
+model = DetrForSegmentation.from_pretrained("facebook/detr-resnet-50-panoptic")
 
-def query(filename):
-    with open(filename, "rb") as f:
-        data = f.read()
-    response = requests.post(API_URL, headers=headers, data=data)
-    return response.json()
+# prepare image for the model
+inputs = feature_extractor(images=image, return_tensors="pt")
 
-def rb(img):
-    # initialiaze io to_bytes converter
-    img_byte_arr = io.BytesIO()
-    # define quality of saved array
-    img.save(img_byte_arr, format='JPEG', subsampling=0, quality=100)
-    # converts image array to bytesarray
-    img_byte_arr = img_byte_arr.getvalue()
-    response = requests.post(API_URL, headers=headers, data=img_byte_arr)
-    return response.json()
+# forward pass
+outputs = model(**inputs)
 
+# use the `post_process_panoptic` method of `DetrFeatureExtractor` to convert to COCO format
+processed_sizes = torch.as_tensor(inputs["pixel_values"].shape[-2:]).unsqueeze(0)
+result = feature_extractor.post_process_panoptic(outputs, processed_sizes)[0]
 
-inputs = gr.inputs.Image(type="pil", label="Upload an image")
-demo = gr.Interface(fn=rb, inputs=inputs, outputs="json")
-demo.launch()
+# the segmentation is stored in a special-format png
+panoptic_seg = Image.open(io.BytesIO(result["png_string"]))
+panoptic_seg = numpy.array(panoptic_seg, dtype=numpy.uint8)
+# retrieve the ids corresponding to each mask
+panoptic_seg_id = rgb_to_id(panoptic_seg)