Tirath5504's picture
Create app.py
b069732 verified
raw
history blame
1.07 kB
import torch
from transformers import ViTForImageClassification, ViTImageProcessor
from PIL import Image
import gradio as gr
model = ViTForImageClassification.from_pretrained('vit-hateful-gesture-classification')
processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224')
class_names = ['cut_throat_gesture', 'finger_gun_to_the_head', 'middle_finger', 'slanted_eyes_gesture', 'swastika']
def predict(image):
inputs = processor(images=image, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs).logits
predicted_class_idx = outputs.argmax(-1).item()
predicted_class = class_names[predicted_class_idx]
return predicted_class
iface = gr.Interface(fn=predict,
inputs=gr.inputs.Image(type="pil"),
outputs=gr.outputs.Label(num_top_classes=1),
title="Hateful Gesture Detection",
description="Upload an image to classify hateful gestures or symbols")
if __name__ == "__main__":
iface.launch(share=True)