Spaces:
Sleeping
Sleeping
File size: 2,266 Bytes
1455ab2 dae620a 23a61b4 dae620a 23a61b4 dae620a 5c05ddf dae620a 4905c12 5c05ddf dae620a 23a61b4 dae620a 5c05ddf dae620a e10e9d9 1455ab2 3a960b2 1455ab2 3a960b2 1455ab2 3a960b2 1455ab2 670cbf6 dae620a 23a61b4 dae620a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
import os
import torch
import torchvision.transforms as transforms
import gradio as gr
from huggingface_hub import hf_hub_download
from PIL import Image
print("Stage 0: Completed")
Fruits = ['Acerola', 'Apple', 'Apricot', 'Avocado', 'Banana', 'Black Berry', 'Blue Berry', 'Cantaloupe', 'Cherry',
'Coconut', 'Fig', 'Grapefruit', 'Grape', 'Guava', 'Kiwi Fruit', 'Lemon', 'Lime', 'Mango', 'Olive', 'Orange',
'Passion Fruit', 'Peach', 'Pear', 'Pineapple', 'Plum', 'Pomegranate', 'Raspberry', 'Strawberry', 'Tomato',
'Watermelon']
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
repo_name = "VinayHajare/fruits30-resnet18"
file_name = "fruit_resnet18(99.40%).pt"
model_path = hf_hub_download(repo_id = repo_name, filename = file_name)
model = torch.load(model_path, map_location=torch.device('cpu'))
model.to(device)
model.eval()
print("Stage 1: Completed ")
transform = transforms.Compose([
transforms.Resize(224),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
def predict_image(image):
image_tensor = transform(Image.fromarray(image)).unsqueeze(0)
image_tensor = image_tensor.to(device)
with torch.no_grad():
output = model(image_tensor)
predicted = torch.argmax(output).item()
return Fruits[predicted]
interface = gr.Interface(
fn = predict_image,
inputs = "image",
outputs = "text",
allow_flagging = "never",
examples = [
os.path.join(os.path.dirname(__file__), "0.jpg"),
os.path.join(os.path.dirname(__file__), "1.jpg"),
os.path.join(os.path.dirname(__file__), "3.jpg"),
os.path.join(os.path.dirname(__file__), "8.jpg"),
os.path.join(os.path.dirname(__file__), "10.jpg"),
os.path.join(os.path.dirname(__file__), "11.jpg"),
os.path.join(os.path.dirname(__file__), "12.jpg"),
os.path.join(os.path.dirname(__file__), "14.jpg"),
os.path.join(os.path.dirname(__file__), "21.jpg"),
os.path.join(os.path.dirname(__file__), "27.jpg"),
os.path.join(os.path.dirname(__file__), "37.jpg"),
],
theme = gr.themes.Soft()
)
print("Stage 2: Completed")
interface.launch(debug = True)
|