Add real2sim translation of the sim2real prediction

app.py

@@ -21,15 +21,20 @@ def pred_pipeline(img, transforms):
         orig_shape = img.shape
         input = transforms(img)
         input = input.unsqueeze(0)
-        output = model(input)
-
-        out_img = make_grid(output,#.detach().cpu(),
+        output_real = sim2real(input)
+        output_syn = real2sim(output_real)
+        out_img_real = make_grid(output_real,
                            nrow=1, normalize=True)  
+        out_syn_real = make_grid(output_syn,
+                           nrow=1, normalize=True)          
+
+
+
         out_transform = Compose([
             T.Resize(orig_shape[:2]),
             T.ToPILImage()
         ])
-        return out_transform(out_img)
+        return out_transform(out_img_real),  out_transform(out_syn_real)
 
 
 
@@ -46,12 +51,17 @@ transform = Compose([
     ])
 
 
-model = GeneratorResNet.from_pretrained('huggan/sim2real_cyclegan', input_shape=(n_channels, image_size, image_size), 
+sim2real = GeneratorResNet.from_pretrained('Chris1/sim2real-512', input_shape=(n_channels, image_size, image_size), 
+                num_residual_blocks=9)
+real2sim = GeneratorResNet.from_pretrained('Chris1/real2sim-512', input_shape=(n_channels, image_size, image_size), 
                 num_residual_blocks=9)
 
 gr.Interface(lambda image: pred_pipeline(image, transform), 
     inputs=gr.inputs.Image( label='input synthetic image'), 
-    outputs=gr.outputs.Image( type="pil",label='style transfer to the real world'),#plot,
+    outputs=[
+            gr.outputs.Image( type="pil",label='GAN sim2real prediction: style transfer of the input to the real world '),
+            gr.outputs.Image( type="pil",label='GAN real2sim prediction: translation to synthetic of the above prediction')
+            ],#plot,
     title = "GTA5(simulated) to Cityscapes (real) translation",
     examples = [
                 [example] for example in glob.glob('./samples/*.png')
@@ -59,6 +69,3 @@ gr.Interface(lambda image: pred_pipeline(image, transform),
     .launch()
 
 
-
-#iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-#iface.launch()