Spaces:
Running
Running
update indentation
Browse files
README.md
CHANGED
@@ -55,15 +55,15 @@ Here’s we define our image classification model prediction function in PyTorch
|
|
55 |
<pre>
|
56 |
<code>def predict(inp):
|
57 |
|
58 |
-
|
|
|
|
|
|
|
|
|
59 |
|
60 |
-
|
61 |
|
62 |
-
|
63 |
-
|
64 |
-
prediction = torch.nn.functional.softmax(model(inp)[0], dim=0)
|
65 |
-
|
66 |
-
return {labels[i]: float(prediction[i]) for i in range(1000)}
|
67 |
</code>
|
68 |
</pre>
|
69 |
</p>
|
|
|
55 |
<pre>
|
56 |
<code>def predict(inp):
|
57 |
|
58 |
+
inp = Image.fromarray(inp.astype('uint8'), 'RGB')
|
59 |
+
|
60 |
+
inp = transforms.ToTensor()(inp).unsqueeze(0)
|
61 |
+
|
62 |
+
with torch.no_grad():
|
63 |
|
64 |
+
prediction = torch.nn.functional.softmax(model(inp)[0], dim=0)
|
65 |
|
66 |
+
return {labels[i]: float(prediction[i]) for i in range(1000)}
|
|
|
|
|
|
|
|
|
67 |
</code>
|
68 |
</pre>
|
69 |
</p>
|