Fix #21.

app.py

@@ -251,7 +251,7 @@ class Solver(object):
             print(f"Error al cargar el checkpoint: {e}.")
             raise Exception(f"Error al cargar el checkpoint: {e}")
 
-    def transfer_style(self, source_image, reference_image, target_domain_index):
+    def transfer_style(self, source_image, reference_image): # Eliminado target_domain_index
         # Asegúrate de que los modelos estén en modo de evaluación
         self.G.eval()
         self.S.eval()
@@ -263,14 +263,18 @@ class Solver(object):
                 transforms.ToTensor(),
                 transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
             ])
+            # Convertir a PIL image antes de la transformación
+            source_image = Image.fromarray(source_image)
+            reference_image = Image.fromarray(reference_image)
+
             source_image = transform(source_image).unsqueeze(0).to(self.device)
             reference_image = transform(reference_image).unsqueeze(0).to(self.device)
 
             # Crear el tensor de dominio objetivo
-            target_domain = torch.tensor([target_domain_index]).to(self.device)
+            # target_domain = torch.tensor([target_domain_index]).to(self.device) # Eliminado
 
             # Codificar el estilo de la imagen de referencia
-            s_ref = self.S(reference_image, target_domain)
+            s_ref = self.S(reference_image, torch.tensor([0]).to(self.device)) # Simplificado
 
             # Generar la imagen con el estilo transferido
             generated_image = self.G(source_image, s_ref)
@@ -280,7 +284,7 @@ class Solver(object):
             return generated_image
 
 # Función principal para la inferencia
-def main(source_image, reference_image, target_domain_index, checkpoint_path, args):
+def main(source_image, reference_image, checkpoint_path, args): # Eliminado target_domain_index
     if source_image is None or reference_image is None:
         raise gr.Error("Por favor, proporciona ambas imágenes (fuente y referencia).")
 
@@ -290,37 +294,37 @@ def main(source_image, reference_image, target_domain_index, checkpoint_path, ar
     solver.load_checkpoint(checkpoint_path)
 
     # Realizar la transferencia de estilo
-    generated_image = solver.transfer_style(source_image, reference_image, target_domain_index)
+    generated_image = solver.transfer_style(source_image, reference_image) # Eliminado target_domain_index
     return generated_image
 
 def gradio_interface():
     # Definir los argumentos (ajustados para la inferencia)
     args = SimpleNamespace(
-        img_size=128,  
-        num_domains=3, 
-        latent_dim=16, 
+        img_size=128,  # Asegúrate de que esto coincida con el tamaño de imagen usado en el entrenamiento
+        num_domains=3, # 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,
     )
 
     # Ruta al checkpoint
-    checkpoint_path = "iter/20500_nets_ema.ckpt"  # Reemplaza con la ruta correcta a tu checkpoint
+    checkpoint_path = "iter/10500_nets_ema.ckpt"  # Reemplaza con la ruta correcta a tu checkpoint
 
     # Crear la interfaz de Gradio
     inputs = [
         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
+        # gr.Radio(choices=[0, 1, 2], label="Target Domain (0: BMW, 1: Corvette, 2: Mazda)", value=0), # Eliminado
     ]
     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)."
+    description = "Transfer the style of one car to another. Upload a source car image and a reference car image." # Modificado
 
     # Crear la interfaz de Gradio
     iface = gr.Interface(
-        fn=lambda source_image, reference_image, target_domain_index: main(source_image, reference_image, target_domain_index, checkpoint_path, args),
+        fn=lambda source_image, reference_image: main(source_image, reference_image, checkpoint_path, args), # Eliminado target_domain_index
         inputs=inputs,
         outputs=outputs,
         title=title,
@@ -330,4 +334,4 @@ def gradio_interface():
 
 if __name__ == '__main__':
     iface = gradio_interface()
-    iface.launch(share=True)
+    iface.launch(share=True)