Spaces:

pytorch
/

MEAL-V2

Runtime error

App Files Files Community

MEAL-V2 / app.py

akhaliq HF staff

Update app.py

5c3e7e7 about 2 years ago

raw history blame contribute delete

No virus

2.7 kB

	import torch
	from PIL import Image
	from torchvision import transforms
	import gradio as gr
	import os
	# list of models: 'mealv1_resnest50', 'mealv2_resnest50', 'mealv2_resnest50_cutmix', 'mealv2_resnest50_380x380', 'mealv2_mobilenetv3_small_075', 'mealv2_mobilenetv3_small_100', 'mealv2_mobilenet_v3_large_100', 'mealv2_efficientnet_b0'
	# load pretrained models, using "mealv2_resnest50_cutmix" as an example
	model = torch.hub.load('AK391/MEAL-V2','meal_v2', 'mealv2_resnest50_cutmix', pretrained=True)
	model.eval()
	# Download an example image from the pytorch website
	torch.hub.download_url_to_file("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")

	# Download ImageNet labels
	os.system("wget https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt")

	def inference(input_image):
	preprocess = transforms.Compose([
	transforms.Resize(256),
	transforms.CenterCrop(224),
	transforms.ToTensor(),
	transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
	])
	input_tensor = preprocess(input_image)
	input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model

	# move the input and model to GPU for speed if available
	if torch.cuda.is_available():
	input_batch = input_batch.to('cuda')
	model.to('cuda')

	with torch.no_grad():
	output = model(input_batch)
	# The output has unnormalized scores. To get probabilities, you can run a softmax on it.
	probabilities = torch.nn.functional.softmax(output[0], dim=0)

	# Read the categories
	with open("imagenet_classes.txt", "r") as f:
	categories = [s.strip() for s in f.readlines()]
	# Show top categories per image
	top5_prob, top5_catid = torch.topk(probabilities, 5)
	result = {}
	for i in range(top5_prob.size(0)):
	result[categories[top5_catid[i]]] = top5_prob[i].item()
	return result

	inputs = gr.inputs.Image(type='pil')
	outputs = gr.outputs.Label(type="confidences",num_top_classes=5)

	title = "MEAL_V2"
	description = "Gradio demo for MEAL_V2, Boosting Tiny and Efficient Models using Knowledge Distillation. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."
	article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2009.08453'>MEAL V2: Boosting Vanilla ResNet-50 to 80%+ Top-1 Accuracy on ImageNet without Tricks</a> \| <a href='https://github.com/szq0214/MEAL-V2'>Github Repo</a></p>"

	examples = [
	['dog.jpg']
	]
	gr.Interface(inference, inputs, outputs, title=title, description=description, article=article, examples=examples, analytics_enabled=False).launch()