Added linear reuse and price layers

app.py

@@ -67,7 +67,53 @@ def predict_brand(img):
 
 
 def estimate_price_and_usage(img):
-    return "Estimated price: 50-100 SEK - Usage: Reuse - Saved C02: 4 kg"
+    query_features = model.encode_image(preprocess(img).unsqueeze(0).to(device))
+
+    # Estimate usage
+    num_classes = 2
+    probe = torch.nn.Linear(
+        query_features.shape[-1],
+        num_classes,
+        dtype=torch.float16,
+        bias=False
+    )
+    # Load weights for the linear layer as a tensor
+    linear_data = torch.load("files/reuse_linear.pth")
+    probe.weight.data = linear_data["weight"]
+
+    # Do inference
+    probe.eval()
+    probe = probe.to(device)
+    output = probe(query_features)
+    print(output)
+    output = torch.softmax(output, dim=-1)
+    output = output.cpu().detach().numpy().astype("float32")
+    reuse = output.argmax(axis=-1)[0]
+    reuse_classes = ["Reuse", "Export"]
+
+    # Estimate price
+    num_classes = 4
+    probe = torch.nn.Linear(
+        query_features.shape[-1],
+        num_classes,
+        dtype=torch.float16,
+        bias=False
+    )
+    # Print output shape for the linear layer
+    # Load weights for the linear layer as a tensor
+    linear_data = torch.load("files/price_linear.pth")
+    probe.weight.data = linear_data["weight"]
+
+    # Do inference
+    probe.eval()
+    probe = probe.to(device)
+    output = probe(query_features)
+    output = torch.softmax(output, dim=-1)
+    output = output.cpu().detach().numpy().astype("float32")
+    price = output.argmax(axis=-1)[0]
+    price_classes = ["<50", "50-100", "100-150", ">150"]
+
+    return f"Estimated price: {price_classes[price]} SEK - Usage: {reuse_classes[reuse]}"
 
 
 def retrieve(query):
@@ -141,7 +187,7 @@ with gr.Blocks(
                 predicted_brand = gr.Textbox(label="Brand", show_label=False)
 
         with gr.Column(variant="compact"):
-            btn_estimate = gr.Button("Estimate Price, Reuse, and Saved C02").style(size="sm")
+            btn_estimate = gr.Button("Estimate Price and Reuse").style(size="sm")
             text_box = gr.Textbox(label="Estimates:", show_label=False)
     with gr.Tab("Image Retrieval"):
         with gr.Row(variant="compact"):

files/price_linear.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7a00e3c8fa9f78af43ab1208e14b818bea5028443e4c5260c6743e05f14b378f
+size 5115

files/reuse_linear.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8998d9093128cab4269649e9dc70541940f0d8a9a92c6e02fe774a6153f5e29b
+size 3067