Update app.py

app.py

@@ -3,17 +3,17 @@ 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
 model_onnx = core.read_model(model="Davinci_eye.onnx")
 compiled_model_onnx = core.compile_model(model=model_onnx, device_name='CPU')
 
-#tfms = transforms.Compose([
- #   transforms.ToTensor(),
-  #  transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # imagenet
-#])
+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)
@@ -26,10 +26,10 @@ def segment_image(filepath):
   image=cv2.imread(filepath)
   image=cv2.cvtColor(image,cv2.COLOR_BGR2RGB)
   image = cv2.resize(image, (512,512))
-  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))
+  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()])
   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)