Make it accept png?

app.py

@@ -35,7 +35,7 @@ def get_base64(bin_file):
 # set_background('ipad_bg.png')
 
 upload_img = (st.file_uploader("输入character sheet", "png", accept_multiple_files=True))
-pose_format = "exr"
+pose_format = "png"
 upload_img2 = (st.file_uploader("pose images", pose_format, accept_multiple_files=True))
 
 # os.system("sh download.sh")

data_loader.py

@@ -225,8 +225,10 @@ class FileDataset(Dataset):
 
         if "labels" in sample:
             # return UDP as 4chn XYZV float tensor
+            if sample["labels"].dtype != np.float32:
+                sample["labels"] = sample["labels"].astype(np.float32) / 255
             sample["labels"] = torch.from_numpy(
-                sample["labels"].transpose((2, 0, 1)).astype(np.float32))
+                sample["labels"].transpose((2, 0, 1)))
             assert (sample["labels"].dtype == torch.float32)
 
         if "image_np" in sample: