Spaces:

teohteiktoe
/

image-classification-general

Sleeping

Update app.py

d77ad2a verified 2 months ago

1.35 kB

# -*- coding: utf-8 -*-
"""image classification general.ipynb
Automatically generated by Colab.
Original file is located at
    https://colab.research.google.com/drive/1IKVz7857OWmXGom1Z2zq884jI84gkSUx
"""

#!pip install gradio
#!pip install transformers torch

import gradio as gr
from transformers import ViTImageProcessor, AutoModelForImageClassification
from PIL import Image
import torch

# Load pre-trained model and feature extractor
model_name = "google/vit-base-patch16-224"
feature_extractor = ViTImageProcessor.from_pretrained(model_name)
model = AutoModelForImageClassification.from_pretrained(model_name)

# Define the prediction function
def classify_image(img):
    inputs = feature_extractor(images=img, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    logits = outputs.logits
    predicted_class_idx = logits.argmax(-1).item()
    predicted_label = model.config.id2label[predicted_class_idx]
    return predicted_label

# Build the Gradio interface
interface = gr.Interface(fn=classify_image,
                         inputs=gr.Image(type="pil"),
                         outputs="text",
                         title="Image Classification with ViT",
                         description="Upload an image and classify it using Vision Transformer (ViT)")

# Launch the app
interface.launch()

	# -- coding: utf-8 --
	"""image classification general.ipynb

	Automatically generated by Colab.

	Original file is located at
	https://colab.research.google.com/drive/1IKVz7857OWmXGom1Z2zq884jI84gkSUx
	"""

	#!pip install gradio
	#!pip install transformers torch

	import gradio as gr
	from transformers import ViTImageProcessor, AutoModelForImageClassification
	from PIL import Image
	import torch

	# Load pre-trained model and feature extractor
	model_name = "google/vit-base-patch16-224"
	feature_extractor = ViTImageProcessor.from_pretrained(model_name)
	model = AutoModelForImageClassification.from_pretrained(model_name)

	# Define the prediction function
	def classify_image(img):
	inputs = feature_extractor(images=img, return_tensors="pt")
	with torch.no_grad():
	outputs = model(**inputs)
	logits = outputs.logits
	predicted_class_idx = logits.argmax(-1).item()
	predicted_label = model.config.id2label[predicted_class_idx]
	return predicted_label

	# Build the Gradio interface
	interface = gr.Interface(fn=classify_image,
	inputs=gr.Image(type="pil"),
	outputs="text",
	title="Image Classification with ViT",
	description="Upload an image and classify it using Vision Transformer (ViT)")

	# Launch the app
	interface.launch()