Davinci_EYE / app.py
alibidaran's picture
Update app.py
96c54a2
raw
history blame
No virus
1.67 kB
from openvino.runtime import Core
import gradio as gr
import numpy as np
from PIL import Image
from torchvision import models,transforms
tfms = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # imagenet
])
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.Soft(primary_hue=gr.themes.colors.amber,secondary_hue=gr.themes.colors.blue),
title="Davinci Eye(CPU version)")
demo.launch()