File size: 1,773 Bytes
70a751e 9da5b91 ccb8223 70a751e 9da5b91 70a751e 9da5b91 70a751e ccb8223 9da5b91 ccb8223 70a751e ccb8223 9da5b91 ccb8223 9da5b91 ccb8223 70a751e 9da5b91 |
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 |
from transformers import ViTFeatureExtractor, ViTForImageClassification
from hugsvision.inference.VisionClassifierInference import VisionClassifierInference
import gradio as gr
import cv2
import numpy as np
# Load the pretrained ViT model and feature extractor
path = "mrm8488/vit-base-patch16-224_finetuned-kvasirv2-colonoscopy"
feature_extractor = ViTFeatureExtractor.from_pretrained(path)
model = ViTForImageClassification.from_pretrained(path)
# Create a VisionClassifierInference instance
classifier = VisionClassifierInference(
feature_extractor=feature_extractor,
model=model,
)
# Define a function to classify and overlay the label on the image
def classify_image_with_overlay(img):
# Predict the label
label = classifier.predict(img_path=img)
# Load the image using OpenCV
image = cv2.imread(img)
# Add a white rectangle for the label
font = cv2.FONT_HERSHEY_SIMPLEX
org = (10, 30)
font_scale = 1
color = (255, 255, 255) # White color
thickness = 2
text_size = cv2.getTextSize(label, font, font_scale, thickness)[0]
cv2.rectangle(image, (org[0] - 10, org[1] - text_size[1] - 10), (org[0] + text_size[0], org[1]), color, cv2.FILLED)
# Put the label text on the white rectangle
cv2.putText(image, label, org, font, font_scale, (0, 0, 0), thickness, cv2.LINE_AA)
# Convert the image to RGB format for Gradio
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
return image_rgb
iface = gr.Interface(
fn=classify_image_with_overlay,
inputs=gr.inputs.Image(),
outputs=gr.outputs.Image(),
live=True,
title="ViT Image Classifier with Overlay",
description="Upload an image for classification with label overlay.",
)
if __name__ == "__main__":
iface.launch()
|