Food_info / app.py
aman012's picture
Update app.py
4e1bbf3 verified
import requests
import json
import os
import transformers
from transformers import ViTFeatureExtractor, ViTForImageClassification
import warnings
from PIL import Image
from io import BytesIO
warnings.filterwarnings('ignore')
api_key = os.environ.get('API_KEY')
import gradio as gr
model_name = "google/vit-base-patch16-224"
feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
model = ViTForImageClassification.from_pretrained(model_name)
def food_identification(img_path):
try:
image = Image.open(img_path)
input = feature_extractor(images= image,return_tensors= 'pt' )
outputs = model(**input)
logits = outputs.logits
# model predicts one of the 1000 ImageNet classes
pred_class_idx = logits.argmax(-1).item()
ans = f"{model.config.id2label[pred_class_idx]}".split(',')[0]
return ans
except FileNotFoundError as e:
print('Got an error -> ', e)
return None
def cal_calculation(fruit_name):
api_url = f'https://api.api-ninjas.com/v1/nutrition?query={fruit_name}'
Headers = { 'X-Api-Key': api_key}
response = requests.get(api_url, headers=Headers)
# we are sending api_key through header because header is more secure as compare to body or api_url .
if response.status_code == 200:
return response.json()
else:
return f"Error: {response.status_code}, Messages: {response.text}"
def json_to_html_table(json_data):
# Start the HTML table
html = '<table border="1" style="border-collapse: collapse; width: 50%;">'
# Add table rows
for item in json_data:
for key, value in item.items():
html += '<tr>'
html += f'<th style="padding: 10px; background-color: #f2f2f2; color: blue; text-align: left;">{key}</th>'
html += f'<td style="padding: 10px; text-align: left;">{value}</td>'
html += '</tr>'
# End the HTML table
html += '</table>'
return html
def display_table(json_input):
json_data = json.loads(json_input)
return json_to_html_table(json_data)
def main_function(image_path):
try:
food_identity = food_identification(image_path)
if not food_identity:
return "<p style='color:red;'>Could not identify the food item.</p>"
nutrition_info = cal_calculation(food_identity)
if isinstance(nutrition_info, str): # Error message from API
return f"<p style='color:red;'>{nutrition_info}</p>"
json_data = json.dumps(nutrition_info)
return display_table(json_data)
except Exception as e:
return f"<p style='color:red;'>Unexpected Error: {str(e)}</p>"
import gradio as gr
iface = gr.Interface(
fn= main_function,
inputs= gr.Image(type='filepath'),
outputs= 'html',
title= 'Food Identification and neutrition info',
description= 'Upload the image of a food to get neutrition info',
allow_flagging='never',
)
iface.launch(share= True)