Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,224 Bytes
c9c230c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
import gradio as gr
import spaces
import torch
from image_loader import load_image_from_url, load_image_from_file
from image_processor import process_image
import logging
# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
torch.set_float32_matmul_precision(["high", "highest"][0])
try:
birefnet = AutoModelForImageSegmentation.from_pretrained(
"ZhengPeng7/BiRefNet", trust_remote_code=True
)
birefnet.to("cuda")
logging.info("BiRefNet model loaded successfully.")
except Exception as e:
logging.error(f"Error loading BiRefNet model: {e}")
raise Exception(f"Error loading BiRefNet model: {e}")
def fn(image_input):
try:
if isinstance(image_input, str): # URL input
img = load_image_from_url(image_input)
else: # File upload
img = load_image_from_file(image_input)
img = img.convert("RGB")
origin = img.copy()
processed_image = process(img)
return (processed_image, origin)
except Exception as e:
logging.error(f"Error in fn function: {e}")
return None, None # Return None or a placeholder image
@spaces.GPU
def process(image):
try:
processed_image = process_image(image, birefnet)
return processed_image
except Exception as e:
logging.error(f"Error in process function: {e}")
raise gr.Error(f"Error processing image: {e}")
def process_file(file_path):
try:
name_path = file_path.rsplit(".", 1)[0] + ".png"
img = load_image_from_file(file_path)
img = img.convert("RGB")
transparent = process(img)
transparent.save(name_path)
logging.info(f"Processed image saved to: {name_path}")
return name_path
except Exception as e:
logging.error(f"Error in process_file function: {e}")
raise gr.Error(f"Error processing file: {e}")
slider1 = gr.ImageSlider(label="Processed Image", type="pil", format="png")
slider2 = gr.ImageSlider(label="Processed Image from URL", type="pil", format="png")
image_upload = gr.Image(label="Upload an image")
image_file_upload = gr.Image(label="Upload an image", type="filepath")
url_input = gr.Textbox(label="Paste an image URL")
output_file = gr.File(label="Output PNG File")
# Example images
try:
chameleon = load_image_from_file("butterfly.jpg")
except Exception as e:
logging.error(f"Error loading example image: {e}")
chameleon = None # Or a placeholder image
url_example = "https://hips.hearstapps.com/hmg-prod/images/gettyimages-1229892983-square.jpg"
tab1 = gr.Interface(fn, inputs=image_upload, outputs=slider1, examples=[chameleon], api_name="image")
tab2 = gr.Interface(fn, inputs=url_input, outputs=slider2, examples=[url_example], api_name="text")
tab3 = gr.Interface(process_file, inputs=image_file_upload, outputs=output_file, examples=["butterfly.jpg"], api_name="png")
demo = gr.TabbedInterface(
[tab1, tab2, tab3], ["Image Upload", "URL Input", "File Output"], title="Background Removal Tool"
)
if __name__ == "__main__":
demo.launch(show_error=True)
|