Davinci_EYE / app.py
alibidaran's picture
Update app.py
2652021
raw
history blame
No virus
2.02 kB
from openvino.runtime import Core
import gradio as gr
import numpy as np
from PIL import Image
import cv2
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
])
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()])
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)
return Image.fromarray(masked_image),Image.fromarray((color_mask))
demo=gr.Interface(fn=segment_image,inputs=gr.Image(type='filepath'),
outputs=[gr.Image(type="pil"),gr.Image(type="pil")],
examples=["R001_ch1_video_03_00-29-13-03.jpg",
"R002_ch1_video_01_01-07-25-19.jpg",
"R003_ch1_video_05_00-22-42-23.jpg",
"R004_ch1_video_01_01-12-22-00.jpg",
"R005_ch1_video_03_00-19-10-11.jpg",
"R006_ch1_video_01_00-45-02-10.jpg",
"R013_ch1_video_03_00-40-17-11.jpg"],
themes=gr.themes.Glass(primary_hue=gr.themes.colors.blue,secondary_hue=gr.themes.colors.blue),
title="Davinci Eye(CPU version)")
demo.launch()