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()