Update app.py

app.py

@@ -3,7 +3,7 @@ import gradio as gr
 import numpy as np 
 from PIL import Image
 import cv2
-from torchvision import models,transforms
+#from torchvision import models,transforms
 core = Core()
 
 # Read model to OpenVINO Runtime
@@ -14,14 +14,22 @@ tfms = transforms.Compose([
     transforms.ToTensor(),
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # imagenet
 ])
+def normalize(img):
+  img=img.astype(np.float32)
+  mean=(0.485, 0.456, 0.406)
+  std=(0.229, 0.224, 0.225)
+  img/=255.0
+  img-=mean
+  img/std
+  return img
 def segment_image(filepath):
   image=cv2.imread(filepath)
   image=cv2.cvtColor(image,cv2.COLOR_BGR2RGB)
   image = cv2.resize(image, (512,512))
-  x=tfms(image.copy()/255.)
-  #ort_input={ort_session.get_inputs()[0].name:x.cpu().unsqueeze(0).float().numpy()}
-  #out=ort_session.run(None,ort_input)
-  out = compiled_model_onnx([x.unsqueeze(0).float().cpu().numpy()])
+  x=normalize(image.copy())
+#ort_input={ort_session.get_inputs()[0].name:x.cpu().unsqueeze(0).float().numpy()}
+#out=ort_session.run(None,ort_input)
+  out = compiled_model_onnx(np.expand_dims(x,axis=0).transpose(0,3,1,2))
   pred_mask=np.squeeze(np.argmax(out[0],1)).astype(np.uint8)
   color_mask=cv2.applyColorMap(pred_mask,cv2.COLORMAP_MAGMA)*10
   masked_image=cv2.addWeighted(image,0.6,color_mask,0.4,0.1)