add vit-base-patch16-224

vit-base-patch16-224/Dockerfile

@@ -0,0 +1,13 @@
+FROM python:3.7.4-slim
+
+WORKDIR /app
+
+COPY requirements.txt /app
+
+RUN pip config set global.index-url https://pypi.mirrors.ustc.edu.cn/simple/ 
+
+RUN pip3 install --trusted-host pypi.python.org -r requirements.txt 
+
+COPY . /app
+
+CMD ["python", "vit.py"]

vit-base-patch16-224/requirements.txt

@@ -0,0 +1,4 @@
+gradio
+huggingface
+torch
+transformers

vit-base-patch16-224/vit.py

@@ -0,0 +1,24 @@
+#图像分类
+import gradio as gr
+from transformers import ViTFeatureExtractor, ViTForImageClassification
+
+feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')
+model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')
+
+def image_classification(image):
+    inputs = feature_extractor(images=image, return_tensors="pt")
+    outputs = model(**inputs)
+    logits = outputs.logits
+    predicted_class_idx = logits.argmax(-1).item()
+    return model.config.id2label[predicted_class_idx]
+
+demo = gr.Interface(fn=image_classification, 
+                    inputs=gr.Image(), 
+                    outputs=gr.Label(num_top_classes=1),
+                    title = "图像分类",
+                    allow_flagging="never",
+                    examples = ['cat.jpeg', 'dog.jpeg', 'zebra.jpeg'])
+
+
+if __name__ == "__main__":
+    demo.queue(concurrency_count=3).launch(server_name = "0.0.0.0", server_port = 7000, max_threads=40)