Upload 6 files

.gitattributes

@@ -1,35 +1,33 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+robot.png filter=lfs diff=lfs merge=lfs -text
+ship.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,12 +1,16 @@
----
-title: Testing
-emoji: 🏆
-colorFrom: green
-colorTo: green
-sdk: gradio
-sdk_version: 5.20.1
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+title: DIS Background Removal
+emoji: 🔥 🌠 🏰
+colorFrom: yellow
+colorTo: blue
+sdk: gradio
+sdk_version: 5.23.1
+python_version: 3.11.10
+app_file: app.py
+pinned: false
+license: apache-2.0
+models:
+- doevent/dis
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,201 +1,153 @@
-from fastapi import FastAPI, File, UploadFile, HTTPException, Form
-from fastapi.responses import JSONResponse
-from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
-import logging
-import torch
-from diffusers import StableDiffusionImg2ImgPipeline, EulerDiscreteScheduler
-from PIL import Image
-import numpy as np
-import cv2
-import mediapipe as mp
-import base64
-from io import BytesIO
-import uvicorn
-
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-app = FastAPI()
-
-# Add CORS middleware to allow requests from Framer
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],  # In production, restrict to your Framer domain
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-pipe = None
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-def initialize_pipeline():
-    global pipe
-    try:
-        logger.info("Starting pipeline initialization...")
-        model_id = "runwayml/stable-diffusion-v1-5"
-        scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
-        pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
-            model_id,
-            scheduler=scheduler,
-            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-            low_cpu_mem_usage=True
-        )
-        pipe = pipe.to(device)
-        logger.info(f"Stable Diffusion pipeline initialized successfully on {device}.")
-    except Exception as e:
-        logger.error(f"Failed to initialize pipeline: {str(e)}", exc_info=True)
-        raise
-
-try:
-    initialize_pipeline()
-except Exception:
-    logger.error("Application failed to start due to pipeline initialization error.")
-    raise
-
-mp_face_detection = mp.solutions.face_detection
-face_detection = mp_face_detection.FaceDetection(min_detection_confidence=0.5)
-
-def extract_face(image: Image.Image):
-    try:
-        logger.info("Starting face extraction.")
-        img_np = np.array(image)
-        img_rgb = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
-        results = face_detection.process(img_rgb)
-        if not results.detections:
-            logger.warning("No face detected in the reference image.")
-            raise HTTPException(status_code=400, detail="No face detected in the reference image.")
-
-        detection = results.detections[0]
-        bbox = detection.location_data.relative_bounding_box
-        ih, iw, _ = img_np.shape
-        x, y, w, h = int(bbox.xmin * iw), int(bbox.ymin * ih), int(bbox.width * iw), int(bbox.height * ih)
-        x, y = max(0, x), max(0, y)
-        w, h = min(w, iw - x), min(h, ih - y)
-        face = img_np[y:y+h, x:x+w]
-        logger.info("Face extracted successfully.")
-        return Image.fromarray(face), (x, y, w, h), image.size
-    except Exception as e:
-        logger.error(f"Error in extract_face: {str(e)}", exc_info=True)
-        raise HTTPException(status_code=500, detail=f"Face extraction failed: {str(e)}")
-
-def scale_coords(coords, original_size, new_size=(512, 512)):
-    orig_w, orig_h = original_size
-    new_w, new_h = new_size
-    x, y, w, h = coords
-    scale_x, scale_y = new_w / orig_w, new_h / orig_h
-    return (int(x * scale_x), int(y * scale_y), int(w * scale_x), int(h * scale_y))
-
-def overlay_face(generated_img, face_img, face_coords):
-    try:
-        logger.info("Starting face overlay.")
-        gen_np = np.array(generated_img)
-        x, y, w, h = face_coords
-        face_np = np.array(face_img.resize((w, h)))
-        gen_bgr = cv2.cvtColor(gen_np, cv2.COLOR_RGB2BGR)
-        mask = np.zeros((h, w), dtype=np.uint8)
-        cv2.rectangle(mask, (0, 0), (w, h), 255, -1)
-        center = (x + w // 2, y + h // 2)
-        face_bgr = cv2.cvtColor(face_np, cv2.COLOR_RGB2BGR)
-        result_bgr = cv2.seamlessClone(face_bgr, gen_bgr, mask, center, cv2.NORMAL_CLONE)
-        result_rgb = cv2.cvtColor(result_bgr, cv2.COLOR_BGR2RGB)
-        logger.info("Face overlay completed.")
-        return Image.fromarray(result_rgb)
-    except Exception as e:
-        logger.error(f"Error in overlay_face: {str(e)}", exc_info=True)
-        raise HTTPException(status_code=500, detail=f"Overlay failed: {str(e)}")
-
-def image_to_base64(image: Image.Image) -> str:
-    try:
-        buffered = BytesIO()
-        image.save(buffered, format="JPEG")  # Changed to JPEG to match Framer client expectation
-        img_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
-        logger.info("Image converted to base64 successfully.")
-        return img_base64
-    except Exception as e:
-        logger.error(f"Error converting image to base64: {str(e)}", exc_info=True)
-        raise HTTPException(status_code=500, detail=f"Base64 conversion failed: {str(e)}")
-
-@app.post("/predict")
-async def predict(
-    prompt: str = Form(...),
-    image: UploadFile = File(...),
-    negative_prompt: str = Form("low quality, blurry"),
-    seed: str = Form("66"),  # Changed to str to match Framer client
-    guidance_scale: str = Form("7.5"),  # Changed to str to match Framer client
-    num_inference_steps: str = Form("10"),  # Changed to str to match Framer client
-    strength: str = Form("0.75")  # Changed to str to match Framer client
-):
-    global pipe
-    try:
-        if pipe is None:
-            logger.error("Pipeline not initialized.")
-            raise HTTPException(status_code=500, detail="Pipeline not initialized.")
-
-        logger.info(f"Received inference request with prompt: {prompt}")
-
-        # Convert string parameters to appropriate types
-        try:
-            seed = int(seed)
-            guidance_scale = float(guidance_scale)
-            num_inference_steps = int(num_inference_steps)
-            strength = float(strength)
-        except ValueError as e:
-            logger.error(f"Invalid parameter format: {str(e)}")
-            raise HTTPException(status_code=400, detail=f"Invalid parameter format: {str(e)}")
-
-        # Load and process uploaded image
-        logger.info("Loading uploaded image...")
-        image_data = await image.read()
-        ref_image = Image.open(BytesIO(image_data)).convert("RGB")
-
-        # Extract face
-        logger.info("Extracting face...")
-        face_img, face_coords, original_size = extract_face(ref_image)
-
-        # Resize reference image to 512x512
-        logger.info("Resizing reference image...")
-        init_image = ref_image.resize((512, 512))
-
-        # Generate image with Stable Diffusion
-        logger.info("Starting image generation...")
-        generator = torch.Generator(device=device).manual_seed(seed)
-        generated_img = pipe(
-            prompt=prompt,
-            image=init_image,
-            strength=strength,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=guidance_scale,
-            negative_prompt=negative_prompt,
-            generator=generator
-        ).images[0]
-        logger.info("Image generation completed.")
-
-        # Scale face coordinates and overlay
-        logger.info("Scaling coordinates and overlaying face...")
-        scaled_coords = scale_coords(face_coords, original_size)
-        final_img = overlay_face(generated_img, face_img, scaled_coords)
-
-        # Convert to base64
-        logger.info("Converting final image to base64...")
-        result_base64 = image_to_base64(final_img)
-
-        logger.info("Inference completed successfully.")
-        return JSONResponse({
-            "result_image": f"data:image/jpeg;base64,{result_base64}"  # Match Framer client expectation
-        })
-    except HTTPException as e:
-        logger.error(f"HTTP Exception: {str(e)}")
-        raise
-    except Exception as e:
-        logger.error(f"Unexpected error in predict: {str(e)}", exc_info=True)
-        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
-
-@app.get("/")
-async def root():
-    return {"status": "API is running"}
-
-if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+import cv2
+import gradio as gr
+import os
+from PIL import Image
+import numpy as np
+import torch
+from torch.autograd import Variable
+from torchvision import transforms
+import torch.nn.functional as F
+import gdown
+import matplotlib.pyplot as plt
+import warnings
+warnings.filterwarnings("ignore")
+
+os.system("git clone https://github.com/xuebinqin/DIS")
+os.system("mv DIS/IS-Net/* .")
+
+# project imports
+from data_loader_cache import normalize, im_reader, im_preprocess 
+from models import *
+
+#Helpers
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+# Download official weights
+if not os.path.exists("saved_models"):
+    os.mkdir("saved_models")
+    os.system("mv isnet.pth saved_models/")
+    
+class GOSNormalize(object):
+    '''
+    Normalize the Image using torch.transforms
+    '''
+    def __init__(self, mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]):
+        self.mean = mean
+        self.std = std
+
+    def __call__(self,image):
+        image = normalize(image,self.mean,self.std)
+        return image
+
+
+transform =  transforms.Compose([GOSNormalize([0.5,0.5,0.5],[1.0,1.0,1.0])])
+
+def load_image(im_path, hypar):
+    im = im_reader(im_path)
+    im, im_shp = im_preprocess(im, hypar["cache_size"])
+    im = torch.divide(im,255.0)
+    shape = torch.from_numpy(np.array(im_shp))
+    return transform(im).unsqueeze(0), shape.unsqueeze(0) # make a batch of image, shape
+
+
+def build_model(hypar,device):
+    net = hypar["model"]#GOSNETINC(3,1)
+
+    # convert to half precision
+    if(hypar["model_digit"]=="half"):
+        net.half()
+        for layer in net.modules():
+            if isinstance(layer, nn.BatchNorm2d):
+                layer.float()
+
+    net.to(device)
+
+    if(hypar["restore_model"]!=""):
+        net.load_state_dict(torch.load(hypar["model_path"]+"/"+hypar["restore_model"], map_location=device))
+        net.to(device)
+    net.eval()  
+    return net
+
+    
+def predict(net,  inputs_val, shapes_val, hypar, device):
+    '''
+    Given an Image, predict the mask
+    '''
+    net.eval()
+
+    if(hypar["model_digit"]=="full"):
+        inputs_val = inputs_val.type(torch.FloatTensor)
+    else:
+        inputs_val = inputs_val.type(torch.HalfTensor)
+
+  
+    inputs_val_v = Variable(inputs_val, requires_grad=False).to(device) # wrap inputs in Variable
+   
+    ds_val = net(inputs_val_v)[0] # list of 6 results
+
+    pred_val = ds_val[0][0,:,:,:] # B x 1 x H x W    # we want the first one which is the most accurate prediction
+
+    ## recover the prediction spatial size to the orignal image size
+    pred_val = torch.squeeze(F.upsample(torch.unsqueeze(pred_val,0),(shapes_val[0][0],shapes_val[0][1]),mode='bilinear'))
+
+    ma = torch.max(pred_val)
+    mi = torch.min(pred_val)
+    pred_val = (pred_val-mi)/(ma-mi) # max = 1
+
+    if device == 'cuda': torch.cuda.empty_cache()
+    return (pred_val.detach().cpu().numpy()*255).astype(np.uint8) # it is the mask we need
+    
+# Set Parameters
+hypar = {} # paramters for inferencing
+
+
+hypar["model_path"] ="./saved_models" ## load trained weights from this path
+hypar["restore_model"] = "isnet.pth" ## name of the to-be-loaded weights
+hypar["interm_sup"] = False ## indicate if activate intermediate feature supervision
+
+##  choose floating point accuracy --
+hypar["model_digit"] = "full" ## indicates "half" or "full" accuracy of float number
+hypar["seed"] = 0
+
+hypar["cache_size"] = [1024, 1024] ## cached input spatial resolution, can be configured into different size
+
+## data augmentation parameters ---
+hypar["input_size"] = [1024, 1024] ## mdoel input spatial size, usually use the same value hypar["cache_size"], which means we don't further resize the images
+hypar["crop_size"] = [1024, 1024] ## random crop size from the input, it is usually set as smaller than hypar["cache_size"], e.g., [920,920] for data augmentation
+
+hypar["model"] = ISNetDIS()
+
+ # Build Model
+net = build_model(hypar, device)
+
+
+def inference(image):
+  image_path = image
+  
+  image_tensor, orig_size = load_image(image_path, hypar) 
+  mask = predict(net, image_tensor, orig_size, hypar, device)
+  
+  pil_mask = Image.fromarray(mask).convert('L')
+  im_rgb = Image.open(image).convert("RGB")
+  
+  im_rgba = im_rgb.copy()
+  im_rgba.putalpha(pil_mask)
+
+  return [im_rgba, pil_mask]
+
+
+title = "Highly Accurate Dichotomous Image Segmentation"
+description = "This is an unofficial demo for DIS, a model that can remove the background from a given image. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below.<br>GitHub: https://github.com/xuebinqin/DIS<br>Telegram bot: https://t.me/restoration_photo_bot<br>[![](https://img.shields.io/twitter/follow/DoEvent?label=@DoEvent&style=social)](https://twitter.com/DoEvent)"
+article = "<div><center><img src='https://visitor-badge.glitch.me/badge?page_id=max_skobeev_dis_cmp_public' alt='visitor badge'></center></div>"
+
+interface = gr.Interface(
+    fn=inference,
+    inputs=gr.Image(type='filepath'),
+    outputs=[gr.Image(type='filepath', format="png"), gr.Image(type='filepath', format="png")],
+    examples=[['robot.png'], ['ship.png']],
+    title=title,
+    description=description,
+    article=article,
+    flagging_mode="never",
+    cache_mode="lazy",
+    ).queue(api_open=True).launch(show_error=True, show_api=True)

requirements.txt

@@ -1,11 +1,8 @@
-torch==2.0.1
-diffusers==0.29.2
-transformers==4.41.2
-numpy==1.26.4
-Pillow==10.3.0
-opencv-python-headless==4.9.0.80
-mediapipe==0.10.14
-fastapi==0.111.0
-uvicorn==0.30.1
-requests==2.32.3
-pydantic==2.7.4
+torch
+torchvision
+requests
+gdown
+matplotlib
+opencv-python
+Pillow
+scikit-image