Spaces:
Build error
Build error
File size: 4,143 Bytes
4f87ea4 122ad84 4f87ea4 122ad84 2b243d1 122ad84 120b17e 4f87ea4 64ff386 4f87ea4 122ad84 |
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 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 |
import torch
import sys
sys.path.append('./ControlNetInpaint/')
from ultralytics import YOLO
from PIL import Image
import gradio as gr
import numpy as np
import os
import cv2
from PIL import Image
from matplotlib import pyplot as plt
from diffusers import StableDiffusionInpaintPipeline, ControlNetModel, UniPCMultistepScheduler
from diffusers.utils import load_image
from ControlNetInpaint.src import *
pipe_sd = StableDiffusionInpaintPipeline.from_pretrained(
"runwayml/stable-diffusion-inpainting",
revision="fp16",
torch_dtype=torch.float16,
)
# speed up diffusion process with faster scheduler and memory optimization
pipe_sd.scheduler = UniPCMultistepScheduler.from_config(pipe_sd.scheduler.config)
#pipe_sd.to('cpu')
# load control net and stable diffusion v1-5
from diffusers import StableDiffusionControlNetInpaintPipeline
controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
"runwayml/stable-diffusion-inpainting", controlnet=controlnet, torch_dtype=torch.float16
)
# speed up diffusion process with faster scheduler and memory optimization
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
# remove following line if xformers is not installed
text_prompt="Transform this image into a work of art by changing its style and color palette. Apply a distinct artistic style, such as impressionism, cubism, or surrealism, to give the image a unique and visually striking appearance. Experiment with brush strokes, textures, and effects to achieve the desired artistic effect while maintaining the essence of the original scene. Additionally, adjust the color palette to evoke a specific mood or theme. For example, infuse warm, earthy tones for a rustic and cozy feel, or opt for vibrant, psychedelic colors for a surreal and otherworldly atmosphere. The goal is to reimagine the image in a creative and expressive way, transforming it into a captivating visual masterpiece."
import os
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
#pipe.to('cuda')
import sys
sys.path.append("..")
from segment_anything import sam_model_registry, SamPredictor
sam_checkpoint = "sam_vit_h_4b8939.pth"
model_type = "vit_h"
device = "cpu"
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)
def image_enhance(image, mask):
image = np.array(image)
mask_image = np.array(mask)
canny_image = cv2.Canny(image, 100, 200)
canny_image = canny_image[:, :, None]
canny_image = np.concatenate([canny_image, canny_image, canny_image], axis=2)
image=Image.fromarray(image)
mask_image=Image.fromarray(mask_image)
canny_image = Image.fromarray(canny_image)
# generate image
generator = torch.manual_seed(42)
new_image = pipe(
text_prompt,
num_inference_steps=20,
generator=generator,
image=image,
control_image=canny_image,
controlnet_conditioning_scale = 0.5,
mask_image=mask_image
).images[0]
return new_image
def function(image):
original_image=image
model = YOLO('best_yolo_2.pt')
image = np.array(image)
results = model.predict(image)
img=Image.fromarray(results[0].plot())
num_result=len(results[0].boxes.cls)
names_output=[]
for i in range(num_result):
name=results[0].names[int(results[0].boxes.cls[i])]
names_output.append(name)
boxes = results[0].boxes
boxes=boxes.xyxy
#image=Image.fromarray(image)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
predictor = SamPredictor(sam)
predictor.set_image(image)
list_mask=[]
for box in boxes:
input_box = np.array(box.cpu())
masks, _, _ = predictor.predict(
point_coords=None,
point_labels=None,
box=input_box[None, :],
multimask_output=False,
)
for mask in masks:
list_mask.append(mask)
for i in range(len(list_mask)):
list_mask[i]=Image.fromarray(list_mask[i])
original_image=image_enhance(original_image,list_mask[i])
return original_image
demo=gr.Interface(fn=function,inputs="image",outputs=["image"])
demo.launch() |