Update app.py

app.py

@@ -25,6 +25,14 @@ model_name_to_fn = {
 
 
 ### 3. Predict function ###
+def fig2img(fig):
+    """Convert a Matplotlib figure to a PIL Image and return it"""
+    import io
+    buf = io.BytesIO()
+    fig.savefig(buf)
+    buf.seek(0)
+    img = Image.open(buf)
+    return img
 
 
 # Create predict function
@@ -67,8 +75,10 @@ def predict(image, model_name: str = "detr",) -> Tuple[Dict, float]:
         # A tensor of shape (height, width) where each value denotes a segment id, filled with -1 if no segment is found
         panoptic_seg = result[0]["segmentation"]
         # Convert the tensor to PIL image
-        plt.imsave("predicted_panoptic_map.png", panoptic_seg, cmap="viridis")
-        output = PIL.Image.open("predicted_panoptic_map.png")
+        plt.plot("predicted_panoptic_map.png", panoptic_seg, cmap="viridis")
+        fig = plt.gcf()
+        output = fig2img(fig)
+        # output = PIL.Image.open("predicted_panoptic_map.png")
         # output = PIL.Image.fromarray(panoptic_seg.cpu().numpy().astype('uint8')).convert('RGB')
     
     elif model_name == "maskformer":
@@ -84,8 +94,10 @@ def predict(image, model_name: str = "detr",) -> Tuple[Dict, float]:
         result = processor.post_process_panoptic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
         # we refer to the demo notebooks for visualization (see "Resources" section in the MaskFormer docs)
         predicted_panoptic_map = result["segmentation"]
-        plt.imsave("predicted_panoptic_map.png", predicted_panoptic_map, cmap="viridis")
-        output = PIL.Image.open("predicted_panoptic_map.png")
+        plt.plot("predicted_panoptic_map.png", predicted_panoptic_map, cmap="viridis")
+        fig = plt.gcf()
+        output = fig2img(fig)
+        # output = PIL.Image.open("predicted_panoptic_map.png")
         # output = PIL.Image.fromarray(predicted_panoptic_map.cpu().numpy().astype('uint8')).convert('RGB')
 
     # Calculate the prediction time
@@ -99,8 +111,8 @@ def predict(image, model_name: str = "detr",) -> Tuple[Dict, float]:
 ### 4. Gradio app ###
 
 # Create title, description and article strings
-title = "Segementation Demo"
-description = "An Mutimodel Segementation Demo"
+title = "Segmentation Demo"
+description = "An Mutimodel Segmentation Demo"
 article = ""
 
 # Create examples list from "examples/" directory