app.py

import pickle
import pandas as p
import numpy as np

import requests
import io
import os
import cv2
import gdown
import tempfile

from PIL import Image, ImageDraw, ImageFont
import PIL

from transparent_background import Remover
import torch
import torch.nn.functional as F

import time

import gradio as gr
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):
        """
        model_bytes: model weights as bytes (downloaded from "https://drive.google.com/file/d/13oBl5MTVcWER3YU4fSxW3ATlVfueFQPY/view?usp=share_link")
        device     : (default cuda:0 if available) specifying device for computation
        """
        
        self.model_path = None
        with tempfile.NamedTemporaryFile(suffix=".pth", delete=False) as tmp_file:
            tmp_file.write(model_bytes)
            self.model_path = tmp_file.name
            
            
        # get the path of the script that defines this class
        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")
        
        self.font_path = font_path
            
        super().__init__(fast=False, jit=False, device=device, ckpt=self.model_path)
    
    def __del__(self):
        if self.model_path is not None and os.path.exists(self.model_path):
            os.remove(self.model_path)
    
    def download(self):
        pass
        
    def predict(self, image, comparison=False, extra=""):
        
        s = time.time()
        prediction = self.raw_predict(image)
        e = time.time()
        #print(f"predict time {e-s:.4f}")
        
        if not comparison:
            return prediction
        else:
            return self.compare(image, prediction, e-s, extra)
        
    def raw_predict(self, image, empty_cache_after_prediction=False):
        
        t1 = time.time()
        out = self.process(image)
        t2 = time.time()
        
        prediction = Image.fromarray(out)
    
        # Crea una nueva imagen RGB con un fondo blanco del mismo tamaño que la original
        new_image = Image.new("RGB", prediction.size, (255, 255, 255))

        # Combina las dos imágenes, reemplazando los píxeles transparentes con blanco
        new_image.paste(prediction, mask=prediction.split()[3])
        
        t3 = time.time()
        
        if empty_cache_after_prediction and "cuda" in self.device:
            torch.cuda.empty_cache()
            
        t4 = time.time()
        
        #print(f"{(t2-t1)*1000:.4f} {(t3-t2)*1000:.4f} {(t4-t3)*1000:.4f}")
        
        return new_image
    
    def compare(self, image1, image2, prediction_time, extra_info=""):
        extra = 80

        concatenated_image = Image.new('RGB', (image1.width + image2.width, image1.height + extra), (255, 255, 255))
        concatenated_image.paste(image1, (0, 0+extra))
        concatenated_image.paste(image2, (image1.width, 0+extra))

        draw = ImageDraw.Draw(concatenated_image)
        
        font = ImageFont.truetype(self.font_path, 20) 
        draw.text((20, 0), f"size:{image1.size}\nmodel time:{prediction_time:.2f}s\n{extra_info}", fill=(0, 0, 0), font=font)

        return concatenated_image
    
    def read_image_from_url(self, url):
        response = requests.get(url)
        image = Image.open(io.BytesIO(response.content)).convert("RGB")
        
        return image
    
    def read_image_from_file(self, file_name):

        image = Image.open(file_name).convert("RGB")

        return image
    
    def read_image_form_bytes(self, image_bytes):

        # Convertir los bytes en imagen
        image = Image.open(io.BytesIO(image_bytes))
        return image
    
    def image_to_bytes(self, image, format="JPEG"):
        image_bytes = io.BytesIO()
        image_rgb = image.convert('RGB')
        image_rgb.save(image_bytes, format=format)
        image_bytes = image_bytes.getvalue()
        
        return image_bytes
    
    @classmethod
    def create_instance_from_model_url(cls, url):
        model_bytes = BackgroundRemover.download_model_from_url(url)
        
        return cls(model_bytes)
    
    @classmethod
    def create_instance_from_model_file(cls, file_path, device=None):
        with open(file_path, 'rb') as f:
            model_bytes = f.read()
        
        return cls(model_bytes, device)
    
    @classmethod
    def download_model_from_url(cls, url):
        with io.BytesIO() as file:
            gdown.download(url, file, quiet=False, fuzzy=True)

            # Get the contents of the file as bytes
            file.seek(0)
            model_bytes = file.read()
        
        return model_bytes
    
def show_image(url: str):
    response = requests.get(url)
    img = Image.open(BytesIO(response.content))
    return img

def do_predictions(url):
    response = requests.get(url)
    img = Image.open(BytesIO(response.content))
        
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    transform_model = BackgroundRemover.create_instance_from_model_file("model_weights.pth")

    # Set up data transformations
    data_transforms = {
        'train': transforms.Compose([
            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((384, 384)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
    }
    
    out = transform_model.predict(img, comparison=False)
    
    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()