Upload 2 files

app.py

@@ -221,7 +221,7 @@ def get_chest_obs(idx=None):
     idx, obs = get_obs_item(dataset_id, idx)
     x = get_fig_arr(postprocess(obs["x"].clone()))
     s = SEX_CAT_CHEST[int(obs["sex"].clone().squeeze().numpy())]
-    f = FIND_CAT[int(obs["finding"].clone().squeeze().numpy())]
+    f = FIND_CAT[obs["finding"].clone().squeeze().numpy().argmax(-1)]
     r = RACE_CAT[obs["race"].clone().squeeze().numpy().argmax(-1)]
     a = (obs["age"].clone().squeeze().numpy() + 1) * 50
     return (idx, x, r, s, f, float(np.round(a, 1)))

vae.py

@@ -6,7 +6,7 @@ import torch.distributions as dist
 import torch.nn.functional as F
 from torch import Tensor, nn
 
-from hps import Hparams
+# from hps import Hparams
 
 EPS = -9  # minimum logscale
 
@@ -85,7 +85,7 @@ class Block(nn.Module):
 
 
 class Encoder(nn.Module):
-    def __init__(self, args: Hparams):
+    def __init__(self, args):
         super().__init__()
         # parse architecture
         stages = []
@@ -135,7 +135,7 @@ class Encoder(nn.Module):
 
 
 class DecoderBlock(nn.Module):
-    def __init__(self, args: Hparams, in_width: int, out_width: int, resolution: int):
+    def __init__(self, args, in_width, out_width, resolution):
         super().__init__()
         bottleneck = int(in_width / args.bottleneck)
         self.res = resolution
@@ -207,7 +207,7 @@ class DecoderBlock(nn.Module):
 
 
 class Decoder(nn.Module):
-    def __init__(self, args: Hparams):
+    def __init__(self, args):
         super().__init__()
         # parse architecture
         stages = []
@@ -335,7 +335,7 @@ class Decoder(nn.Module):
 
 
 class DGaussNet(nn.Module):
-    def __init__(self, args: Hparams):
+    def __init__(self, args):
         super(DGaussNet, self).__init__()
         self.x_loc = nn.Conv2d(
             args.widths[0], args.input_channels, kernel_size=1, stride=1
@@ -438,7 +438,7 @@ class DGaussNet(nn.Module):
 
 
 class HVAE(nn.Module):
-    def __init__(self, args: Hparams):
+    def __init__(self, args):
         super().__init__()
         args.vr = "light" if "ukbb" in args.hps else None  # hacky
         self.encoder = Encoder(args)