Gun_detection / app.py
NKASG's picture
Create app.py
fd27a87 verified
raw
history blame
1.03 kB
import gradio as gr
from transformers import AutoTokenizer, AutoModelForImageClassification
from PIL import Image
import requests
import torch
# Load model from Hugging Face model hub
model_name = "your-username/your-model-name" # Replace with your model's name on Hugging Face
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForImageClassification.from_pretrained(model_name)
# Define function for image preprocessing and prediction
def process_image(image):
# Load and preprocess image
image = Image.open(image)
inputs = tokenizer(image, return_tensors="pt", padding=True, truncation=True)
# Make prediction
outputs = model(**inputs)
predicted_class = torch.argmax(outputs.logits, dim=1)
return predicted_class.item()
# Create Gradio interface
inputs = gr.inputs.Image()
output = gr.outputs.Label(num_top_classes=1)
interface = gr.Interface(fn=process_image, inputs=inputs, outputs=output, capture_session=True)
# Deploy the Gradio interface
interface.launch(share=True)