Update app.py

app.py

@@ -1,14 +1,14 @@
-from threading import Thread
 import gradio as gr
 import torch
-from transformers import AutoModel, AutoProcessor
-from transformers import StoppingCriteria, TextIteratorStreamer, StoppingCriteriaList
 import numpy as np
 import torch.nn.functional as F
+import PIL
+from threading import Thread
+from transformers import AutoModel, AutoProcessor
+from transformers import StoppingCriteria, TextIteratorStreamer, StoppingCriteriaList
 from torchvision.transforms.functional import normalize
 from huggingface_hub import hf_hub_download
 from briarmbg import BriaRMBG
-import PIL
 from PIL import Image
 from typing import Tuple
 
@@ -106,31 +106,30 @@ def process(image):
     orig_image = Image.fromarray(image)
     w,h = orig_im_size = orig_image.size
     image = resize_image(orig_image)
-    return image.tobytes()
-    # im_np = np.array(image)
-    # im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2,0,1)
-    # im_tensor = torch.unsqueeze(im_tensor,0)
-    # im_tensor = torch.divide(im_tensor,255.0)
-    # im_tensor = normalize(im_tensor,[0.5,0.5,0.5],[1.0,1.0,1.0])
-    # if torch.cuda.is_available():
-    #     im_tensor=im_tensor.cuda()
-
-    # #inference
-    # result=net(im_tensor)
-    # # post process
-    # result = torch.squeeze(F.interpolate(result[0][0], size=(h,w), mode='bilinear') ,0)
-    # ma = torch.max(result)
-    # mi = torch.min(result)
-    # result = (result-mi)/(ma-mi)    
-    # # image to pil
-    # im_array = (result*255).cpu().data.numpy().astype(np.uint8)
-    # pil_im = Image.fromarray(np.squeeze(im_array))
-    # # paste the mask on the original image
-    # new_im = Image.new("RGBA", pil_im.size, (0,0,0,0))
-    # new_im.paste(orig_image, mask=pil_im)
-    # # new_orig_image = orig_image.convert('RGBA')
-
-    # return new_im
+    im_np = np.array(image)
+    im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2,0,1)
+    im_tensor = torch.unsqueeze(im_tensor,0)
+    im_tensor = torch.divide(im_tensor,255.0)
+    im_tensor = normalize(im_tensor,[0.5,0.5,0.5],[1.0,1.0,1.0])
+    if torch.cuda.is_available():
+        im_tensor=im_tensor.cuda()
+
+    #inference
+    result=net(im_tensor)
+    # post process
+    result = torch.squeeze(F.interpolate(result[0][0], size=(h,w), mode='bilinear') ,0)
+    ma = torch.max(result)
+    mi = torch.min(result)
+    result = (result-mi)/(ma-mi)    
+    # image to pil
+    im_array = (result*255).cpu().data.numpy().astype(np.uint8)
+    pil_im = Image.fromarray(np.squeeze(im_array))
+    # paste the mask on the original image
+    new_im = Image.new("RGBA", pil_im.size, (0,0,0,0))
+    new_im.paste(orig_image, mask=pil_im)
+    # new_orig_image = orig_image.convert('RGBA')
+
+    return new_im.tobytes()
 
 
 title = """<h1 style="text-align: center;">Product description generator</h1>"""
@@ -149,7 +148,7 @@ with gr.Blocks(css=css) as demo:
             chat = gr.Chatbot(show_label=False)
             submit = gr.Button(value="Upload", variant="primary")
         with gr.Column():
-            output = gr.Image(type="pil", interactive=False)
+            output = gr.Image(type="pil")
         
     response_handler = (
         response,