Fixing model outputs

app.py

@@ -15,19 +15,15 @@ labels = [x for x in response.text.split("\n") if x != ""]
 model = pull_from_hub(repo_id="Dimitre/mobilenet_v3_small")
 
 def preprocess(image):
-    print(image.shape)
     image = image.reshape((-1, 224, 224, 3)) # (batch_size, height, width, num_channels)
-    print(image.shape)
     return image / 255.
 
 def postprocess(prediction):
-    # return {labels[i]: prediction[i] for i in range(len(labels))}
-    return {labels[i]: 0 for i in range(len(labels))}
+    return {labels[i]: prediction[i] for i in range(len(labels))}
 
 def predict_fn(image):
     image = preprocess(image)
-    prediction = model(image)
-    print(prediction)
+    prediction = model(image)[0].numpy()
     scores = postprocess(prediction)
     return scores