Spaces:

onnx
/

FCN

Runtime error

App Files Files Community

FCN / app.py

akhaliq HF staff

Update app.py

1266b7d about 2 years ago

raw history blame contribute delete

No virus

3.91 kB

	from PIL import Image
	import numpy as np
	import torch
	from torchvision import transforms, models
	from onnx import numpy_helper
	import os
	import onnxruntime as rt
	from matplotlib.colors import hsv_to_rgb
	import cv2
	import gradio as gr

	preprocess = transforms.Compose([
	transforms.ToTensor(),
	transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
	])

	# Start from ORT 1.10, ORT requires explicitly setting the providers parameter if you want to use execution providers
	# other than the default CPU provider (as opposed to the previous behavior of providers getting set/registered by default
	# based on the build flags) when instantiating InferenceSession.
	# For example, if NVIDIA GPU is available and ORT Python package is built with CUDA, then call API as following:
	# onnxruntime.InferenceSession(path/to/model, providers=['CUDAExecutionProvider'])
	os.system("wget https://github.com/AK391/models/raw/main/vision/object_detection_segmentation/fcn/model/fcn-resnet101-11.onnx")
	sess = rt.InferenceSession("fcn-resnet101-11.onnx")

	outputs = sess.get_outputs()


	classes = [line.rstrip('\n') for line in open('voc_classes.txt')]
	num_classes = len(classes)

	def get_palette():
	# prepare and return palette
	palette = [0] * num_classes * 3

	for hue in range(num_classes):
	if hue == 0: # Background color
	colors = (0, 0, 0)
	else:
	colors = hsv_to_rgb((hue / num_classes, 0.75, 0.75))

	for i in range(3):
	palette[hue * 3 + i] = int(colors[i] * 255)

	return palette

	def colorize(labels):
	# generate colorized image from output labels and color palette
	result_img = Image.fromarray(labels).convert('P', colors=num_classes)
	result_img.putpalette(get_palette())
	return np.array(result_img.convert('RGB'))

	def visualize_output(image, output):
	assert(image.shape[0] == output.shape[1] and \
	image.shape[1] == output.shape[2]) # Same height and width
	assert(output.shape[0] == num_classes)

	# get classification labels
	raw_labels = np.argmax(output, axis=0).astype(np.uint8)

	# comput confidence score
	confidence = float(np.max(output, axis=0).mean())

	# generate segmented image
	result_img = colorize(raw_labels)

	# generate blended image
	blended_img = cv2.addWeighted(image[:, :, ::-1], 0.5, result_img, 0.5, 0)

	result_img = Image.fromarray(result_img)
	blended_img = Image.fromarray(blended_img)

	return confidence, result_img, blended_img, raw_labels

	def inference(img):
	input_image = Image.open(img)
	orig_tensor = np.asarray(input_image)
	input_tensor = preprocess(input_image)
	input_tensor = input_tensor.unsqueeze(0)
	input_tensor = input_tensor.detach().cpu().numpy()
	output_names = list(map(lambda output: output.name, outputs))
	input_name = sess.get_inputs()[0].name
	detections = sess.run(output_names, {input_name: input_tensor})
	output, aux = detections
	conf, result_img, blended_img, _ = visualize_output(orig_tensor, output[0])
	return blended_img

	title="Fully Convolutional Network"
	description="FCNs are a model for real-time neural network for class-wise image segmentation. As the name implies, every weight layer in the network is convolutional. The final layer has the same height/width as the input image, making FCNs a useful tool for doing dense pixel-wise predictions without a significant amount of postprocessing. Being fully convolutional also provides great flexibility in the resolutions this model can handle. This specific model detects 20 different classes. The models have been pre-trained on the COCO train2017 dataset on this class subset."
	examples=[["examplefcn.png"]]
	gr.Interface(inference,gr.inputs.Image(type="filepath"),gr.outputs.Image(type="pil"),title=title,description=description,examples=examples).launch(enable_queue=True)