Fix #18 app.py

app.py

@@ -290,46 +290,41 @@ def main(args, checkpoint_path, source_image, reference_image, target_domain_ind
 
     return generated_image
 
-def gradio_interface(checkpoint_path="iter/20500_nets_ema.ckpt", img_size=128, num_domains=3): # Agregamos los valores por defecto
-    # Interfaz de Gradio
+def gradio_interface():
+    # Definir los argumentos (ajustados para la inferencia)
+    args = SimpleNamespace(
+        img_size=128,  
+        num_domains=3, 
+        latent_dim=16, 
+        style_dim=64,
+        num_workers=0,  # Establecer en 0 para evitar problemas en algunos entornos
+        seed=8365,
+    )
+
+    # Ruta al checkpoint
+    checkpoint_path = "iter/20500_nets_ema.ckpt"  # Reemplaza con la ruta correcta a tu checkpoint
+
+    # Crear la interfaz de Gradio
     inputs = [
-        gr.Image(label="Source Image", type="pil"), # Especificamos el tipo de imagen como PIL
-        gr.Image(label="Reference Image", type="pil"),
-        gr.Radio(choices=["BMW", "Corvette", "Mazda"], label="Target Domain", default="BMW")
+        gr.Image(label="Source Image (Car to change style)"),
+        gr.Image(label="Reference Image (Style to transfer)"),
+        gr.Radio(choices=[0, 1, 2], label="Target Domain (0: BMW, 1: Corvette, 2: Mazda)", value=0), # Cambiado a value
     ]
-    outputs = gr.Image(label="Generated Image")
-
-    def process_images(source_image, reference_image, target_domain):
-        # Mapear el dominio seleccionado a un índice
-        domain_index = {"BMW": 0, "Corvette": 1, "Mazda": 2}[target_domain]
-
-        # Definir los argumentos (ajustados para la inferencia)
-        args = SimpleNamespace(
-            img_size=img_size,  # Asegúrate de que esto coincida con el tamaño de imagen usado en el entrenamiento
-            num_domains=num_domains, #args.num_domains, # Cambiado a 3 para que coincida con el checkpoint del MappingNetwork
-            latent_dim=16, # Puedes ajustar esto si es necesario
-            style_dim=64,
-            num_workers=0,  # Establecer en 0 para evitar problemas en algunos entornos
-            seed=8365,
-        )
-        try:
-            # Llamar a la función principal para realizar la inferencia
-            generated_image = main(args, checkpoint_path, source_image, reference_image, domain_index)
-            return generated_image
-        except Exception as e:
-            print(f"Error during processing: {e}")
-            return None # Devolvemos None en caso de error
+    outputs = gr.Image(label="Generated Image (Car with transferred style)")
+
+    title = "AutoStyleGAN: Car Style Transfer"
+    description = "Transfer the style of one car to another. Upload a source car image and a reference car image. Select the target domain (car brand)."
 
+    # Crear la interfaz de Gradio
     iface = gr.Interface(
-        fn=process_images,
+        fn=lambda source_image, reference_image, target_domain_index: main(source_image, reference_image, target_domain_index, checkpoint_path, args),
         inputs=inputs,
         outputs=outputs,
-        title="AutoStyleGAN Demo",
-        description="Transfer the style of a reference car image to a source car image.  Select the target car domain.",
+        title=title,
+        description=description,
     )
     return iface
 
 if __name__ == '__main__':
-    # Lanzar la interfaz de Gradio
     iface = gradio_interface()
     iface.launch()