year

app.py

@@ -4,7 +4,7 @@ import torch.nn.functional as F
 import torchvision.transforms as transforms
 from PIL import Image
 
-from constants import MAKES_MODELS, PRICE_BIN_LABELS
+from constants import MAKES_MODELS, PRICE_BIN_LABELS, YEARS
 
 model = torch.jit.load("mobilenetv2_432000_calib.pt")
 model.eval()
@@ -26,13 +26,24 @@ def classify(img: Image.Image):
     price_bins = dict(
         zip(PRICE_BIN_LABELS, F.softmax(outputs["price_bin"], dim=-1)[0].tolist())
     )
+    years = dict(
+        zip(
+            [str(year) for year in YEARS] + ["Unknown"],
+            F.softmax(outputs["year"], dim=-1)[0].tolist(),
+        )
+    )
     make_models = dict(
         zip(
             ([f"{make} {model}" for make, model in MAKES_MODELS] + ["Unknown"]),
             F.softmax(outputs["make_model"], dim=-1)[0].tolist(),
         )
     )
-    return f"${int(round(outputs['price_median'].item()))}", price_bins, make_models
+    return (
+        f"${int(round(outputs['price_median'].item()))}",
+        price_bins,
+        years,
+        make_models,
+    )
 
 
 iface = gr.Interface(
@@ -40,8 +51,9 @@ iface = gr.Interface(
     inputs=gr.Image(shape=(224, 224), type="pil"),
     outputs=[
         gr.Text(label="Price Prediction"),
-        gr.Label(label="Price Bin"),
-        gr.Label(label="Make/Model"),
+        gr.Label(label="Price Bin", num_top_classes=5),
+        gr.Label(label="Year", num_top_classes=5),
+        gr.Label(label="Make/Model", num_top_classes=10),
     ],
 )
 iface.launch()