Update app.py

app.py

@@ -25,7 +25,22 @@ model.fc = nn.Linear(2048, 21)
 model.load_state_dict(torch.load('resnet_best.pth', map_location=torch.device('cpu')), strict=True)
 model.eval()
 
-
-#if uploaded_file is not None:
-#    if '.jpg' in uploaded_file.name.lower() or '.png' in uploaded_file.name.lower():
- #
+if uploaded_file is not None:
+    if '.jpg' in uploaded_file.name.lower() or '.png' in uploaded_file.name.lower():
+        st.write(uploaded_file.name)
+        img = Image.open(uploaded_file)
+        st.image(img)
+        img = np.array(img)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # OpenCV images have a different color profile. So remember to switch to RGB, that our resnet model understands.
+
+        cust_transform = A.Compose([A.Resize(height=256, width=256, p=1.0),ToTensorV2(p=1.0)], p=1.0)
+        tensor = cust_transform(image=img)
+        tensor = tensor['image'].float().resize(1,3,256,256)
+
+        custom_pred = model.forward(tensor).detach().numpy() # Forward is the python method defined inside the resnet.
+        custom_pred
+
+        st.write(f'Predicted: {id2class[np.argmax(custom_pred)]}')
+    elif '.csv' in uploaded_file.name:
+        dataframe = pd.read_csv(uploaded_file)
+        st.write(dataframe)