Pipeline

app.py

@@ -1,24 +1,15 @@
-from transformers import EfficientNetImageProcessor, EfficientNetForImageClassification
+from transformers import pipeline
 from PIL import Image
-import torch
-import torch.nn.functional as F
 
-device = torch.device("cpu")
-
-
-processor = EfficientNetImageProcessor.from_pretrained("models/dennisjooo/Birds-Classifier-EfficientNetB2")
-model = EfficientNetForImageClassification.from_pretrained("models/dennisjooo/Birds-Classifier-EfficientNetB2").to(device)
+#device = torch.device("cpu")
 
+#processor = EfficientNetImageProcessor.from_pretrained("models/dennisjooo/Birds-Classifier-EfficientNetB2")
+#model = EfficientNetForImageClassification.from_pretrained("models/dennisjooo/Birds-Classifier-EfficientNetB2").to(device)
 
 def predict(image):
-    inputs = processor(images=image, return_tensors="pt").to(device)
-    outputs = model(**inputs)
-    logits = outputs.logits
-    predicted_class_prob = F.softmax(logits, dim=-1).detach().cpu().numpy().max()
-    predicted_class_idx = logits.argmax(-1).item()
-    label = model.config.id2label[predicted_class_idx].split(",")[0]
-    return {label: float(predicted_class_prob)}
-
+    pipe = pipeline("image-classification", model="models/dennisjooo/Birds-Classifier-EfficientNetB2")
+    predictions = pipe(img)
+    return {p["label"]: p["score"] for p in predictions}
 
 import gradio as gr