Update app.py

app.py

@@ -1,5 +1,5 @@
 import pickle
-import pandas as pd
+import pandas as p
 import numpy as np
 
 import requests
@@ -14,6 +14,7 @@ import PIL
 
 from transparent_background import Remover
 import torch
+import torch.nn.functional as F
 
 import time
 
@@ -22,6 +23,9 @@ from PIL import Image
 import requests
 from io import BytesIO
 
+from torchvision import datasets, models, transforms
+
+
 class BackgroundRemover(Remover):
     def __init__(self, model_bytes, device=None):
         """
@@ -36,7 +40,7 @@ class BackgroundRemover(Remover):
             
             
         # get the path of the script that defines this class
-        script_path = os.path.abspath(__file__)
+        script_path = "" #os.path.abspath(__file__)
 
         # construct the path to the arial.ttf file relative to the script location
         font_path = os.path.join(os.path.dirname(script_path), "arial.ttf")
@@ -152,9 +156,7 @@ class BackgroundRemover(Remover):
             model_bytes = file.read()
         
         return model_bytes
-
-
-
+    
 def show_image(url: str):
     response = requests.get(url)
     img = Image.open(BytesIO(response.content))
@@ -171,50 +173,25 @@ def do_predictions(url):
     # Set up data transformations
     data_transforms = {
         'train': transforms.Compose([
-            #transforms.Resize(512), #256
-            #transforms.CenterCrop(480), # 224
-            #transforms.Resize((256, 256)),
             transforms.Resize((384, 384)),
             transforms.ToTensor(),
             transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
         ]),
         'val': transforms.Compose([
-            #transforms.Resize(512), #256
-            #transforms.CenterCrop(480), # 224
-            #transforms.Resize((256, 256)),
-            transforms.Resize((384, 284)),
+            transforms.Resize((384, 384)),
             transforms.ToTensor(),
             transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
         ]),
     }
-
-    # Crear un modelo con la misma arquitectura
-    detect_model = models.resnet50(weights=None)  # Cambiar 'pretrained' por 'weights'
-    num_ftrs = detect_model.fc.in_features
-    num_classes = 2
-    detect_model.fc = nn.Linear(num_ftrs, num_classes)
-    detect_model = detect_model.to(device)
-
-    # Cargar los pesos guardados
-    model_weights_path = 'white_background_detection/resnet50_finetuned_weights.pth'
-    detect_model.load_state_dict(torch.load(model_weights_path))
-
-    # Cambiar el modelo a modo de evaluación
-    detect_model.eval()
-
-    print("")
     
-    prediction, predicted_probability, inference_time = predict_single_image_detection(img, detect_model, data_transforms['val'], "cuda:0")
+    out = transform_model.predict(img, comparison=False)
     
-    if prediction=="real":
-        out = transform_model.predict(img, comparison=False)
-        return prediction, predicted_probability, img, out, 
-    else:
-        return prediction, predicted_probability, img, None
-
-iface = gr.Interface(fn=do_predictions, inputs="text", outputs=["text", "text", "image", "image"], examples=[["https://http2.mlstatic.com/D_NQ_NP_2X_823376-MLU29226703936_012019-F.webp"],
-                                                                                                             ["https://http2.mlstatic.com/D_781350-MLA53584851929_022023-F.jpg"]])
-#iface.outputs[0].set_title("Predicción")
-#iface.outputs[1].set_title("Clase")
-#iface.outputs[2].set_title("Probabilidad")
+    return img, out
+
+iface = gr.Interface(fn=do_predictions, inputs="text",
+                     examples=[["https://http2.mlstatic.com/D_NQ_NP_2X_823376-MLU29226703936_012019-F.webp"],
+                               ["https://http2.mlstatic.com/D_781350-MLA53584851929_022023-F.jpg"]],
+                     outputs=["image", "image"],
+                    )
+
 iface.launch(share=True)