Update app.py

app.py

@@ -1,4 +1,4 @@
-from openvino.runtime import Core
+#from openvino.runtime import Core
 import gradio as gr 
 import numpy as np 
 from PIL import Image
@@ -6,14 +6,101 @@ import cv2
 from torchvision import models,transforms
 from typing import Iterable
 import gradio as gr
+from torch import nn 
 from gradio.themes.base import Base
 from gradio.themes.utils import colors, fonts, sizes
 import time
-core = Core()
+import intel_extension_for_pytorch as ipex
 
+#core = Core()
+def conv(in_channels, out_channels):
+    return nn.Sequential(
+        nn.ReflectionPad2d(1),
+        nn.Conv2d(in_channels, out_channels, kernel_size=3),
+        nn.BatchNorm2d(out_channels),
+        nn.ReLU(inplace=True)
+    )
+class resconv(nn.Module):
+  def __init__(self,in_features,out_features):
+    super(resconv,self).__init__()
+    self.block=nn.Sequential(
+        nn.ReflectionPad2d(1),
+        nn.Conv2d(in_features,out_features,3),
+        nn.InstanceNorm2d(out_features),
+        nn.ReLU(inplace=True),
+        nn.ReflectionPad2d(1),
+        nn.Conv2d(out_features,out_features,3),
+        nn.InstanceNorm2d(out_features),
+        nn.ReLU(inplace=True),
+
+    )
+  def forward(self,x):
+      return x+self.block(x)
+
+def up_conv(in_channels, out_channels):
+    return nn.Sequential(
+        nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2),
+        nn.BatchNorm2d(out_channels),
+        nn.ReLU(inplace=True)
+    )
+
+class ResnUnet(nn.Module):
+  def __init__(self, out_channels=32,number_of_block=9):
+    super().__init__()
+    out_features=64
+    channels=3
+    model=[nn.ReflectionPad2d(3),nn.Conv2d(3,out_features,7),nn.InstanceNorm2d(out_features),
+    nn.ReLU(inplace=True),nn.MaxPool2d(3,stride=2)]
+    model+=[resconv(out_features,out_features)]
+    model+=[nn.Conv2d(out_features,out_features*2,3,stride=2,padding=1),nn.InstanceNorm2d(out_features),
+    nn.ReLU(inplace=True)]
+    model+=[resconv(out_features*2,out_features*2)]
+    model+=[nn.Conv2d(out_features*2,out_features*4,3,stride=2,padding=1),nn.InstanceNorm2d(out_features),
+    nn.ReLU(inplace=True)]
+    model+=[resconv(out_features*4,out_features*4)]
+    model+=[nn.Conv2d(out_features*4,out_features*8,3,stride=2,padding=1),nn.InstanceNorm2d(out_features),
+    nn.ReLU(inplace=True)]
+    model+=[resconv(out_features*8,out_features*8)]
+    out_features*=8
+    input_features=out_features
+    for _ in range(4):
+      out_features//=2
+      model+=[
+              nn.Upsample(scale_factor=2),
+              nn.Conv2d(input_features,out_features,3,stride=1,padding=1),
+              nn.InstanceNorm2d(out_features),
+              nn.ReLU(inplace=True)
+      ]
+      input_features=out_features
+    model+=[nn.ReflectionPad2d(3),nn.Conv2d(32,32,7),
+    ]
+    self.model=nn.Sequential(*model)
+  def forward(self,x):
+    return self.model(x)
+
+
+model=ResnUnet().to('cpu')
+
+# Load the state_dict
+state_dict = torch.load('/content/real_model1_onnx_compat.pt',map_location='cpu')
+
+# Create a new state_dict without the 'module.' prefix
+new_state_dict = {}
+for key, value in state_dict.items():
+    new_key = key.replace("module.", "")  # Remove the 'module.' prefix
+    new_state_dict[new_key] = value
+
+# Load the new state_dict into your model
+model.load_state_dict(new_state_dict)
+model.eval()
+
+model = ipex.optimize(model, weights_prepack=False)
+
+model = torch.compile(model, backend="ipex")    
+    
 # Read model to OpenVINO Runtime
-model_ir = core.read_model(model="Davinci_eye.onnx")
-compiled_model_ir = core.compile_model(model=model_ir, device_name='CPU')
+#model_ir = core.read_model(model="Davinci_eye.onnx")
+#compiled_model_ir = core.compile_model(model=model_ir, device_name='CPU')
 
 tfms = transforms.Compose([
     transforms.ToTensor(),
@@ -21,8 +108,8 @@ tfms = transforms.Compose([
 ])
 color_map={
     (251,244,5): 1,
-    (37,250,5):2, 
-    (0,21,209):3, 
+    (37,250,5):2,
+    (0,21,209):3,
     (172,21,2): 4,
     (172,21,229): 5,
     (6,254,249): 6,
@@ -52,7 +139,7 @@ color_map={
     (20,20,0):30,
     (20,20,0):31,
      }
-
+colormap={v:[i for i in k] for k,v in color_map.items()}
 items = {
     1: "HarmonicAce_Head",
     2: "HarmonicAce_Body",
@@ -127,27 +214,29 @@ class Davinci_Eye(Base):
 
 davincieye = Davinci_Eye()
 
-colormap={v:[i for i in k] for k,v in color_map.items()}
 
 def convert_mask_to_rgb(pred_mask):
+
     rgb_mask=np.zeros((pred_mask.shape[0],pred_mask.shape[1],3),dtype=np.uint8)
     for k,v in colormap.items():
         rgb_mask[pred_mask==k]=v
     return rgb_mask
 
 
-
 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())
+  image = cv2.resize(image, (224,224))
+  x=tfms(image.copy()/255.)
+  with torch.no_grad():
+    mask=model(x.unsqueeze(0).float())
 #ort_input={ort_session.get_inputs()[0].name:x.cpu().unsqueeze(0).float().numpy()}
 #out=ort_session.run(None,ort_input)
-  out = compiled_model_ir(x.unsqueeze(0).float().cpu().numpy())
-  pred_mask=np.squeeze(np.argmax(out[0],1)).astype(np.uint8)
+  _,pred_mask=torch.max(mask,dim=1)
+  pred_mask=pred_mask[0].numpy()
+  pred_mask=pred_mask.astype(np.uint8)
   color_mask=convert_mask_to_rgb(pred_mask)
-  masked_image=cv2.addWeighted(image,0.6,color_mask,0.4,0.1)
+  masked_image=cv2.addWeighted(image,0.3,color_mask,0.8,0.2)
   pred_keys=pred_mask[np.nonzero(pred_mask)]
   objects=[items[k] for k in pred_keys]
   surgery_items=np.unique(np.array(objects),axis=0)
@@ -156,6 +245,7 @@ def segment_image(filepath):
         surg+=item+","+" "
   return Image.fromarray(masked_image),surg
 
+
 demo=gr.Interface(fn=segment_image,inputs=gr.Image(type='filepath'),
                   outputs=[gr.Image(type="pil"),gr.Text()],
                   examples=["R001_ch1_video_03_00-29-13-03.jpg",