init

requirements.txt

@@ -0,0 +1,7 @@
+torch
+torchvision
+numpy
+huggingface-hub
+opencv-python
+timm>=0.6.12
+transformers

view_angle_app.py

@@ -0,0 +1,25 @@
+import gradio as gr 
+from torchvision import transforms
+from transformers import ViTForImageClassification
+
+model_path = "Inf009/view-angle"
+model = ViTForImageClassification.from_pretrained(model_path)
+model.eval()
+val_transforms = transforms.Compose(
+        [
+            transforms.Resize((224, 224)),
+            transforms.CenterCrop(224),
+            transforms.ToTensor(),
+        ]
+    )
+
+def predict_view_angle(image):
+    image = val_transforms(image)
+    outputs = model(image.unsqueeze(0)).logits.squeeze(0).sigmoid().detach().numpy()
+    indices = sorted(range(len(outputs)), key=lambda x: outputs[x], reverse=True)
+    predict_tags = ["45度俯视", "俯视", "正视"]
+    return predict_tags[indices[0]]
+
+    
+app = gr.Interface(fn=predict_view_angle, inputs=gr.inputs.Image(type="pil"), outputs=gr.outputs.Textbox())
+app.launch()