Update app/model.py

app/model.py

@@ -1,24 +1,24 @@
 from transformers import ViTImageProcessor, ViTForImageClassification
-from PIL import Image
 import torch
+from PIL import Image
 
 MODEL_NAME = "google/vit-base-patch16-224"
 
-# Load processor and model once at startup
+# Load once at startup
 processor = ViTImageProcessor.from_pretrained(MODEL_NAME)
 model = ViTForImageClassification.from_pretrained(MODEL_NAME)
 
+
 def analyze_image(image: Image.Image):
     inputs = processor(images=image, return_tensors="pt")
     with torch.no_grad():
         outputs = model(**inputs)
-
     logits = outputs.logits
     predicted_class_idx = logits.argmax(-1).item()
     label = model.config.id2label[predicted_class_idx]
-    score = torch.nn.functional.softmax(logits, dim=-1)[0, predicted_class_idx].item()
+    confidence = torch.nn.functional.softmax(logits, dim=-1)[0, predicted_class_idx].item()
 
     return {
         "label": label,
-        "confidence": round(score, 3)
+        "confidence": round(confidence, 4)
     }