|
from transformers import AutoFeatureExtractor, ResNetForImageClassification |
|
import torch |
|
|
|
|
|
|
|
|
|
|
|
feature_extractor = AutoFeatureExtractor.from_pretrained("microsoft/resnet-50") |
|
model = ResNetForImageClassification.from_pretrained("microsoft/resnet-50") |
|
|
|
import gradio as gr |
|
def segment(image): |
|
inputs = feature_extractor(image, return_tensors="pt") |
|
|
|
with torch.no_grad(): |
|
logits = model(**inputs).logits |
|
probs = torch.nn.Softmax(dim=1)(logits) |
|
|
|
labels = {model.config.id2label[idx] : flaot(prob) for idx, prob in enumerate(probs[0])} |
|
print(labels) |
|
|
|
|
|
predicted_label = logits.argmax(-1).item() |
|
return labels |
|
|
|
gr.Interface(fn=segment, inputs="image", outputs="text").launch() |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|