add segment model

app.py

@@ -3,7 +3,7 @@ import numpy as np
 import torch
 from torchvision.transforms import Compose
 import cv2
-from dpt.models import DPTDepthModel
+from dpt.models import DPTDepthModel, DPTSegmentationModel
 from dpt.transforms import Resize, NormalizeImage, PrepareForNet
 import os
 
@@ -11,16 +11,31 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print("device: %s" % device)
 default_models = {
     "dpt_hybrid": "weights/dpt_hybrid-midas-501f0c75.pt",
+    "segment_hybrid": "weights/dpt_hybrid-ade20k-53898607.pt"
     }
 torch.backends.cudnn.enabled = True
 torch.backends.cudnn.benchmark = True
-net_w = net_h = 384
-model = DPTDepthModel(
+
+depth_model = DPTDepthModel(
             path=default_models["dpt_hybrid"],
             backbone="vitb_rn50_384",
             non_negative=True,
             enable_attention_hooks=False,
     )
+
+depth_model.eval()
+depth_model.to(device)
+
+seg_model = DPTSegmentationModel(
+            150,
+            path=default_models["segment_hybrid"],
+            backbone="vitb_rn50_384",
+        )
+seg_model.eval()
+seg_model.to(device)
+
+# Transform
+net_w = net_h = 384
 normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
 transform = Compose(
         [
@@ -38,8 +53,6 @@ transform = Compose(
         ]
     )
 
-model.eval()
-model.to(device)
 
 def write_depth(depth, bits=1, absolute_depth=False):
     """Write depth map to pfm and png file.
@@ -67,7 +80,8 @@ def write_depth(depth, bits=1, absolute_depth=False):
         return out.astype("uint8")
     elif bits == 2:
         return out.astype("uint16")
-    
+
+
 
 def DPT(image):
     img_input = transform({"image": image})["image"]
@@ -75,7 +89,7 @@ def DPT(image):
     with torch.no_grad():
         sample = torch.from_numpy(img_input).to(device).unsqueeze(0)
 
-        prediction = model.forward(sample)
+        prediction = depth_model.forward(sample)
         prediction = (
             torch.nn.functional.interpolate(
                 prediction.unsqueeze(1),
@@ -90,6 +104,26 @@ def DPT(image):
 
         depth_img = write_depth(prediction, bits=2)
         return depth_img
+    
+def Segment(image):
+    img_input = transform({"image": image})["image"]
+
+        # compute
+    with torch.no_grad():
+        sample = torch.from_numpy(img_input).to(device).unsqueeze(0)
+        # if optimize == True and device == torch.device("cuda"):
+        #     sample = sample.to(memory_format=torch.channels_last)
+        #     sample = sample.half()
+
+        out = seg_model.forward(sample)
+
+        prediction = torch.nn.functional.interpolate(
+            out, size=image.shape[:2], mode="bicubic", align_corners=False
+        )
+        prediction = torch.argmax(prediction, dim=1) + 1
+        prediction = prediction.squeeze().cpu().numpy()
+
+    return prediction
 
 title = " AISeed AI Application Demo "
 description = "# A Demo of Deep Learning for Depth Estimation"
@@ -99,16 +133,19 @@ with gr.Blocks() as demo:
     demo.title = title
     gr.Markdown(description)
     with gr.Row():
-        im = gr.Image(label="Input Image")
-        im_2 = gr.Image(label="Depth Image")
         with gr.Column():
-
+            
+            im_2 = gr.Image(label="Depth Image")
+            im_3 = gr.Image(label="Segment Image")
+        with gr.Column():
+            im = gr.Image(label="Input Image")
             btn1 = gr.Button(value="Depth Estimator")
             btn1.click(DPT, inputs=[im], outputs=[im_2])
+            btn2 = gr.Button(value="Segment")
+            btn2.click(Segment, inputs=[im], outputs=[im_3])
             gr.Examples(examples=example_list,
                 inputs=[im],
-                outputs=[im_2],
-                fn=DPT)
+                outputs=[im_2])
     
 if __name__ == "__main__":
     demo.launch()