Spaces:
Sleeping
Sleeping
| from transformers import ViTImageProcessor, ViTForImageClassification | |
| import torch | |
| import gradio as gr | |
| feature_extractor = ViTImageProcessor.from_pretrained("model_artifacts") | |
| model = ViTForImageClassification.from_pretrained("model_artifacts") | |
| labels = ['Chevrolet Equinox', | |
| 'Chevrolet Silverado 1500', | |
| 'Ford Escape', | |
| 'Ford Explorer', | |
| 'Ford F-150', | |
| 'GMC Sierra 1500', | |
| 'Honda CR-V', | |
| 'Jeep Compass', | |
| 'Jeep Grand Cherokee', | |
| 'Jeep Wrangler', | |
| 'Mazda CX-5', | |
| 'Nissan Rogue', | |
| 'RAM 1500', | |
| 'RAM 2500', | |
| 'Toyota Camry'] | |
| def classify(im): | |
| features = feature_extractor(im, return_tensors='pt') | |
| logits = model(features["pixel_values"])[-1] | |
| probability = torch.nn.functional.softmax(logits, dim=-1) | |
| probs = probability[0].detach().numpy() | |
| confidences = {label: float(probs[i]) for i, label in enumerate(labels)} | |
| return confidences | |
| description = """ | |
| Simple car recognition model. Can recognize one of the followings: | |
| Chevrolet Equinox | |
| Chevrolet Silverado 1500 | |
| Ford Escape | |
| Ford Explorer | |
| Ford F-150 | |
| GMC Sierra 1500 | |
| Honda CR-V | |
| Jeep Compass | |
| Jeep Grand Cherokee | |
| Jeep Wrangler | |
| Mazda CX-5 | |
| Nissan Rogue | |
| RAM 1500 | |
| RAM 2500 | |
| Toyota Camry | |
| """ | |
| interface = gr.Interface(fn=classify, | |
| inputs="image", | |
| outputs="label", | |
| title="Car classification demo :)", | |
| description=description ) | |
| interface.launch() | |