Update app.py

app.py

@@ -25,7 +25,7 @@ import albumentations as A # Albumentations is a computer vision tool that boost
 from albumentations.pytorch.transforms import ToTensorV2
 
 model = models.resnet50(pretrained=False)
-model.fc = nn.Linear(2048, num_classes)
+model.fc = nn.Linear(2048, 21)
 model.load_state_dict(torch.load('resnet_best.pth', map_location=torch.device('cpu')), strict=True)
 
 st.title("some big ML function")
@@ -37,6 +37,7 @@ if uploaded_file is not None:
     img = Image.open(uploaded_file)
     st.image(img)
     img = np.array(img)
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
     
   elif ".csv" in uploaded_file.name:
     dataframe = pd.read_csv(uploaded_file)