File size: 2,019 Bytes
70a751e 9da5b91 ccb8223 639e21c 70a751e 9da5b91 70a751e 9da5b91 70a751e 639e21c ccb8223 639e21c 17b5eb0 fef230d 639e21c fef230d 17b5eb0 fef230d 17b5eb0 fef230d 00f8749 639e21c 70a751e ccb8223 9da5b91 00f8749 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 55 56 57 58 |
from transformers import ViTFeatureExtractor, ViTForImageClassification
from hugsvision.inference.VisionClassifierInference import VisionClassifierInference
import gradio as gr
import cv2
import numpy as np
from PIL import Image
# 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):
# Convert the numpy array image to a PIL image
img_pil = Image.fromarray(img)
# Predict the label and probability
label, probability = classifier.predict_image(img=img_pil)
# Load the image using OpenCV
image = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)
# 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(f"{label}: {probability:.2f}", 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, f"{label}: {probability:.2f}", org, font, font_scale, (0, 0, 0), thickness, cv2.LINE_AA)
# Resize the image to a larger size
image = cv2.resize(image, (800, 800)) # Defina o tamanho desejado aqui
return image
iface = gr.Interface(
fn=classify_image_with_overlay,
inputs=gr.inputs.Image(),
outputs=gr.outputs.Image(type="numpy"),
live=True,
title="ViT Image Classifier with Overlay",
description="Upload an image for classification with label overlay.",
)
if __name__ == "__main__":
iface.launch()
|